Add GPU indexing to the live service

CMakeLists.txt

@@ -22,9 +22,14 @@ include(AlwaysColourCompilation)
 include_directories(include)
 
 # Find Python
-find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
+find_package(Python3 COMPONENTS Interpreter Development.Module REQUIRED)
 include_directories(${Python3_INCLUDE_DIRS})
 
+# Make sure we can install ffbidx into correct location.
+if(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT)
+    set(CMAKE_INSTALL_PREFIX "${CMAKE_BINARY_DIR}/install" CACHE PATH "Install path" FORCE)
+endif()
+
 # Dependency fetching
 set(FETCHCONTENT_QUIET OFF)
 include(FetchContent)
@@ -48,6 +53,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)
+
+# Override install path for ffbidx to site-packages
+set(PYTHON_MODULE_PATH ${Python3_SITEARCH} CACHE PATH "Python install path" FORCE)
+
+FetchContent_Declare(
+  ffbidx
+  GIT_REPOSITORY https://github.com/paulscherrerinstitute/fast-feedback-indexer.git
+  GIT_TAG main
+)
+FetchContent_MakeAvailable(ffbidx)
+
 find_package(fmt REQUIRED)
 find_package(spdlog REQUIRED)
 
@@ -128,4 +148,13 @@ add_subdirectory(h5read)
 add_subdirectory(baseline)
 add_subdirectory(spotfinder)
 
+# define a ssx_index 'executable' that actually just runs the python code (for style consistency)
 configure_file(${CMAKE_SOURCE_DIR}/scripts/ssx_index ${CMAKE_BINARY_DIR}/bin/ssx_index COPYONLY)
+
+# run 'pip install .' to install the c++ extension for ssx_index
+install(CODE "
+  execute_process(
+    COMMAND ${Python3_EXECUTABLE} -m pip install .
+    WORKING_DIRECTORY \"${CMAKE_SOURCE_DIR}\"
+  )
+")

baseline/indexer/index_module.cpp

@@ -1,11 +1,13 @@
 #include <nanobind/eigen/dense.h>
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/array.h>
+#include <nanobind/stl/string.h>
 #include <nanobind/stl/tuple.h>
 #include <nanobind/stl/vector.h>
 
 #include <dx2/crystal.hpp>
 #include <dx2/detector.hpp>
+#include <dx2/detector_attenuations.hpp>
 #include <dx2/reflection.hpp>
 #include <experimental/mdspan>
 #include <tuple>
@@ -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"),

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"
 }

src/ffs/service.py

@@ -10,14 +10,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
 
@@ -38,17 +45,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):
@@ -175,6 +276,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,
@@ -197,13 +307,41 @@ 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:
             self.log.warning(
                 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}"
@@ -271,6 +409,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)}")
 
@@ -309,6 +449,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)

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(