Ispyb integration vmxm grid scan

pyproject.toml

@@ -94,7 +94,7 @@ dev = [
     "tox-uv",
     "types-mock",
     "types-requests",
-    "sphinxcontrib.mermaid", # To build dodal docs
+    "sphinxcontrib.mermaid",     # To build dodal docs
 ]
 
 [project.scripts]

src/mx_bluesky/beamlines/i02_1/__init__.py

@@ -0,0 +1,3 @@
+from mx_bluesky.beamlines.i02_1.i02_1_gridscan_plan import i02_1_gridscan_plan
+
+__all__ = ["i02_1_gridscan_plan"]

src/mx_bluesky/beamlines/i02_1/composites.py

@@ -0,0 +1,16 @@
+from pathlib import Path
+
+from mx_bluesky.beamlines.i02_1.parameters.gridscan import SpecifiedTwoDGridScan
+
+
+class I02_1FgsParams(SpecifiedTwoDGridScan):  # noqa: N801
+    """For VMXm gridscans, GDA currently takes the snapshots and provides bluesky with a path, and
+    sends over the grid parameters"""
+
+    path_to_xtal_snapshot: Path
+    beam_size_x: float
+    beam_size_y: float
+    microns_per_pixel_x: float
+    microns_per_pixel_y: float
+    upper_left_x: int  # position of X,Y for the top left of the grid, in pixels
+    upper_left_y: int

src/mx_bluesky/beamlines/i02_1/external_interaction/callbacks/gridscan/ispyb_callback.py

@@ -0,0 +1,88 @@
+from collections.abc import Sequence
+
+from mx_bluesky.beamlines.i02_1.composites import I02_1FgsParams
+from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_mapping import (
+    construct_comment_for_gridscan,
+)
+from mx_bluesky.common.external_interaction.callbacks.grid.gridscan.ispyb_callback import (
+    GridscanISPyBCallback as CommonGridscanISPyBCallback,
+)
+from mx_bluesky.common.external_interaction.ispyb.data_model import (
+    DataCollectionGridInfo,
+    DataCollectionInfo,
+    Orientation,
+    ScanDataInfo,
+)
+from mx_bluesky.common.utils.log import ISPYB_ZOCALO_CALLBACK_LOGGER
+
+
+def _make_comment(x_steps: int, y_steps: int) -> str:
+    return f"Diffraction grid scan of {x_steps} by {y_steps}."
+
+
+class GridscanISPyBCallback(CommonGridscanISPyBCallback):
+    def _get_scan_infos(self, doc) -> Sequence[ScanDataInfo]:
+        """
+        For VMXm, grid information is available immediately after the plan is triggered.
+        """
+        assert isinstance(self.params, I02_1FgsParams)
+        assert self.ispyb_ids.data_collection_ids, "No current data collection"
+        assert self.data_collection_group_info, "No data collection group"
+        data = doc["data"]
+        scan_data_infos = []
+
+        for grid_num in range(self.params.num_grids):
+            omega = data.get("gonio-omega", self.params.omega_starts_deg[grid_num])
+
+            ISPYB_ZOCALO_CALLBACK_LOGGER.info(
+                f"Generating dc info for gridplane XY, omega {omega}"
+            )
+            data_collection_number = self.params.detector_params.run_number
+            file_template = f"{self.params.detector_params.prefix}_{data_collection_number}_master.h5"
+            # Snapshots have already been taken in GDA
+
+            data_collection_info = DataCollectionInfo(
+                xtal_snapshot1=str(self.params.path_to_xtal_snapshot),
+                xtal_snapshot2=str(self.params.path_to_xtal_snapshot),
+                xtal_snapshot3=str(self.params.path_to_xtal_snapshot),
+                n_images=self.params.num_images,
+                data_collection_number=data_collection_number,
+                file_template=file_template,
+            )
+            data_collection_grid_info = DataCollectionGridInfo(
+                dx_in_mm=self.params.x_step_size_um / 1000,
+                dy_in_mm=self.params.y_step_sizes_um[grid_num] / 1000,
+                steps_x=self.params.x_steps,
+                steps_y=self.params.y_steps[grid_num],
+                microns_per_pixel_x=self.params.microns_per_pixel_x,
+                microns_per_pixel_y=self.params.microns_per_pixel_y,
+                snapshot_offset_x_pixel=self.params.upper_left_x,
+                snapshot_offset_y_pixel=self.params.upper_left_y,
+                orientation=Orientation.HORIZONTAL,
+                snaked=True,
+            )
+            data_collection_info.comments = construct_comment_for_gridscan(
+                data_collection_grid_info
+            )
+
+            data_collection_id = self.ispyb_ids.data_collection_ids[0]
+
+            self.data_collection_group_info.comments = _make_comment(
+                self.params.x_steps, self.params.y_steps[0]
+            )
+
+            self._populate_axis_info(data_collection_info, doc["data"])
+
+            scan_data_info = ScanDataInfo(
+                data_collection_info=data_collection_info,
+                data_collection_id=data_collection_id,
+                data_collection_grid_info=data_collection_grid_info,
+            )
+
+            scan_data_infos.append(scan_data_info)
+
+        ISPYB_ZOCALO_CALLBACK_LOGGER.info(
+            "Updating ispyb data collection after loading grid params"
+        )
+
+        return scan_data_infos

