mx-bluesky 1.5.8__py3-none-any.whl → 1.5.10__py3-none-any.whl

mx_bluesky/common/device_setup_plans/setup_zebra_and_shutter.py

@@ -0,0 +1,245 @@
+from typing import Protocol, runtime_checkable
+
+import bluesky.plan_stubs as bps
+from bluesky.utils import MsgGenerator
+from dodal.devices.zebra.zebra import (
+    ArmDemand,
+    EncEnum,
+    I03Axes,
+    RotationDirection,
+    Zebra,
+)
+from dodal.devices.zebra.zebra_controlled_shutter import (
+    ZebraShutter,
+    ZebraShutterControl,
+)
+
+from mx_bluesky.common.parameters.constants import ZEBRA_STATUS_TIMEOUT
+from mx_bluesky.common.utils.log import LOGGER
+
+"""Plans in this file will work as intended if the zebra has the following configuration:
+- A fast shutter is connected through TTL inputs from the Zebra.
+- When the zebra shutter is set to auto mode, the IOC sets the Zebra's SOFT_IN1 signal high.
+- When the zebra shutter is set to manual mode, the IOC sets the Zebra's SOFT_IN1 signal low.
+"""
+
+
+@runtime_checkable
+class GridscanSetupDevices(Protocol):
+    zebra: Zebra
+    sample_shutter: ZebraShutter
+
+
+def setup_zebra_for_gridscan(
+    composite: GridscanSetupDevices,  # XRC gridscan's generic trigger setup expects a composite rather than individual devices
+    group="setup_zebra_for_gridscan",
+    wait=True,
+    ttl_input_for_detector_to_use: None | int = None,
+) -> MsgGenerator:
+    """
+    Configure the zebra for an MX XRC gridscan by allowing the zebra to trigger the fast shutter and detector via signals
+    sent from the motion controller.
+
+    Args:
+        composite: Composite device containing a zebra and zebra shutter
+        group: Bluesky group to use when waiting on completion
+        wait: If true, block until completion
+        ttl_input_for_detector_to_use: If the zebra isn't using the TTL_DETECTOR zebra input, manually
+        specify which TTL input is being used for the desired detector
+
+    This plan assumes that the motion controller, as part of its gridscan PLC, will send triggers as required to the zebra's
+    IN4_TTL and IN3_TTL to control the fast_shutter and detector respectively
+
+    """
+    zebra = composite.zebra
+    ttl_detector = ttl_input_for_detector_to_use or zebra.mapping.outputs.TTL_DETECTOR
+    # Set shutter to automatic and to trigger via motion controller GPIO signal (IN4_TTL)
+    yield from configure_zebra_and_shutter_for_auto_shutter(
+        zebra, composite.sample_shutter, zebra.mapping.sources.IN4_TTL, group=group
+    )
+
+    yield from bps.abs_set(
+        zebra.output.out_pvs[ttl_detector],
+        zebra.mapping.sources.IN3_TTL,
+        group=group,
+    )
+
+    if wait:
+        yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT)
+
+
+def set_shutter_auto_input(zebra: Zebra, input: int, group="set_shutter_trigger"):
+    """Set the signal that controls the shutter. We use the second input to the
+    Zebra's AND_GATE_FOR_AUTO_SHUTTER for this input. ZebraShutter control mode must be in auto for this input to take control
+
+    For more details see the ZebraShutter device."""
+    auto_gate = zebra.mapping.AND_GATE_FOR_AUTO_SHUTTER
+    auto_shutter_control = zebra.logic_gates.and_gates[auto_gate]
+    yield from bps.abs_set(auto_shutter_control.sources[2], input, group)
+
+
+def configure_zebra_and_shutter_for_auto_shutter(
+    zebra: Zebra, zebra_shutter: ZebraShutter, input: int, group="use_automatic_shutter"
+):
+    """Set the shutter to auto mode, and configure the zebra to trigger the shutter on
+    an input source. For the input, use one of the source constants in zebra.py
+
+    When the shutter is in auto/manual, logic in EPICS sets the Zebra's
+    SOFT_IN1 to low/high respectively. The Zebra's AND_GATE_FOR_AUTO_SHUTTER should be used to control the shutter while in auto mode.
+    To do this, we need (AND_GATE_FOR_AUTO_SHUTTER = SOFT_IN1 AND input), where input is the zebra signal we want to control the shutter when in auto mode.
+    """
+
+    # Set shutter to auto mode
+    yield from bps.abs_set(
+        zebra_shutter.control_mode, ZebraShutterControl.AUTO, group=group
+    )
+
+    auto_gate = zebra.mapping.AND_GATE_FOR_AUTO_SHUTTER
+
+    # Set first input of AND_GATE_FOR_AUTO_SHUTTER to SOFT_IN1, which is high when shutter is in auto mode
+    # Note the Zebra should ALWAYS be setup this way. See https://github.com/DiamondLightSource/mx-bluesky/issues/551
+    yield from bps.abs_set(
+        zebra.logic_gates.and_gates[auto_gate].sources[1],
+        zebra.mapping.sources.SOFT_IN1,
+        group=group,
+    )
+
+    # Set the second input of AND_GATE_FOR_AUTO_SHUTTER to the requested zebra input source
+    yield from set_shutter_auto_input(zebra, input, group=group)
+
+
+def tidy_up_zebra_after_gridscan(
+    zebra: Zebra,
+    zebra_shutter: ZebraShutter,
+    group="tidy_up_zebra_after_gridscan",
+    wait=True,
+    ttl_input_for_detector_to_use: int | None = None,
+) -> MsgGenerator:
+    """
+    Set the zebra back to a state which is expected by GDA.
+
+    Args:
+        zebra: Zebra device.
+        zebra_shutter: Zebra shutter device.
+        group: Bluesky group to use when waiting on completion.
+        wait: If true, block until completion.
+        ttl_input_for_detector_to_use: If the zebra isn't using the TTL_DETECTOR zebra input, manually
+        specify which TTL input is being used for the desired detector.
+    """
+
+    LOGGER.info("Tidying up Zebra")
+
+    ttl_detector = ttl_input_for_detector_to_use or zebra.mapping.outputs.TTL_DETECTOR
+
+    yield from bps.abs_set(
+        zebra.output.out_pvs[ttl_detector],
+        zebra.mapping.sources.PC_PULSE,
+        group=group,
+    )
+    yield from bps.abs_set(
+        zebra_shutter.control_mode, ZebraShutterControl.MANUAL, group=group
+    )
+    yield from set_shutter_auto_input(zebra, zebra.mapping.sources.PC_GATE, group=group)
+
+    if wait:
+        yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT)
+
+
+def setup_zebra_for_rotation(
+    zebra: Zebra,
+    zebra_shutter: ZebraShutter,
+    axis: EncEnum = I03Axes.OMEGA,
+    start_angle: float = 0,
+    scan_width: float = 360,
+    shutter_opening_deg: float = 2.5,
+    shutter_opening_s: float = 0.04,
+    direction: RotationDirection = RotationDirection.POSITIVE,
+    group: str = "setup_zebra_for_rotation",
+    wait: bool = True,
+    ttl_input_for_detector_to_use: int | None = None,
+):
+    """Set up the Zebra to collect a rotation dataset. Any plan using this is
+    responsible for setting the smargon velocity appropriately so that the desired
+    image width is achieved with the exposure time given here.
+
+    Parameters:
+        zebra:              The zebra device to use
+        axis:               Encoder enum representing which axis to use for position
+                            compare. Currently always omega.
+        start_angle:        Position at which the scan should begin, in degrees.
+        scan_width:         Total angle through which to collect, in degrees.
+        shutter_opening_deg:How many degrees of rotation it takes for the fast shutter
+                            to open. Increases the gate width.
+        shutter_opening_s:  How many seconds it takes for the fast shutter to open. The
+                            detector pulse is delayed after the shutter signal by this
+                            amount.
+        direction:          RotationDirection enum for positive or negative.
+                            Defaults to Positive.
+        group:              A name for the group of statuses generated
+        wait:               Block until all the settings have completed
+        ttl_input_for_detector_to_use: If the zebra isn't using the TTL_DETECTOR zebra input,
+                            manually specify which TTL input is being used for the desired detector
+    """
+
+    ttl_detector = ttl_input_for_detector_to_use or zebra.mapping.outputs.TTL_DETECTOR
+
+    if not isinstance(direction, RotationDirection):
+        raise ValueError(
+            "Disallowed rotation direction provided to Zebra setup plan. "
+            "Use RotationDirection.POSITIVE or RotationDirection.NEGATIVE."
+        )
+    yield from bps.abs_set(zebra.pc.dir, direction.value, group=group)
+    LOGGER.info("ZEBRA SETUP: START")
+    # Set gate start, adjust for shutter opening time if necessary
+    LOGGER.info(f"ZEBRA SETUP: degrees to adjust for shutter = {shutter_opening_deg}")
+    LOGGER.info(f"ZEBRA SETUP: start angle start: {start_angle}")
+    LOGGER.info(f"ZEBRA SETUP: start angle adjusted, gate start set to: {start_angle}")
+    yield from bps.abs_set(zebra.pc.gate_start, start_angle, group=group)
+    # set gate width to total width
+    yield from bps.abs_set(
+        zebra.pc.gate_width, scan_width + shutter_opening_deg, group=group
+    )
+    LOGGER.info(
+        f"Pulse start set to shutter open time, set to: {abs(shutter_opening_s)}"
+    )
+    yield from bps.abs_set(zebra.pc.pulse_start, abs(shutter_opening_s), group=group)
+    # Set gate position to be angle of interest
+    yield from bps.abs_set(zebra.pc.gate_trigger, axis.value, group=group)
+    # Set shutter to automatic and to trigger via PC_GATE
+    yield from configure_zebra_and_shutter_for_auto_shutter(
+        zebra, zebra_shutter, zebra.mapping.sources.PC_GATE, group=group
+    )
+    # Trigger the detector with a pulse
+    yield from bps.abs_set(
+        zebra.output.out_pvs[ttl_detector],
+        zebra.mapping.sources.PC_PULSE,
+        group=group,
+    )
+
+    LOGGER.info(f"ZEBRA SETUP: END - {'' if wait else 'not'} waiting for completion")
+    if wait:
+        yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT)
+
+
+def tidy_up_zebra_after_rotation_scan(
+    zebra: Zebra,
+    zebra_shutter: ZebraShutter,
+    group="tidy_up_zebra_after_rotation",
+    wait=True,
+):
+    """
+    Set the zebra back to a state which is expected by GDA.
+
+    Args:
+        zebra: Zebra device.
+        zebra_shutter: Zebra shutter device.
+        group: Bluesky group to use when waiting on completion.
+        wait: If true, block until completion.
+    """
+
+    yield from bps.abs_set(zebra.pc.arm, ArmDemand.DISARM, group=group)
+    yield from bps.abs_set(
+        zebra_shutter.control_mode, ZebraShutterControl.MANUAL, group=group
+    )
+    if wait:
+        yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT)

