(Don't review) Working branch for vmxm testing

docs/developer/hyperion/reference/gridscan.puml

@@ -97,13 +97,13 @@ class GridScanParamsCommon {
     z2_start_mm
 }
 class PandAGridScanParams
-class ZebraGridScanParams
+class ZebraGridScanParamsThreeD
 
 AbstractExperimentBase <|-- AbstractExperimentWithBeamParams
 AbstractExperimentWithBeamParams <|-- GridScanParamsCommon
 GridScanParamsCommon <|-- PandAGridScanParams
-GridScanParamsCommon <|-- ZebraGridScanParams
+GridScanParamsCommon <|-- ZebraGridScanParamsThreeD
 
-HyperionThreeDGridScan --> ZebraGridScanParams : generates
+HyperionThreeDGridScan --> ZebraGridScanParamsThreeD : generates
 HyperionThreeDGridScan --> PandAGridScanParams : generates
 @enduml

pyproject.toml

@@ -48,7 +48,7 @@ dependencies = [
     "ophyd >= 1.10.5",
     "ophyd-async >= 0.10.0a2",
     "bluesky >= 1.13.1",
-    "dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git",
+    "dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git@8ff1c172cad8c043e3d8318c05162c791f6e0b67",
 ]
 
 
@@ -115,7 +115,7 @@ typeCheckingMode = "standard"
 # Run pytest with all our checkers, and don't spam us with massive tracebacks on error
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"
-timeout = 1
+timeout = 60
 markers = [
     "dlstbx: marks tests as requiring dlstbx (deselect with '-m \"not dlstbx\"')",
     "skip_log_setup: marks tests so that loggers are not setup before the test.",

src/mx_bluesky/beamlines/i02_1/__init__.py

@@ -0,0 +1,5 @@
+from mx_bluesky.beamlines.i02_1.i02_1_flyscan_xray_centre_plan import (
+    i02_1_flyscan_xray_centre,
+)
+
+__all__ = ["i02_1_flyscan_xray_centre"]

src/mx_bluesky/beamlines/i02_1/constants.py

@@ -0,0 +1,7 @@
+from pydantic.dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class I02_1_Constants:
+    GRAYLOG_PORT = 12232
+    LOG_FILE_NAME = "i02_1_bluesky.log"

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

@@ -0,0 +1,28 @@
+import bluesky.plan_stubs as bps
+from dodal.devices.zebra.zebra import Zebra
+
+ZEBRA_STATUS_TIMEOUT = 30
+
+
+# Control Eiger from motion controller
+def setup_zebra_for_xrc_flyscan(zebra: Zebra, group="setup_zebra_for_xrc", wait=True):
+    yield from bps.abs_set(
+        zebra.output.out_pvs[zebra.mapping.outputs.TTL_EIGER],
+        zebra.mapping.sources.IN1_TTL,
+    )
+    if wait:
+        yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT)
+
+
+# State of zebra expected by GDA
+def tidy_up_zebra_after_gridscan(
+    zebra: Zebra, group="tidy_up_vmxm_zebra_after_gridscan", wait=False
+):
+    yield from bps.abs_set(
+        zebra.output.out_pvs[zebra.mapping.outputs.TTL_EIGER],
+        zebra.mapping.sources.OR1,
+        group=group,
+    )
+
+    if wait:
+        yield from bps.wait(group)

src/mx_bluesky/beamlines/i02_1/i02_1_flyscan_xray_centre_plan.py

@@ -0,0 +1,109 @@
+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.eiger import EigerDetector
+from dodal.devices.fast_grid_scan import (
+    set_fast_grid_scan_params as set_flyscan_params_plan,
+)
+from dodal.devices.i02_1.fast_grid_scan import ZebraGridScanParamsTwoD
+from dodal.devices.i02_1.sample_motors import SampleMotors
+from dodal.devices.synchrotron import Synchrotron
+from dodal.devices.zebra.zebra import Zebra
+
+from mx_bluesky.beamlines.i02_1.device_setup_plans.setup_zebra import (
+    setup_zebra_for_xrc_flyscan,
+    tidy_up_zebra_after_gridscan,
+)
+from mx_bluesky.common.experiment_plans.common_flyscan_xray_centre_plan import (
+    BeamlineSpecificFGSFeatures,
+    FlyScanBaseComposite,
+    common_flyscan_xray_centre,
+    construct_beamline_specific_FGS_features,
+)
+from mx_bluesky.common.external_interaction.callbacks.common.callback_util import (
+    create_gridscan_callbacks,
+)
+from mx_bluesky.common.parameters.device_composites import SampleStageWithOmega
+from mx_bluesky.common.parameters.gridscan import SpecifiedThreeDGridScan
+from mx_bluesky.common.utils.log import LOGGER
+
+
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
+class FlyScanXRayCentreComposite(
+    FlyScanBaseComposite[ZebraGridScanParamsTwoD, SampleStageWithOmega]
+):
+    """All devices which are directly or indirectly required by this plan"""
+
+    zebra: Zebra
+
+
+def construct_i02_1_specific_features(
+    fgs_composite: FlyScanXRayCentreComposite,
+    parameters: SpecifiedThreeDGridScan,
+) -> BeamlineSpecificFGSFeatures:
+    signals_to_read_pre_flyscan = [
+        fgs_composite.synchrotron.synchrotron_mode,
+        fgs_composite.sample_stage,
+    ]
+    signals_to_read_during_collection = [
+        fgs_composite.eiger.bit_depth,
+    ]
+
+    return construct_beamline_specific_FGS_features(
+        _zebra_triggering_setup,
+        partial(_tidy_plan, group="flyscan_zebra_tidy", wait=True),
+        partial(
+            set_flyscan_params_plan,
+            fgs_composite.grid_scan,
+            parameters.FGS_params,
+        ),
+        fgs_composite.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_xrc_flyscan(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_flyscan_xray_centre(
+    parameters: SpecifiedThreeDGridScan,
+    eiger: EigerDetector = inject("eiger"),
+    zebra_fast_grid_scan: ZebraFastGridScanTwoD = inject("ZebraFastGridScanTwoD"),
+    synchrotron: Synchrotron = inject("synchrotron"),
+    zebra: Zebra = inject("zebra"),
+    sample_motors: SampleMotors = inject("sample_motors"),
+) -> MsgGenerator:
+    """BlueAPI entry point for XRC grid scans"""
+
+    # Composites have to be made this way until https://github.com/DiamondLightSource/dodal/issues/874
+    # is done and we can properly use composite devices in BlueAPI
+    composite = FlyScanXRayCentreComposite(
+        eiger,
+        synchrotron,
+        sample_motors,
+        zebra_fast_grid_scan,
+        zebra,
+    )
+
+    beamline_specific = construct_i02_1_specific_features(composite, parameters)
+
+    @bpp.subs_decorator(create_gridscan_callbacks())
+    def decorated_flyscan_plan():
+        yield from common_flyscan_xray_centre(composite, parameters, beamline_specific)
+
+    yield from decorated_flyscan_plan()

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

@@ -13,7 +13,7 @@
 from dodal.devices.detector.detector_motion import DetectorMotion
 from dodal.devices.eiger import EigerDetector
 from dodal.devices.fast_grid_scan import (
-    ZebraFastGridScan,
+    ZebraFastGridScanThreeD,
     set_fast_grid_scan_params,
 )
 from dodal.devices.flux import Flux
@@ -43,17 +43,10 @@
 from mx_bluesky.common.experiment_plans.oav_snapshot_plan import (
     setup_beamline_for_OAV,
 )
-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,
-)
-from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import (
-    GridscanNexusFileCallback,
+from mx_bluesky.common.external_interaction.callbacks.common.callback_util import (
+    create_gridscan_callbacks,
 )
 from mx_bluesky.common.parameters.constants import (
-    EnvironmentConstants,
     OavConstants,
     PlanGroupCheckpointConstants,
     PlanNameConstants,
@@ -80,7 +73,7 @@ def i04_grid_detect_then_xray_centre(
     backlight: Backlight = inject("backlight"),
     beamstop: Beamstop = inject("beamstop"),
     dcm: BaseDCM = inject("dcm"),
-    zebra_fast_grid_scan: ZebraFastGridScan = inject("zebra_fast_grid_scan"),
+    zebra_fast_grid_scan: ZebraFastGridScanThreeD = inject("zebra_fast_grid_scan"),
     flux: Flux = inject("flux"),
     oav: OAV = inject("oav"),
     pin_tip_detection: PinTipDetection = inject("pin_tip_detection"),
@@ -114,15 +107,14 @@ def i04_grid_detect_then_xray_centre(
     composite = GridDetectThenXRayCentreComposite(
         eiger,
         synchrotron,
-        zocalo,
         smargon,
+        zebra_fast_grid_scan,
         aperture_scatterguard,
         attenuator,
         backlight,
         beamstop,
         dcm,
         detector_motion,
-        zebra_fast_grid_scan,
         flux,
         oav,
         pin_tip_detection,
@@ -132,12 +124,13 @@ def i04_grid_detect_then_xray_centre(
         zebra,
         robot,
         sample_shutter,
+        zocalo,
     )
 
     def tidy_beamline_if_not_udc():
         if not udc:
             yield from get_ready_for_oav_and_close_shutter(
-                composite.smargon,
+                composite.sample_stage,
                 composite.backlight,
                 composite.aperture_scatterguard,
                 composite.detector_motion,
@@ -190,20 +183,6 @@ def get_ready_for_oav_and_close_shutter(
     yield from bps.wait(group)
 
 
-def create_gridscan_callbacks() -> tuple[
-    GridscanNexusFileCallback, GridscanISPyBCallback
-]:
-    return (
-        GridscanNexusFileCallback(param_type=SpecifiedThreeDGridScan),
-        GridscanISPyBCallback(
-            param_type=GridCommon,
-            emit=ZocaloCallback(
-                PlanNameConstants.DO_FGS, EnvironmentConstants.ZOCALO_ENV
-            ),
-        ),
-    )
-
-
 def construct_i04_specific_features(
     xrc_composite: GridDetectThenXRayCentreComposite,
     xrc_parameters: SpecifiedThreeDGridScan,
@@ -216,9 +195,7 @@ def construct_i04_specific_features(
         xrc_composite.synchrotron.synchrotron_mode,
         xrc_composite.s4_slit_gaps.xgap,
         xrc_composite.s4_slit_gaps.ygap,
-        xrc_composite.smargon.x,
-        xrc_composite.smargon.y,
-        xrc_composite.smargon.z,
+        xrc_composite.sample_stage,
         xrc_composite.dcm.energy_in_kev,
     ]
 
@@ -239,16 +216,15 @@ def construct_i04_specific_features(
     )
     set_flyscan_params_plan = partial(
         set_fast_grid_scan_params,
-        xrc_composite.zebra_fast_grid_scan,
+        xrc_composite.grid_scan,
         xrc_parameters.FGS_params,
     )
-    fgs_motors = xrc_composite.zebra_fast_grid_scan
+    fgs_motors = xrc_composite.grid_scan
     return construct_beamline_specific_FGS_features(
         setup_zebra_for_gridscan,
         tidy_plan,
         set_flyscan_params_plan,
         fgs_motors,
         signals_to_read_pre_flyscan,
         signals_to_read_during_collection,
-        get_xrc_results_from_zocalo=True,
     )