src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py

@@ -1,24 +1,33 @@
 from functools import partial
+from pathlib import Path
 
 import bluesky.preprocessors as bpp
 import pydantic
 from bluesky.utils import MsgGenerator
 from dodal.beamlines.i02_1 import ZebraFastGridScanTwoD
 from dodal.common import inject
 from dodal.devices.attenuator.attenuator import ReadOnlyAttenuator
+from dodal.devices.beamlines.i02_1.flux import Flux
 from dodal.devices.common_dcm import DoubleCrystalMonochromatorBase
 from dodal.devices.fast_grid_scan import (
     set_fast_grid_scan_params as set_flyscan_params_plan,
 )
-from dodal.devices.flux import Flux
-from dodal.devices.s4_slit_gaps import S4SlitGaps
+from dodal.devices.motors import XYZWrappedOmegaStage
+from dodal.devices.slits import Slits
 from dodal.devices.undulator import BaseUndulator
 from dodal.devices.zebra.zebra import Zebra
+from pydantic import BaseModel
+from pydantic_extra_types.semantic_version import SemanticVersion
+from semver import Version
 
+from mx_bluesky.beamlines.i02_1.composites import I02_1FgsParams
 from mx_bluesky.beamlines.i02_1.device_setup_plans.setup_zebra import (
     setup_zebra_for_gridscan,
     tidy_up_zebra_after_gridscan,
 )
