Add vmxm gridscan plan

src/mx_bluesky/beamlines/i02_1/device_setup_plans/setup_zebra.py

@@ -0,0 +1,45 @@
+import bluesky.plan_stubs as bps
+from dodal.devices.zebra.zebra import Zebra
+
+from mx_bluesky.common.parameters.constants import PlanGroupCheckpointConstants
+
+ZEBRA_STATUS_TIMEOUT = 30
+
+
+# Control Eiger from motion controller. Fast shutter is configured in GDA
+def setup_zebra_for_gridscan(
+    zebra: Zebra,
+    ttl_detector: int | None = None,
+    group=PlanGroupCheckpointConstants.SETUP_ZEBRA_FOR_GRIDSCAN,
+    wait=True,
+):
+    """
+    Assumes that the motion controller, as part of its gridscan PLC, will send triggers as required to the zebra's
+    IN1_TTL to control the detector. The fast shutter is configured in GDA, don't need to touch it in Bluesky for now.
+    """
+    ttl_detector = ttl_detector or zebra.mapping.outputs.TTL_EIGER
+    yield from bps.abs_set(
+        zebra.output.out_pvs[ttl_detector],
+        zebra.mapping.sources.IN1_TTL,
+    )
+    if wait:
+        yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT)
+
+
+def tidy_up_zebra_after_gridscan(
+    zebra: Zebra,
+    ttl_detector: int | None = None,
+    group=PlanGroupCheckpointConstants.TIDY_ZEBRA_AFTER_GRIDSCAN,
+    wait=False,
+):
+    """Revert zebra to state expected by GDA"""
+    ttl_detector = ttl_detector or zebra.mapping.outputs.TTL_EIGER
+
+    yield from bps.abs_set(
+        zebra.output.out_pvs[ttl_detector],
+        zebra.mapping.sources.OR1,
+        group=group,
+    )
+
+    if wait:
+        yield from bps.wait(group)

src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py

