Add ispyb integration for vmxm grid scans

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/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,4 +1,5 @@
 from functools import partial
+from pathlib import Path
 
 import bluesky.preprocessors as bpp
 import pydantic
@@ -14,11 +15,18 @@
 from dodal.devices.s4_slit_gaps import S4SlitGaps
 from dodal.devices.undulator import BaseUndulator
 from dodal.devices.zebra.zebra import Zebra
+from pydantic import BaseModel, PrivateAttr
+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,13 +36,15 @@
 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.constants import (
     EnvironmentConstants,
     PlanNameConstants,
@@ -43,21 +53,20 @@
     FlyScanEssentialDevices,
     GonioWithOmegaType,
 )
-from mx_bluesky.common.parameters.gridscan import GenericGrid
 from mx_bluesky.common.utils.log import LOGGER
 
 
-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),
             ),
         ),
     )
@@ -122,16 +131,42 @@ 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):
+    gda_parameter_version: SemanticVersion
+    visit: str
+    file_name: str
+    storage_directory: Path
+    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
+    sample_id: int | None = None
+
+    # Internal parameter version compatible with this external model
+    _internal_param_version: str = PrivateAttr(default="6.0.0")
+
+
 def i02_1_gridscan_plan(
-    parameters: SpecifiedTwoDGridScan,
+    parameters: I02_1FgsParams,
     composite: FlyScanXRayCentreComposite = inject(""),
 ) -> MsgGenerator:
     """BlueAPI entry point for i02-1 grid scans"""
 
     beamline_specific = construct_i02_1_specific_features(composite, parameters)
-    callbacks = create_gridscan_callbacks()
+    callbacks = create_gridscan_callbacks(parameters)
 
     @bpp.subs_decorator(callbacks)
+    @ispyb_activation_decorator(parameters)
     def decorated_flyscan_plan():
         yield from common_flyscan_xray_centre(composite, parameters, beamline_specific)
 

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)
@@ -270,16 +276,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

@@ -31,7 +31,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.

src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py

@@ -129,25 +129,28 @@ def activity_gated_event(self, doc: Event) -> Event:
         return self.tag_doc(doc)
 
     def _handle_ispyb_hardware_read(self, doc) -> Sequence[ScanDataInfo]:
+        _data = doc["data"]
+
         assert self.params, "Event handled before activity_gated_start received params"
         ISPYB_ZOCALO_CALLBACK_LOGGER.info(
             "ISPyB handler received event from read hardware"
         )
-        synchrotron_mode = doc["data"]["synchrotron-synchrotron_mode"]
+        synchrotron_mode = _data["synchrotron-synchrotron_mode"]
         assert isinstance(synchrotron_mode, SynchrotronMode)
         hwscan_data_collection_info = DataCollectionInfo(
-            undulator_gap1=doc["data"]["undulator-current_gap"],
+            undulator_gap1=_data["undulator-current_gap"],
             synchrotron_mode=synchrotron_mode.value,
-            slitgap_horizontal=doc["data"]["s4_slit_gaps-xgap"],
-            slitgap_vertical=doc["data"]["s4_slit_gaps-ygap"],
+            slitgap_horizontal=_data["s4_slit_gaps-xgap"],
+            slitgap_vertical=_data["s4_slit_gaps-ygap"],
         )
         hwscan_data_collection_info = _update_based_on_energy(
             doc, self.params.detector_params, hwscan_data_collection_info
         )
+
         hwscan_position_info = DataCollectionPositionInfo(
-            pos_x=float(doc["data"]["gonio-x"]),
-            pos_y=float(doc["data"]["gonio-y"]),
-            pos_z=float(doc["data"]["gonio-z"]),
+            pos_x=float(_data["gonio-x"]),
+            pos_y=float(_data.get("gonio-y")),
+            pos_z=float(_data["gonio-z"]),
         )
         scan_data_infos = self.populate_info_for_update(
             hwscan_data_collection_info, hwscan_position_info, self.params
@@ -160,18 +163,36 @@ def _handle_ispyb_hardware_read(self, doc) -> Sequence[ScanDataInfo]:
     def _handle_ispyb_transmission_flux_read(
         self, doc: Event
     ) -> Sequence[ScanDataInfo]:
+        _data = doc["data"]
+
         assert self.params
-        aperture = doc["data"]["aperture_scatterguard-selected_aperture"]
-        beamsize_x_mm = doc["data"]["beamsize-x_um"] / 1000
-        beamsize_y_mm = doc["data"]["beamsize-y_um"] / 1000
+        aperture = _data.get(
+            "aperture_scatterguard-selected_aperture", "Not implemented"
+        )
+        beamsize_x_mm = _data.get("beamsize-x_um", None)
+        if beamsize_x_mm:
+            beamsize_x_mm = beamsize_x_mm / 1000
+        beamsize_y_mm = _data.get("beamsize-y_um", None)
+        if beamsize_y_mm:
+            beamsize_y_mm = beamsize_y_mm / 1000
+        if not (beamsize_x_mm and beamsize_y_mm):
+            # VMXm don't have a beamsize device in dodal yet, they get beamsize sent in from GDA
+            try:
+                beamsize_x_mm = self.params.beam_size_x  # type: ignore
+                beamsize_y_mm = self.params.beam_size_y  # type: ignore
+            except Exception:
+                ISPYB_ZOCALO_CALLBACK_LOGGER.warning(
+                    "ISPyB callbacks couldn't get beamsize"
+                )
+
         hwscan_data_collection_info = DataCollectionInfo(
             beamsize_at_samplex=beamsize_x_mm,
             beamsize_at_sampley=beamsize_y_mm,
-            flux=doc["data"]["flux-flux_reading"],
-            detector_mode="ROI" if doc["data"]["eiger_cam_roi_mode"] else "FULL",
-            ispyb_detector_id=doc["data"]["eiger-ispyb_detector_id"],
+            flux=_data["flux-flux_reading"],
+            detector_mode="ROI" if _data["eiger_cam_roi_mode"] else "FULL",
+            ispyb_detector_id=_data["eiger-ispyb_detector_id"],
         )
-        if transmission := doc["data"]["attenuator-actual_transmission"]:
+        if transmission := _data["attenuator-actual_transmission"]:
             # Ispyb wants the transmission in a percentage, we use fractions
             hwscan_data_collection_info.transmission = transmission * 100
         hwscan_data_collection_info = _update_based_on_energy(