diff --git a/.dockerignore b/.dockerignore index 687c8fe3..6a56063f 100644 --- a/.dockerignore +++ b/.dockerignore @@ -13,6 +13,7 @@ Makefile *.tar .*_cache ENV/ +*.so # This is better, but doesn't work with podman. sigh. # * # !**/CMakeLists.txt @@ -24,3 +25,4 @@ ENV/ # !**/*.cc # !**/*.cpp # !**/*.h + diff --git a/CMakeLists.txt b/CMakeLists.txt index 88113295..5751f88a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -43,10 +43,9 @@ include(AlwaysColourCompilation) include_directories(include) -# Find Python -find_package(Python3 COMPONENTS Interpreter Development REQUIRED) -include_directories(${Python3_INCLUDE_DIRS}) +find_package(Python COMPONENTS Interpreter Development.Module NumPy REQUIRED) +# Dependency fetching # Find packages that are in conda environment find_package(fmt REQUIRED) find_package(spdlog REQUIRED) @@ -72,6 +71,21 @@ FetchContent_Declare( ) FetchContent_MakeAvailable(argparse) +# Required for Python module build of ffbidx +set(BUILD_FAST_INDEXER ON CACHE BOOL "" FORCE) +set(PYTHON_MODULE ON CACHE BOOL "" FORCE) +set(PYTHON_MODULE_RPATH ON CACHE BOOL "" FORCE) + +# ffbidx uses this to find site-packages +set(PYTHON_MODULE_PATH ${Python_SITEARCH} CACHE PATH "ffbidx python site-packages path" FORCE) + +FetchContent_Declare( + ffbidx + GIT_REPOSITORY https://github.com/paulscherrerinstitute/fast-feedback-indexer.git + GIT_TAG main +) +FetchContent_MakeAvailable(ffbidx) + # GTest - try find_package first, fallback to FetchContent find_package(GTest QUIET) if(NOT GTest_FOUND AND NOT TARGET GTest::gtest_main) @@ -191,5 +205,3 @@ add_subdirectory(h5read) add_subdirectory(baseline) add_subdirectory(spotfinder) add_subdirectory(integrator) - -configure_file(${CMAKE_SOURCE_DIR}/scripts/ssx_index ${CMAKE_BINARY_DIR}/bin/ssx_index COPYONLY) diff --git a/Dockerfile b/Dockerfile index 9e30623f..782d057f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -24,21 +24,22 @@ RUN micromamba create -y -f /opt/runtime-environment.yml -p /opt/ffs # Copy source COPY . /opt/ffs_src +ENV CMAKE_GENERATOR Ninja # Build the C++/CUDA backend WORKDIR /opt/build -RUN cmake /opt/ffs_src -G Ninja -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/opt/ffs -DHDF5_ROOT=/opt/ffs -DCUDA_ARCH=80 -RUN cmake --build . --target spotfinder --target spotfinder32 +RUN cmake /opt/ffs_src \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_PREFIX=/opt/ffs \ + -DHDF5_ROOT=/opt/ffs \ + -DPython3_ROOT_DIR=/opt/ffs \ + -DCUDA_ARCH=80 -RUN cmake --install . --component Runtime +RUN cmake --build . -# Build and install dx2 submodule -WORKDIR /opt/build_dx2 -RUN cmake /opt/ffs_src/dx2 -G Ninja -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/opt/ffs -RUN cmake --build . && cmake --install . +RUN cmake --install . # Install Python package -WORKDIR /opt/build RUN SETUPTOOLS_SCM_PRETEND_VERSION_FOR_FFS=1.0 /opt/ffs/bin/pip3 install /opt/ffs_src # Now copy this into an isolated runtime container @@ -46,7 +47,7 @@ FROM nvcr.io/nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu24.04 LABEL org.opencontainers.image.title="fast-feedback-service" \ org.opencontainers.image.description="GPU-accelerated fast-feedback X-ray diffraction analysis service" \ - org.opencontainers.image.authors="Nicholas Devenish , Dimitrios Vlachos " \ + org.opencontainers.image.authors="Nicholas Devenish , Dimitrios Vlachos , James Beilsten-Edmands " \ org.opencontainers.image.source="https://github.com/DiamondLightSource/fast-feedback-service" \ org.opencontainers.image.licenses="BSD-3-Clause" diff --git a/baseline/indexer/CMakeLists.txt b/baseline/indexer/CMakeLists.txt index d9fee60b..6cdc8e34 100644 --- a/baseline/indexer/CMakeLists.txt +++ b/baseline/indexer/CMakeLists.txt @@ -60,9 +60,11 @@ target_link_libraries(baseline_indexer # Create Python module nanobind_add_module(index index_module.cpp) target_include_directories(index PRIVATE ${CMAKE_SOURCE_DIR}/baseline/predictor) -target_link_libraries(index PRIVATE Eigen3::Eigen dx2 ${Python3_LIBRARIES}) +target_link_libraries(index PRIVATE Eigen3::Eigen dx2 Python::Module) -set(PYTHON_MODULE_DIR ${CMAKE_SOURCE_DIR}/src/ffs) -add_custom_command(TARGET index POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy $ ${PYTHON_MODULE_DIR} + +# Install into the runtime Python sitearch +install(TARGETS index + LIBRARY DESTINATION ${Python_SITEARCH}/ffs + COMPONENT Runtime ) diff --git a/baseline/indexer/index_module.cpp b/baseline/indexer/index_module.cpp index d2989fc7..c782ae87 100644 --- a/baseline/indexer/index_module.cpp +++ b/baseline/indexer/index_module.cpp @@ -1,11 +1,13 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -220,6 +222,9 @@ NB_MODULE(index, m) { .def_prop_ro("delpsi", [](const IndexingResult &r) { return r.delpsi; }) .def_prop_ro("rmsds", [](const IndexingResult &r) { return r.rmsds; }); m.def("make_panel", &make_panel, "Create a configured Panel object"); + m.def("calculate_mu_for_material_at_wavelength", + &calculate_mu_for_material_at_wavelength, + "Calculate the absorption coefficient from material and wavelength"); m.def("ssx_xyz_to_rlp", &ssx_xyz_to_rlp, nb::arg("xyzobs_px"), diff --git a/build.sh b/build.sh index 532eccbe..4c8ad4ab 100755 --- a/build.sh +++ b/build.sh @@ -161,6 +161,17 @@ build_directory() { make -j"$JOBS" fi ) + + # Install + print_status "Installing $description..." + ( + cd "$build_dir" + if [[ "$BUILD_CMD" == "ninja" ]]; then + ninja "install" + else + make "install" -j"$JOBS" + fi + ) print_success "Successfully built $description" } diff --git a/chart/values.yaml b/chart/values.yaml index 3db44ca4..281ff602 100644 --- a/chart/values.yaml +++ b/chart/values.yaml @@ -27,4 +27,4 @@ podSecurityContext: # Container-level security context containerSecurityContext: # k8s-i24-beamline user - runAsUser: 36145 \ No newline at end of file + runAsUser: 36145 diff --git a/environment.yml b/environment.yml index bb4a1e71..f05a02e3 100644 --- a/environment.yml +++ b/environment.yml @@ -15,6 +15,7 @@ dependencies: - hdf5=1.14.6 - hdf5-external-filter-plugins=0.1.0 - nlohmann_json=3.12.0 + - numpy=2.3.4 - pre-commit=4.3.0 - pydantic=2.12.3 - pytest=8.4.2 diff --git a/pyproject.toml b/pyproject.toml index 551870d6..dce4a512 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ readme = "README.md" authors = [ { name = "Nicholas Devenish", email = "ndevenish@gmail.com" }, { name = "Dimitri Vlachos", email = "dimitrios.vlachos@diamond.ac.uk" }, + { name = "James Beilsten-Edmands", email = "james.beilsten-edmands@diamond.ac.uk"}, ] requires-python = ">=3.12" @@ -22,6 +23,9 @@ dependencies = [ "zocalo>=1.2.0", ] +[project.scripts] +ssx_index = "ffs.ssx_index:run" + [dependency-groups] dev = [ "pytest", diff --git a/runtime-environment.yml b/runtime-environment.yml index 6082a291..0b3d54e9 100644 --- a/runtime-environment.yml +++ b/runtime-environment.yml @@ -7,6 +7,7 @@ dependencies: - gemmi=0.7.3 - hdf5=1.14.6 - hdf5-external-filter-plugins=0.1.0 + - h5py - python=3.13.7 - pip - zocalo @@ -14,3 +15,4 @@ dependencies: - pydantic - rich - spdlog=1.15.3 + - numpy=2.3.4 diff --git a/src/ffs/service.py b/src/ffs/service.py index f2ac44d6..a572ec6c 100644 --- a/src/ffs/service.py +++ b/src/ffs/service.py @@ -11,14 +11,21 @@ import threading import time from datetime import datetime +from enum import Enum from pathlib import Path -from typing import Iterator, Optional +from typing import Iterator, Literal, Optional, Union +import gemmi +import numpy as np +import pydantic import workflows.recipe -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, Field, PrivateAttr, ValidationError from rich.logging import RichHandler from workflows.services.common_service import CommonService +import ffs.index +from ffs.ssx_index import GPUIndexer + logger = logging.getLogger(__name__) logger.level = logging.DEBUG @@ -39,17 +46,111 @@ class PiaRequest(BaseModel): detector_distance: float d_min: float | None = None d_max: float | None = None + unit_cell: tuple[float, float, float, float, float, float] | None = None + detector: str = "Eiger16M" + + @pydantic.validator("unit_cell", pre=True) + def check_unit_cell(cls, v): + if not v: + return None + orig_v = v + if isinstance(v, str): + v = v.replace(",", " ").split() + v = [float(v) for v in v] + try: + assert len(v) == 6 + except Exception: + raise ValueError(f"Invalid unit_cell {orig_v}") + return v + + +class Material(str, Enum): + Si = "Si" + CdTe = "CdTe" + + +class DetectorParameters(BaseModel): + """ + Define a set of detector metadata that derived classes + need to provide. + This class is not to be instantiated directly. + """ + + detector_type: str + thickness: float + material: Material + pixel_size_x: float + pixel_size_y: float + image_size_x: int + image_size_y: int + # mu cache not serialized to dict/json, ok as mu not needed for spotfinder + _mu_cache: dict = PrivateAttr(default_factory=dict) + + def __init_subclass__(cls, **kwargs): + # enforce setting of defaults for all fields in subclasses. + super().__init_subclass__(**kwargs) + missing_defaults = [ + name + for name, field in cls.__fields__.items() + if field.default is None and field.default_factory is None + ] + if missing_defaults: + raise TypeError( + f"{cls.__name__} must define default values for all fields: missing {missing_defaults}" + ) + + def calculate_mu(self, wavelength: float) -> float: + if wavelength not in self._mu_cache: + self._mu_cache[wavelength] = ( + ffs.index.calculate_mu_for_material_at_wavelength( + self.material, wavelength + ) + ) + return self._mu_cache[wavelength] + + +class Eiger16M(DetectorParameters): + detector_type: Literal["Eiger16M"] + thickness: float = 0.45 + material: Material = Material.Si + pixel_size_x: float = 0.075 + pixel_size_y: float = 0.075 + image_size_x: int = 4148 + image_size_y: int = 4362 + + +class Eiger4M(DetectorParameters): + detector_type: Literal["Eiger4M"] + thickness: float = 0.45 + material: Material = Material.Si + pixel_size_x: float = 0.075 + pixel_size_y: float = 0.075 + image_size_x: int = 2068 + image_size_y: int = 2162 + + +class Eiger9MCdTe(DetectorParameters): + detector_type: Literal["Eiger9MCdTe"] + thickness: float = 0.75 + material: Material = Material.CdTe + pixel_size_x: float = 0.075 + pixel_size_y: float = 0.075 + image_size_x: int = 3108 + image_size_y: int = 3262 class DetectorGeometry(BaseModel): - pixel_size_x: float = 0.075 # Default value for Eiger - pixel_size_y: float = 0.075 # Default value for Eiger distance: float beam_center_x: float beam_center_y: float + detector: Union[Eiger9MCdTe, Eiger16M, Eiger4M] = Field( + ..., discriminator="detector_type" + ) def to_json(self): - return json.dumps(self.dict(), indent=4) + d = self.dict() + d.update(self.detector.dict()) + return json.dumps(d, indent=4) def _setup_rich_logging(level=logging.DEBUG): @@ -203,6 +304,15 @@ def initializing(self): ) self._spotfinder_executable = _find_spotfinder() self._order_resolver = MessageOrderResolver(self.log) + ## Initialise the fast-feedback-indexer + self.indexer = None + self.output_for_index = False # Only turn on when we have confirmed all the things we need (cell, etc) + try: + self.indexer = GPUIndexer() + except ModuleNotFoundError: + self.log.debug( + "ffbidx not found, has the fast-feedback-indexer module been built and sourced?" + ) def gpu_per_image_analysis( self, @@ -225,6 +335,7 @@ def gpu_per_image_analysis( distance=parameters.detector_distance, beam_center_x=parameters.xBeam, beam_center_y=parameters.yBeam, + detector={"detector_type": parameters.detector}, ) self.log.debug("{detector_geometry.to_json()=}") except ValidationError as e: @@ -232,6 +343,33 @@ def gpu_per_image_analysis( f"Rejecting PIA request for {parameters.dcgid}/{parameters.message_index}({parameters.dcid}): Invalid detector parameters \n{e}" ) + if self.indexer and parameters.unit_cell and parameters.wavelength: + ## We have all we need to index, so make up to date models. + cell = gemmi.UnitCell(*parameters.unit_cell) + self.indexer.cell = np.reshape( + np.array(cell.orth.mat, dtype="float32"), (3, 3) + ) ## Cell as an orthogonalisation matrix + ## convert beam centre to correct units (given in mm, want in px). + px_size_x = detector_geometry.detector.pixel_size_x + px_size_y = detector_geometry.detector.pixel_size_y + mu = detector_geometry.detector.calculate_mu(parameters.wavelength) + self.indexer.panel = ffs.index.make_panel( + detector_geometry.distance, + detector_geometry.beam_center_x / px_size_x, + detector_geometry.beam_center_y / px_size_y, + px_size_x, + px_size_y, + detector_geometry.detector.image_size_x, + detector_geometry.detector.image_size_y, + detector_geometry.detector.thickness, + mu, + ) + self.indexer.wavelength = parameters.wavelength + self.output_for_index = ( + True # The indexer has been configured, so can run the spotfinder + ) + # with --output-for-index and capture the results in read_and_send. + start_time = time.monotonic() self.log.info( f"Gotten PIA request for {parameters.dcgid}/{parameters.message_index}({parameters.dcid}): {parameters.filename}/:{parameters.start_frame_index}-{parameters.start_frame_index + parameters.number_of_frames}" @@ -298,6 +436,8 @@ def gpu_per_image_analysis( command.extend(["--dmin", str(parameters.d_min)]) if parameters.d_max: command.extend(["--dmax", str(parameters.d_max)]) + if self.output_for_index: + command.extend(["--output-for-index"]) self.log.info(f"Running: {' '.join(str(x) for x in command)}") @@ -336,6 +476,14 @@ def read_and_send() -> None: data["file-seen-at"] = time.time() # XRC has one-based-indexing data["file-number"] += 1 + ## Do indexing + if self.output_for_index: + xyzobs_px = np.array(data["spot_centers"]) + indexing_result = self.indexer.index(xyzobs_px) + self.log.info(indexing_result.model_dump_json(indent=2)) + result = indexing_result.model_dump() + data.update(result) + del data["spot_centers"] # don't send this data array onwards. self.log.info(f"Sending: {data}") rw.set_default_channel("result") rw.send_to("result", data) diff --git a/src/ffs/ssx_index.py b/src/ffs/ssx_index.py index 72a954b7..ea8e195e 100644 --- a/src/ffs/ssx_index.py +++ b/src/ffs/ssx_index.py @@ -240,7 +240,7 @@ def write_table(self, filename): ## rlp, flags, xyzobs.mm.value -def run(args): +def run(args=None): st = time.time() parser = argparse.ArgumentParser( prog="index", diff --git a/tests/test_gpu_ssx_index.py b/tests/test_gpu_ssx_index.py index fea21144..a2cb39d3 100644 --- a/tests/test_gpu_ssx_index.py +++ b/tests/test_gpu_ssx_index.py @@ -157,8 +157,8 @@ def test_gpu_ssx_index(dials_data, tmp_path): cwd=tmp_path, ) assert not proc.stderr - assert tmp_path / "indexed.refl" - assert tmp_path / "indexed_crystals.json" + assert (tmp_path / "indexed.refl").exists() + assert (tmp_path / "indexed_crystals.json").exists() with open(tmp_path / "indexed_crystals.json", "r") as f: crystals = json.load(f) for i, (crystal, expected) in enumerate(