@@ -0,0 +1,138 @@
+from functools import partial
+
+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.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.undulator import BaseUndulator
+from dodal.devices.zebra.zebra import Zebra
+
+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.parameters.gridscan import SpecifiedTwoDGridScan
+from mx_bluesky.common.experiment_plans.common_flyscan_xray_centre_plan import (
+    BeamlineSpecificFGSFeatures,
+    common_flyscan_xray_centre,
+    construct_beamline_specific_fast_gridscan_features,
+)
+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 (
+    GridscanNexusFileCallback,
+)
+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.utils.log import LOGGER
+
+
+def create_gridscan_callbacks() -> tuple[
+    GridscanNexusFileCallback, GridscanISPyBCallback
+]:
+    return (
+        GridscanNexusFileCallback(param_type=SpecifiedTwoDGridScan),
+        GridscanISPyBCallback(
+            param_type=GenericGrid,
+            emit=ZocaloCallback(
+                PlanNameConstants.DO_FGS,
+                EnvironmentConstants.ZOCALO_ENV,
+                generate_start_info_from_omega_map,
+            ),
+        ),
+    )
+
+
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
+class FlyScanXRayCentreComposite(FlyScanEssentialDevices[GonioWithOmegaType]):
+    """All devices which are directly or indirectly required by this plan"""
+
+    zebra: Zebra
+    zebra_fast_grid_scan: ZebraFastGridScanTwoD
+    dcm: DoubleCrystalMonochromatorBase
+    attenuator: ReadOnlyAttenuator
+    flux: Flux
+    undulator: BaseUndulator
+    s4_slit_gaps: S4SlitGaps
+
+
+def construct_i02_1_specific_features(
+    fgs_composite: FlyScanXRayCentreComposite,
+    parameters: SpecifiedTwoDGridScan,
+) -> BeamlineSpecificFGSFeatures:
+    signals_to_read_pre_flyscan = [
+        fgs_composite.synchrotron.synchrotron_mode,
+        fgs_composite.gonio,
+        fgs_composite.dcm.energy_in_keV,
+        fgs_composite.undulator.current_gap,
+        fgs_composite.s4_slit_gaps,
+    ]
+
+    signals_to_read_during_collection = [
+        fgs_composite.attenuator.actual_transmission,
+        fgs_composite.flux.flux_reading,
+        fgs_composite.dcm.energy_in_keV,
+        fgs_composite.eiger.bit_depth,
+        fgs_composite.eiger.cam.roi_mode,
+        fgs_composite.eiger.ispyb_detector_id,
+    ]
+
+    return construct_beamline_specific_fast_gridscan_features(
+        partial(_zebra_triggering_setup),
+        partial(_tidy_plan, fgs_composite, group="flyscan_zebra_tidy", wait=True),
+        partial(
+            set_flyscan_params_plan,
+            fgs_composite.zebra_fast_grid_scan,
+            parameters.fast_gridscan_params,
+        ),
+        fgs_composite.zebra_fast_grid_scan,
+        signals_to_read_pre_flyscan,
+        signals_to_read_during_collection,  # type: ignore # See : https://github.com/bluesky/bluesky/issues/1809
+    )
+
+
+def _zebra_triggering_setup(fgs_composite: FlyScanXRayCentreComposite, _):
+    yield from setup_zebra_for_gridscan(fgs_composite.zebra)
+
+
+def _tidy_plan(
+    fgs_composite: FlyScanXRayCentreComposite, group, wait=True
+) -> MsgGenerator:
+    LOGGER.info("Tidying up Zebra")
+    yield from tidy_up_zebra_after_gridscan(fgs_composite.zebra)
+
+
+def i02_1_gridscan_plan(
+    parameters: SpecifiedTwoDGridScan,
+    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()
+
+    @bpp.subs_decorator(callbacks)
+    def decorated_flyscan_plan():
+        yield from common_flyscan_xray_centre(composite, parameters, beamline_specific)
+
+    yield from decorated_flyscan_plan()

src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py

@@ -1,4 +1,5 @@
 from dodal.devices.beamlines.i02_1.fast_grid_scan import ZebraGridScanParamsTwoD
+from pydantic import model_validator
 
 from mx_bluesky.common.parameters.components import SplitScan, WithOptionalEnergyChange
 from mx_bluesky.common.parameters.gridscan import SpecifiedGrids
@@ -24,5 +25,17 @@ def fast_gridscan_params(self) -> ZebraGridScanParamsTwoD:
             z1_start_mm=self.z_starts_um[0] / 1000,
             set_stub_offsets=self._set_stub_offsets,
             transmission_fraction=0.5,
-            dwell_time_ms=self.exposure_time_s,
+            dwell_time_ms=self.exposure_time_s * 1000,
         )
+
+    @model_validator(mode="after")
+    def validate_y_axes(self):
+        _err_str = "must be length 1 for 2D scans"
+        if len(self.y_steps) != 1:
+            raise ValueError(f"{self.y_steps=} {_err_str}")
+        if len(self.y_step_sizes_um) != 1:
+            raise ValueError(f"{self.y_step_sizes_um=} {_err_str}")
+        if len(self.omega_starts_deg) != 1:
+            raise ValueError(f"{self.omega_starts_deg=} {_err_str}")
+
+        return self

src/mx_bluesky/common/experiment_plans/common_flyscan_xray_centre_plan.py

@@ -30,7 +30,9 @@
     FlyScanEssentialDevices,
     GonioWithOmegaType,
 )
-from mx_bluesky.common.parameters.gridscan import SpecifiedThreeDGridScan
+from mx_bluesky.common.parameters.gridscan import (
+    SpecifiedGrids,
+)
 from mx_bluesky.common.utils.exceptions import (
     SampleError,
 )