mx_bluesky/common/device_setup_plans/xbpm_feedback.py

@@ -16,7 +16,7 @@ def unpause_xbpm_feedback_and_set_transmission_to_1(
                                       the beam in position
         attenuator (BinaryFilterAttenuator): The attenuator used to set transmission
     """
-    yield from bps.mv(xbpm_feedback.pause_feedback, Pause.RUN, attenuator, 1.0)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+    yield from bps.mv(xbpm_feedback.pause_feedback, Pause.RUN, attenuator, 1.0)
 
 
 def check_and_pause_feedback(
@@ -35,11 +35,11 @@ def check_and_pause_feedback(
                                                turning XBPM feedback off.
 
     """
-    yield from bps.mv(attenuator, 1.0)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+    yield from bps.mv(attenuator, 1.0)
     LOGGER.info("Waiting for XBPM feedback to be stable")
     yield from bps.trigger(xbpm_feedback, wait=True)
     LOGGER.info(
         f"XPBM feedback in position, pausing and setting transmission to {desired_transmission_fraction}"
     )
-    yield from bps.mv(xbpm_feedback.pause_feedback, Pause.PAUSE)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
-    yield from bps.mv(attenuator, desired_transmission_fraction)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+    yield from bps.mv(xbpm_feedback.pause_feedback, Pause.PAUSE)
+    yield from bps.mv(attenuator, desired_transmission_fraction)

mx_bluesky/common/experiment_plans/common_flyscan_xray_centre_plan.py

@@ -8,11 +8,12 @@ import bluesky.plan_stubs as bps
 import bluesky.preprocessors as bpp
 import numpy as np
 from bluesky.protocols import Readable
-from bluesky.utils import MsgGenerator
+from bluesky.utils import FailedStatus, MsgGenerator
 from dodal.common.beamlines.commissioning_mode import read_commissioning_mode
 from dodal.devices.fast_grid_scan import (
     FastGridScanCommon,
     FastGridScanThreeD,
+    GridScanInvalidError,
 )
 from dodal.devices.zocalo import ZocaloResults
 from dodal.devices.zocalo.zocalo_results import (
@@ -270,14 +271,18 @@ def run_gridscan(
         yield from beamline_specific.read_pre_flyscan_plan()
 
     LOGGER.info("Setting fgs params")
-    yield from beamline_specific.set_flyscan_params_plan()
 
-    LOGGER.info("Waiting for gridscan validity check")
-    yield from wait_for_gridscan_valid(beamline_specific.fgs_motors)
+    try:
+        yield from beamline_specific.set_flyscan_params_plan()
+    except FailedStatus as e:
+        if isinstance(e.__cause__, GridScanInvalidError):
+            raise SampleException(
+                "Scan invalid - gridscan not valid for detected pin position"
+            ) from e
 
     LOGGER.info("Waiting for arming to finish")
     yield from bps.wait(PlanGroupCheckpointConstants.GRID_READY_FOR_DC)
-    yield from bps.stage(fgs_composite.eiger)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+    yield from bps.stage(fgs_composite.eiger, wait=True)
 
     yield from kickoff_and_complete_gridscan(
         beamline_specific.fgs_motors,
@@ -293,23 +298,6 @@ def run_gridscan(
         yield from bps.abs_set(beamline_specific.fgs_motors.z_steps, 0, wait=False)
 
 
-def wait_for_gridscan_valid(fgs_motors: FastGridScanCommon, timeout=0.5):
-    LOGGER.info("Waiting for valid fgs_params")
-    SLEEP_PER_CHECK = 0.1
-    times_to_check = int(timeout / SLEEP_PER_CHECK)
-    for _ in range(times_to_check):
-        scan_invalid = yield from bps.rd(fgs_motors.scan_invalid)
-        pos_counter = yield from bps.rd(fgs_motors.position_counter)
-        LOGGER.debug(
-            f"Scan invalid: {scan_invalid} and position counter: {pos_counter}"
-        )
-        if not scan_invalid and pos_counter == 0:
-            LOGGER.info("Gridscan scan valid and position counter reset")
-            return
-        yield from bps.sleep(SLEEP_PER_CHECK)
-    raise SampleException("Scan invalid - pin too long/short/bent and out of range")
-
-
 def _xrc_result_in_boxes_to_result_in_mm(
     xrc_result: XrcResult, parameters: SpecifiedThreeDGridScan
 ) -> XRayCentreResult:

mx_bluesky/common/experiment_plans/inner_plans/do_fgs.py

@@ -94,14 +94,16 @@ def kickoff_and_complete_gridscan(
         md={
             "subplan_name": plan_name,
             "omega_to_scan_spec": {
-                GridscanPlane.OMEGA_XY: scan_points[0],
-                GridscanPlane.OMEGA_XZ: scan_points[1],
+                # These have to be cast to strings due to a bug in orsjon. See
+                # https://github.com/ijl/orjson/issues/414
+                str(GridscanPlane.OMEGA_XY): scan_points[0],
+                str(GridscanPlane.OMEGA_XZ): scan_points[1],
             },
         }
     )
     @bpp.contingency_decorator(
         except_plan=lambda e: (yield from bps.stop(detector)),  # type: ignore # Fix types in ophyd-async (https://github.com/DiamondLightSource/mx-bluesky/issues/855)
-        else_plan=lambda: (yield from bps.unstage(detector)),
+        else_plan=lambda: (yield from bps.unstage(detector, wait=True)),
     )
     def _decorated_do_fgs():
         yield from _wait_for_zocalo_to_stage_then_do_fgs(

mx_bluesky/common/experiment_plans/inner_plans/read_hardware.py

@@ -4,7 +4,7 @@ import bluesky.plan_stubs as bps
 from bluesky.protocols import Readable
 from dodal.devices.aperturescatterguard import ApertureScatterguard
 from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator
-from dodal.devices.common_dcm import BaseDCM
+from dodal.devices.common_dcm import DoubleCrystalMonochromator
 from dodal.devices.eiger import EigerDetector
 from dodal.devices.flux import Flux
 from dodal.devices.s4_slit_gaps import S4SlitGaps
@@ -43,7 +43,7 @@ def standard_read_hardware_pre_collection(
     undulator: Undulator,
     synchrotron: Synchrotron,
     s4_slit_gaps: S4SlitGaps,
-    dcm: BaseDCM,
+    dcm: DoubleCrystalMonochromator,
     smargon: Smargon,
 ):
     LOGGER.info("Reading status of beamline for callbacks, pre collection.")
@@ -63,7 +63,7 @@ def standard_read_hardware_during_collection(
     aperture_scatterguard: ApertureScatterguard,
     attenuator: BinaryFilterAttenuator,
     flux: Flux,
-    dcm: BaseDCM,
+    dcm: DoubleCrystalMonochromator,
     detector: EigerDetector,
 ):
     signals_to_read_during_collection = [

mx_bluesky/common/experiment_plans/oav_grid_detection_plan.py

@@ -10,7 +10,7 @@ from bluesky.utils import MsgGenerator
 from dodal.devices.oav.oav_detector import OAV
 from dodal.devices.oav.pin_image_recognition import PinTipDetection
 from dodal.devices.oav.pin_image_recognition.utils import NONE_VALUE
-from dodal.devices.oav.utils import PinNotFoundException, wait_for_tip_to_be_found
+from dodal.devices.oav.utils import PinNotFoundError, wait_for_tip_to_be_found
 from dodal.devices.smargon import Smargon
 
 from mx_bluesky.common.device_setup_plans.setup_oav import (
@@ -108,7 +108,7 @@ def grid_detection_plan(
         yield from bps.sleep(HardwareConstants.OAV_REFRESH_DELAY)
 
         tip_x_px, tip_y_px = yield from catch_exception_and_warn(
-            PinNotFoundException, wait_for_tip_to_be_found, pin_tip_detection
+            PinNotFoundError, wait_for_tip_to_be_found, pin_tip_detection
         )
 
         LOGGER.info(f"Tip is at x,y: {tip_x_px},{tip_y_px}")

mx_bluesky/common/external_interaction/callbacks/common/ispyb_mapping.py

@@ -50,9 +50,10 @@ def populate_remaining_data_collection_info(
     data_collection_info.xbeam = beam_position[0]
     data_collection_info.ybeam = beam_position[1]
     data_collection_info.start_time = get_current_time_string()
-    # temporary file template until nxs filewriting is integrated and we can use
-    # that file name
-    data_collection_info.file_template = f"{params.detector_params.prefix}_{data_collection_info.data_collection_number}_master.h5"
+    if data_collection_info.data_collection_number is not None:
+        # Do not write the file template if we don't have sufficient information - for gridscans we  may not
+        # know the data collection number until later
+        data_collection_info.file_template = f"{params.detector_params.prefix}_{data_collection_info.data_collection_number}_master.h5"
     return data_collection_info
 
 

mx_bluesky/common/external_interaction/callbacks/xray_centre/ispyb_callback.py

@@ -105,6 +105,7 @@ class GridscanISPyBCallback(BaseISPyBCallback):
         self._start_of_fgs_uid: str | None = None
         self._processing_start_time: float | None = None
         self._grid_plane_to_id_map: dict[GridscanPlane, int] = {}
+        self._grid_plane_to_width_map: dict[GridscanPlane, int] = {}
         self.data_collection_group_info: DataCollectionGroupInfo | None
 
     def activity_gated_start(self, doc: RunStart):
@@ -131,9 +132,7 @@ class GridscanISPyBCallback(BaseISPyBCallback):
                     data_collection_info=populate_remaining_data_collection_info(
                         "MX-Bluesky: Xray centring 1 -",
                         None,
-                        DataCollectionInfo(
-                            data_collection_number=self.params.detector_params.run_number,
-                        ),
+                        DataCollectionInfo(),
                         self.params,
                     ),
                 ),
@@ -141,11 +140,7 @@ class GridscanISPyBCallback(BaseISPyBCallback):
                     data_collection_info=populate_remaining_data_collection_info(
                         "MX-Bluesky: Xray centring 2 -",
                         None,
-                        DataCollectionInfo(
-                            data_collection_number=(
-                                self.params.detector_params.run_number + 1
-                            ),
-                        ),
+                        DataCollectionInfo(),
                         self.params,
                     )
                 ),
@@ -192,6 +187,15 @@ class GridscanISPyBCallback(BaseISPyBCallback):
         assert self.params, "ISPyB handler didn't receive parameters!"
         assert self.data_collection_group_info, "No data collection group"
         data = doc["data"]
+        omega = doc["data"]["smargon-omega"]
+        grid_plane = _smargon_omega_to_xyxz_plane(omega)
+        ISPYB_ZOCALO_CALLBACK_LOGGER.info(
+            f"Generating dc info for gridplane {grid_plane}, omega {omega}"
+        )
+        data_collection_number = self.data_collection_number_from_gridplane(grid_plane)
+        file_template = (
+            f"{self.params.detector_params.prefix}_{data_collection_number}_master.h5"
+        )
         data_collection_info = DataCollectionInfo(
             xtal_snapshot1=data.get("oav-grid_snapshot-last_path_full_overlay"),
             xtal_snapshot2=data.get("oav-grid_snapshot-last_path_outer"),
@@ -200,6 +204,8 @@ class GridscanISPyBCallback(BaseISPyBCallback):
                 data["oav-grid_snapshot-num_boxes_x"]
                 * data["oav-grid_snapshot-num_boxes_y"]
             ),
+            data_collection_number=data_collection_number,
+            file_template=file_template,
         )
         microns_per_pixel_x = data["oav-microns_per_pixel_x"]
         microns_per_pixel_y = data["oav-microns_per_pixel_y"]
@@ -219,21 +225,22 @@ class GridscanISPyBCallback(BaseISPyBCallback):
             data_collection_grid_info
         )
 
-        if self.data_collection_group_info.comments:
-            self.data_collection_group_info.comments += (
-                f"by {data_collection_grid_info.steps_y}."
-            )
-        else:
-            self.data_collection_group_info.comments = (
-                f"Diffraction grid scan of "
-                f"{data_collection_grid_info.steps_x} "
-                f"by {data_collection_grid_info.steps_y} "
-            )
-
+        # Snapshots may be triggered in a different order to gridscans, so save
+        # the mapping to the data collection id in order to trigger Zocalo correctly.
         data_collection_id = self.ispyb_ids.data_collection_ids[
-            self._oav_snapshot_event_idx
+            0 if grid_plane == GridscanPlane.OMEGA_XY else 1
         ]
-        self._populate_axis_info(data_collection_info, doc["data"]["smargon-omega"])
+        self._grid_plane_to_id_map[grid_plane] = data_collection_id
+        self._grid_plane_to_width_map[grid_plane] = data_collection_grid_info.steps_y
+
+        y_steps = self._grid_plane_to_width_map.get(GridscanPlane.OMEGA_XY, "_")
+        z_steps = self._grid_plane_to_width_map.get(GridscanPlane.OMEGA_XZ, "_")
+        self.data_collection_group_info.comments = (
+            f"Diffraction grid scan of {data_collection_grid_info.steps_x} by "
+            f"{y_steps} by {z_steps}."
+        )
+
+        self._populate_axis_info(data_collection_info, omega)
 
         scan_data_info = ScanDataInfo(
             data_collection_info=data_collection_info,
@@ -243,10 +250,6 @@ class GridscanISPyBCallback(BaseISPyBCallback):
         ISPYB_ZOCALO_CALLBACK_LOGGER.info(
             "Updating ispyb data collection after oav snapshot."
         )
-        grid_plane = _smargon_omega_to_xyxz_plane(doc["data"]["smargon-omega"])
-        # Snapshots may be triggered in a different order to gridscans, so save
-        # the mapping to the data collection id in order to trigger Zocalo correctly.
-        self._grid_plane_to_id_map[grid_plane] = data_collection_id
 
         self._oav_snapshot_event_idx += 1
         return [scan_data_info]
@@ -315,6 +318,7 @@ class GridscanISPyBCallback(BaseISPyBCallback):
             )
             self.data_collection_group_info = None
             self._grid_plane_to_id_map.clear()
+            self._grid_plane_to_width_map.clear()
             return super().activity_gated_stop(doc)
         return self.tag_doc(doc)
 
@@ -325,6 +329,11 @@ class GridscanISPyBCallback(BaseISPyBCallback):
             doc["grid_plane_to_id_map"] = self._grid_plane_to_id_map
         return doc  # type: ignore
 
+    def data_collection_number_from_gridplane(self, plane) -> int:
+        assert self.params
+        base_number = self.params.detector_params.run_number
+        return base_number if plane == GridscanPlane.OMEGA_XY else base_number + 1
+
 
 def generate_start_info_from_omega_map() -> ZocaloInfoGenerator:
     """

mx_bluesky/common/parameters/device_composites.py

@@ -4,7 +4,7 @@ from dodal.devices.aperturescatterguard import (
 )
 from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator
 from dodal.devices.backlight import Backlight
-from dodal.devices.common_dcm import BaseDCM
+from dodal.devices.common_dcm import DoubleCrystalMonochromator
 from dodal.devices.detector.detector_motion import DetectorMotion
 from dodal.devices.eiger import EigerDetector
 from dodal.devices.fast_grid_scan import (
@@ -51,7 +51,7 @@ class GridDetectThenXRayCentreComposite(FlyScanEssentialDevices):
     attenuator: BinaryFilterAttenuator
     backlight: Backlight
     beamstop: Beamstop
-    dcm: BaseDCM
+    dcm: DoubleCrystalMonochromator
     detector_motion: DetectorMotion
     zebra_fast_grid_scan: ZebraFastGridScanThreeD
     flux: Flux

mx_bluesky/hyperion/device_setup_plans/setup_zebra.py

@@ -1,120 +1,22 @@
 import bluesky.plan_stubs as bps
 from dodal.devices.zebra.zebra import (
     ArmDemand,
-    EncEnum,
-    I03Axes,
-    RotationDirection,
     Zebra,
 )
 from dodal.devices.zebra.zebra_controlled_shutter import (
     ZebraShutter,
-    ZebraShutterControl,
 )
 
-from mx_bluesky.common.parameters.constants import ZEBRA_STATUS_TIMEOUT
-from mx_bluesky.common.utils.log import LOGGER
-from mx_bluesky.phase1_zebra.device_setup_plans.setup_zebra import (
+from mx_bluesky.common.device_setup_plans.setup_zebra_and_shutter import (
     configure_zebra_and_shutter_for_auto_shutter,
 )
+from mx_bluesky.common.parameters.constants import ZEBRA_STATUS_TIMEOUT
 
 
 def arm_zebra(zebra: Zebra):
     yield from bps.abs_set(zebra.pc.arm, ArmDemand.ARM, wait=True)
 
 
-def tidy_up_zebra_after_rotation_scan(
-    zebra: Zebra,
-    zebra_shutter: ZebraShutter,
-    group="tidy_up_zebra_after_rotation",
-    wait=True,
-):
-    yield from bps.abs_set(zebra.pc.arm, ArmDemand.DISARM, group=group)
-    yield from bps.abs_set(
-        zebra_shutter.control_mode, ZebraShutterControl.MANUAL, group=group
-    )
-    if wait:
-        yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT)
-
-
-def setup_zebra_for_rotation(
-    zebra: Zebra,
-    zebra_shutter: ZebraShutter,
-    axis: EncEnum = I03Axes.OMEGA,
-    start_angle: float = 0,
-    scan_width: float = 360,
-    shutter_opening_deg: float = 2.5,
-    shutter_opening_s: float = 0.04,
-    direction: RotationDirection = RotationDirection.POSITIVE,
-    group: str = "setup_zebra_for_rotation",
-    wait: bool = True,
-):
-    """Set up the Zebra to collect a rotation dataset. Any plan using this is
-    responsible for setting the smargon velocity appropriately so that the desired
-    image width is achieved with the exposure time given here.
-
-    Parameters:
-        zebra:              The zebra device to use
-        axis:               I03 axes enum representing which axis to use for position
-                            compare. Currently always omega.
-        start_angle:        Position at which the scan should begin, in degrees.
-        scan_width:         Total angle through which to collect, in degrees.
-        shutter_opening_deg:How many degrees of rotation it takes for the fast shutter
-                            to open. Increases the gate width.
-        shutter_opening_s:  How many seconds it takes for the fast shutter to open. The
-                            detector pulse is delayed after the shutter signal by this
-                            amount.
-        direction:          RotationDirection enum for positive or negative.
-                            Defaults to Positive.
-        group:              A name for the group of statuses generated
-        wait:               Block until all the settings have completed
-    """
-
-    if not isinstance(direction, RotationDirection):
-        raise ValueError(
-            "Disallowed rotation direction provided to Zebra setup plan. "
-            "Use RotationDirection.POSITIVE or RotationDirection.NEGATIVE."
-        )
-    yield from bps.abs_set(zebra.pc.dir, direction.value, group=group)
-    LOGGER.info("ZEBRA SETUP: START")
-    # Set gate start, adjust for shutter opening time if necessary
-    LOGGER.info(f"ZEBRA SETUP: degrees to adjust for shutter = {shutter_opening_deg}")
-    LOGGER.info(f"ZEBRA SETUP: start angle start: {start_angle}")
-    LOGGER.info(f"ZEBRA SETUP: start angle adjusted, gate start set to: {start_angle}")
-    yield from bps.abs_set(zebra.pc.gate_start, start_angle, group=group)
-    # set gate width to total width
-    yield from bps.abs_set(
-        zebra.pc.gate_width, scan_width + shutter_opening_deg, group=group
-    )
-    LOGGER.info(
-        f"Pulse start set to shutter open time, set to: {abs(shutter_opening_s)}"
-    )
-    yield from bps.abs_set(zebra.pc.pulse_start, abs(shutter_opening_s), group=group)
-    # Set gate position to be angle of interest
-    yield from bps.abs_set(zebra.pc.gate_trigger, axis.value, group=group)
-    # Set shutter to automatic and to trigger via PC_GATE
-    yield from configure_zebra_and_shutter_for_auto_shutter(
-        zebra, zebra_shutter, zebra.mapping.sources.PC_GATE, group=group
-    )
-    # Trigger the detector with a pulse
-    yield from bps.abs_set(
-        zebra.output.out_pvs[zebra.mapping.outputs.TTL_DETECTOR],
-        zebra.mapping.sources.PC_PULSE,
-        group=group,
-    )
-    # Don't use the fluorescence detector
-    yield from bps.abs_set(
-        zebra.output.out_pvs[zebra.mapping.outputs.TTL_XSPRESS3],
-        zebra.mapping.sources.DISCONNECT,
-        group=group,
-    )
-    yield from bps.abs_set(
-        zebra.output.pulse_1.input, zebra.mapping.sources.DISCONNECT, group=group
-    )
-    LOGGER.info(f"ZEBRA SETUP: END - {'' if wait else 'not'} waiting for completion")
-    if wait:
-        yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT)
-
-
 def setup_zebra_for_panda_flyscan(
     zebra: Zebra,
     zebra_shutter: ZebraShutter,