+from mx_bluesky.beamlines.i02_1.external_interaction.callbacks.gridscan.ispyb_callback import (
+    GridscanISPyBCallback,
+)
 from mx_bluesky.beamlines.i02_1.parameters.gridscan import SpecifiedTwoDGridScan
 from mx_bluesky.common.experiment_plans.common_flyscan_xray_centre_plan import (
     BeamlineSpecificFGSFeatures,
@@ -28,43 +37,50 @@
 from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import (
     ZocaloCallback,
 )
-from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import (
-    GridscanISPyBCallback,
-    generate_start_info_from_omega_map,
-)
-from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import (
+from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback import (
     GridscanNexusFileCallback,
 )
+from mx_bluesky.common.external_interaction.callbacks.grid.gridscan.ispyb_callback import (
+    ispyb_activation_decorator,
+)
+from mx_bluesky.common.external_interaction.callbacks.grid.utils import (
+    generate_start_info_from_num_grids,
+)
+from mx_bluesky.common.parameters.components import (
+    IspybExperimentType,
+    get_param_version,
+)
 from mx_bluesky.common.parameters.constants import (
     EnvironmentConstants,
     PlanNameConstants,
 )
 from mx_bluesky.common.parameters.device_composites import (
     FlyScanEssentialDevices,
-    GonioWithOmegaType,
 )
-from mx_bluesky.common.parameters.gridscan import GenericGrid
+from mx_bluesky.common.parameters.gridscan import PositiveFloat
 from mx_bluesky.common.utils.log import LOGGER
 
+DEFAULT_BOX_SIZE_UM = 2
 
-def create_gridscan_callbacks() -> tuple[
-    GridscanNexusFileCallback, GridscanISPyBCallback
-]:
+
+def create_gridscan_callbacks(
+    params: I02_1FgsParams,
+) -> tuple[GridscanNexusFileCallback, GridscanISPyBCallback]:
     return (
         GridscanNexusFileCallback(param_type=SpecifiedTwoDGridScan),
         GridscanISPyBCallback(
-            param_type=GenericGrid,
+            param_type=I02_1FgsParams,
             emit=ZocaloCallback(
                 PlanNameConstants.DO_FGS,
                 EnvironmentConstants.ZOCALO_ENV,
-                generate_start_info_from_omega_map,
+                lambda: generate_start_info_from_num_grids(params),
             ),
         ),
     )
 
 
 @pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
-class I021FlyScanXRayCentreComposite(FlyScanEssentialDevices[GonioWithOmegaType]):
+class I021FlyScanXRayCentreComposite(FlyScanEssentialDevices[XYZWrappedOmegaStage]):
     """All devices which are directly or indirectly required by this plan"""
 
     zebra: Zebra
@@ -73,7 +89,7 @@ class I021FlyScanXRayCentreComposite(FlyScanEssentialDevices[GonioWithOmegaType]
     attenuator: ReadOnlyAttenuator
     flux: Flux
     undulator: BaseUndulator
-    s4_slit_gaps: S4SlitGaps
+    s4_slit_gaps: Slits
 
 
 def construct_i02_1_specific_features(
@@ -122,17 +138,82 @@ def _tidy_plan(
     yield from tidy_up_zebra_after_gridscan(fgs_composite.zebra)
 
 
+PARAMETER_VERSION = Version.parse("1.0.0")
+
+
+def get_internal_param_version() -> SemanticVersion:
+    return SemanticVersion.validate_from_str(str(PARAMETER_VERSION))
+
+
+class ExternalGridScanParams(BaseModel):
+    visit: str
+    file_name: str
+    storage_directory: str
+    exposure_time_s: float
+    snapshot_directory: Path
+    x_start_um: float
+    y_start_um: float
+    z_start_um: float
+    x_steps: int
+    y_steps: int
+    beam_size_x: float
+    beam_size_y: float
+    microns_per_pixel_x: float
+    microns_per_pixel_y: float
+    upper_left_x: int
+    upper_left_y: int
+    detector_distance_mm: float
+    sample_id: int
+
+    # GDA branch needs to update for these params
+    x_step_size_um: PositiveFloat
+    y_step_sizes_um: list[PositiveFloat]
+    omega_start_deg: int
+
+
+def get_internal_params(params: ExternalGridScanParams) -> I02_1FgsParams:
+    return I02_1FgsParams(
+        y_starts_um=[params.y_start_um],
+        x_start_um=params.x_start_um,
+        z_starts_um=[params.z_start_um],
+        omega_starts_deg=[params.omega_start_deg],
+        sample_id=params.sample_id,
+        visit=params.visit,
+        parameter_model_version=get_param_version(),
+        file_name=params.file_name,
+        storage_directory=params.storage_directory,
+        x_steps=params.x_steps,
+        y_steps=[params.y_steps],
+        path_to_xtal_snapshot=params.snapshot_directory,
+        beam_size_x=params.beam_size_x,
+        beam_size_y=params.beam_size_y,
+        microns_per_pixel_x=params.microns_per_pixel_x,
+        microns_per_pixel_y=params.microns_per_pixel_y,
+        upper_left_x=params.upper_left_x,
+        upper_left_y=params.upper_left_y,
+        detector_distance_mm=params.detector_distance_mm,
+        ispyb_experiment_type=IspybExperimentType.SAD,
+        x_step_size_um=params.x_step_size_um,
+        y_step_sizes_um=params.y_step_sizes_um,
+        use_roi_mode=False,
+        box_size_um=DEFAULT_BOX_SIZE_UM,
+    )
+
+
 def i02_1_gridscan_plan(
-    parameters: SpecifiedTwoDGridScan,
+    parameters: ExternalGridScanParams,
     composite: I021FlyScanXRayCentreComposite = inject(""),
 ) -> MsgGenerator:
     """BlueAPI entry point for i02-1 grid scans"""
 
-    beamline_specific = construct_i02_1_specific_features(composite, parameters)
-    callbacks = create_gridscan_callbacks()
+    params = get_internal_params(parameters)
+
+    beamline_specific = construct_i02_1_specific_features(composite, params)
+    callbacks = create_gridscan_callbacks(params)
 
     @bpp.subs_decorator(callbacks)
+    @ispyb_activation_decorator(params)
     def decorated_flyscan_plan():
-        yield from common_flyscan_xray_centre(composite, parameters, beamline_specific)
+        yield from common_flyscan_xray_centre(composite, params, beamline_specific)
 
     yield from decorated_flyscan_plan()

src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py

@@ -62,16 +62,19 @@
 from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import (
     ZocaloCallback,
 )
-from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import (
-    GridscanISPyBCallback,
-    generate_start_info_from_omega_map,
+from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import (
+    GridDetectAndScanISPyBCallback,
 )
-from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import (
+from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.nexus_callback import (
     GridscanNexusFileCallback,
 )
+from mx_bluesky.common.external_interaction.callbacks.grid.utils import (
+    generate_start_info_from_omega_map,
+)
 from mx_bluesky.common.parameters.components import get_param_version
 from mx_bluesky.common.parameters.constants import (
     EnvironmentConstants,
+    GridscanParamConstants,
     OavConstants,
     PlanGroupCheckpointConstants,
     PlanNameConstants,
@@ -208,7 +211,10 @@ def _inner_grid_detect_then_xrc():
         # Hyperion handles its callbacks differently to BlueAPI-managed plans, see
         # https://github.com/DiamondLightSource/mx-bluesky/issues/1117
         flyscan_event_handler = XRayCentreEventHandler()
-        callbacks = *create_gridscan_callbacks(), flyscan_event_handler
+        callbacks = (
+            *create_gridscan_callbacks(),
+            flyscan_event_handler,
+        )
 
         @bpp.subs_decorator(callbacks)
         @verify_undulator_gap_before_run_decorator(composite)
@@ -271,16 +277,18 @@ def get_ready_for_oav_and_close_shutter(
 
 
 def create_gridscan_callbacks() -> tuple[
-    GridscanNexusFileCallback, GridscanISPyBCallback
+    GridscanNexusFileCallback, GridDetectAndScanISPyBCallback
 ]:
     return (
         GridscanNexusFileCallback(param_type=SpecifiedThreeDGridScan),
-        GridscanISPyBCallback(
+        GridDetectAndScanISPyBCallback(
             param_type=GenericGrid,
             emit=ZocaloCallback(
                 PlanNameConstants.DO_FGS,
                 EnvironmentConstants.ZOCALO_ENV,
-                generate_start_info_from_omega_map,
+                lambda: generate_start_info_from_omega_map(
+                    [GridscanParamConstants.OMEGA_1, GridscanParamConstants.OMEGA_2]
+                ),
             ),
         ),
     )

src/mx_bluesky/common/experiment_plans/common_grid_detect_then_xray_centre_plan.py

@@ -32,7 +32,7 @@
     GridDetectionCallback,
     GridParamUpdate,
 )
-from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import (
+from mx_bluesky.common.external_interaction.callbacks.grid.grid_detect_and_scan.ispyb_callback import (
     ispyb_activation_wrapper,
 )
 from mx_bluesky.common.parameters.constants import (

src/mx_bluesky/common/experiment_plans/inner_plans/do_fgs.py

@@ -1,4 +1,4 @@
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from time import time
 
 import bluesky.plan_stubs as bps
@@ -65,7 +65,7 @@ def kickoff_and_complete_gridscan(
     detector: EigerDetector,  # Once Eiger inherits from StandardDetector, use that type instead
     synchrotron: Synchrotron,
     scan_points: list[AxesPoints[Axis]],
-    omega_starts_deg: list[float],
+    omega_starts_deg: Sequence[float],
     plan_during_collection: Callable[[], MsgGenerator] | None = None,
 ):
     """Triggers a grid scan motion program and waits for completion, accounting for synchrotron topup.