@@ -114,7 +116,7 @@ def construct_beamline_specific_fast_gridscan_features(
 
 def common_flyscan_xray_centre(
     composite: FlyScanEssentialDevices[GonioWithOmegaType],
-    parameters: SpecifiedThreeDGridScan,
+    parameters: SpecifiedGrids,
     beamline_specific: BeamlineSpecificFGSFeatures,
 ) -> MsgGenerator:
     """Main entry point of the MX-Bluesky x-ray centering flyscan
@@ -156,7 +158,7 @@ def _decorated_flyscan():
         @bpp.finalize_decorator(lambda: _overall_tidy())
         def run_gridscan_and_tidy(
             fgs_composite: FlyScanEssentialDevices[GonioWithOmegaType],
-            params: SpecifiedThreeDGridScan,
+            params: SpecifiedGrids,
             beamline_specific: BeamlineSpecificFGSFeatures,
         ) -> MsgGenerator:
             yield from beamline_specific.setup_trigger_plan(fgs_composite, parameters)
@@ -174,12 +176,13 @@ def run_gridscan_and_tidy(
 
 def run_gridscan(
     fgs_composite: FlyScanEssentialDevices[GonioWithOmegaType],
-    parameters: SpecifiedThreeDGridScan,
+    parameters: SpecifiedGrids,
     beamline_specific: BeamlineSpecificFGSFeatures,
 ):
-    # Currently gridscan only works for omega 0, see https://github.com/DiamondLightSource/mx-bluesky/issues/410
     with TRACER.start_span("moving_omega_to_0"):
-        yield from bps.abs_set(fgs_composite.gonio.omega, 0)
+        yield from bps.abs_set(
+            fgs_composite.gonio.omega, parameters.omega_starts_deg[0], wait=True
+        )
 
     with TRACER.start_span("ispyb_hardware_readings"):
         yield from beamline_specific.read_pre_flyscan_plan()

src/mx_bluesky/common/parameters/constants.py

@@ -144,6 +144,8 @@ class PlanGroupCheckpointConstants:
     READY_FOR_OAV = "ready_for_oav"
     PREPARE_APERTURE = "prepare_aperture"
     SETUP_ZEBRA_FOR_ROTATION = "setup_zebra_for_rotation"
+    SETUP_ZEBRA_FOR_GRIDSCAN = "setup_zebra_for_gridscan"
+    TIDY_ZEBRA_AFTER_GRIDSCAN = "tidy_zebra_after_gridscan"
 
 
 # Eventually replace below with https://github.com/DiamondLightSource/mx-bluesky/issues/798

src/mx_bluesky/common/parameters/gridscan.py

@@ -4,7 +4,7 @@
 from typing import Annotated, Generic, TypeVar
 
 from dodal.devices.aperturescatterguard import ApertureValue
-from dodal.devices.detector.det_dim_constants import EIGER2_X_9M_SIZE, EIGER2_X_16M_SIZE
+from dodal.devices.detector.det_dim_constants import EIGER2_X_4M_SIZE, EIGER2_X_16M_SIZE
 from dodal.devices.detector.detector import DetectorParams
 from dodal.devices.fast_grid_scan import (
     GridScanParamsCommon,
@@ -32,7 +32,7 @@
 )
 
 DETECTOR_SIZE_PER_BEAMLINE = {
-    "i02-1": EIGER2_X_9M_SIZE,
+    "i02-1": EIGER2_X_4M_SIZE,
     "dev": EIGER2_X_16M_SIZE,
     "i03": EIGER2_X_16M_SIZE,
     "i04": EIGER2_X_16M_SIZE,
@@ -251,6 +251,9 @@ def validate_y_and_z_axes(self):
             raise ValueError(f"{self.y_steps=} {_err_str}")
         if len(self.y_step_sizes_um) != 2:
             raise ValueError(f"{self.y_step_sizes_um=} {_err_str}")
+        if len(self.omega_starts_deg) != 2:
+            raise ValueError(f"{self.omega_starts_deg=} {_err_str}")
+
         return self
 
     @property