mx-bluesky 1.4.6__py3-none-any.whl → 1.4.8__py3-none-any.whl

mx_bluesky/common/parameters/components.py

@@ -5,7 +5,7 @@ from abc import abstractmethod
 from collections.abc import Sequence
 from enum import StrEnum
 from pathlib import Path
-from typing import Literal, SupportsInt, cast
+from typing import Literal, Self, SupportsInt, cast
 
 from dodal.devices.aperturescatterguard import ApertureValue
 from dodal.devices.detector import (
@@ -114,12 +114,32 @@ class MxBlueskyParameters(BaseModel):
 
 
 class WithSnapshot(BaseModel):
+    """
+    Configures how snapshot images are created.
+    Attributes:
+        snapshot_directory: Path to the directory where snapshot images will be stored
+        snapshot_omegas_deg: list of omega values at which snapshots will be taken. For
+            gridscans, this attribute is ignored.
+        use_grid_snapshots: This may be specified for rotation snapshots to speed up rotation
+            execution. If set to True then rotation snapshots are generated from the
+            previously captured grid snapshots. Otherwise they are captured using
+            freshly captured snapshots during the rotation plan.
+    """
+
     snapshot_directory: Path
     snapshot_omegas_deg: list[float] | None = None
+    use_grid_snapshots: bool = False
 
     @property
     def take_snapshots(self) -> bool:
-        return bool(self.snapshot_omegas_deg)
+        return bool(self.snapshot_omegas_deg) or self.use_grid_snapshots
+
+    @model_validator(mode="after")
+    def _validate_omegas_with_grid_snapshots(self) -> Self:
+        assert not self.use_grid_snapshots or self.snapshot_omegas_deg is None, (
+            "snapshot_omegas may not be specified with use_grid_snapshots"
+        )
+        return self
 
 
 class WithOptionalEnergyChange(BaseModel):

mx_bluesky/common/parameters/constants.py

@@ -55,9 +55,12 @@ class PlanNameConstants:
     ROBOT_LOAD_AND_SNAPSHOTS = "robot_load_and_snapshots"
     # Rotation scan
     ROTATION_MULTI = "multi_rotation_wrapper"
+    ROTATION_MULTI_OUTER = "multi_rotation_outer"
     ROTATION_OUTER = "rotation_scan_with_cleanup"
     ROTATION_MAIN = "rotation_scan_main"
     FLYSCAN_RESULTS = "xray_centre_results"
+    SET_ENERGY = "set_energy"
+    UNNAMED_RUN = "unnamed_run"
 
 
 @dataclass(frozen=True)
@@ -65,17 +68,12 @@ class EnvironmentConstants:
     ZOCALO_ENV = ZOCALO_ENV_FROM_DODAL
 
 
-@dataclass(frozen=True)
-class TriggerConstants:
-    ZOCALO = "trigger_zocalo_on"
-
-
 @dataclass(frozen=True)
 class HardwareConstants:
     OAV_REFRESH_DELAY = 0.3
     PANDA_FGS_RUN_UP_DEFAULT = 0.17
     CRYOJET_MARGIN_MM = 0.2
-    THAWING_TIME = 20
+    THAWING_TIME = 40
     TIP_OFFSET_UM = 0
 
     # Value quoted in https://www.dectris.com/en/detectors/x-ray-detectors/eiger2/eiger2-for-synchrotrons/eiger2-x/,
@@ -92,6 +90,7 @@ class GridscanParamConstants:
     OMEGA_1 = 0.0
     OMEGA_2 = 90.0
     PANDA_RUN_UP_DISTANCE_MM = 0.2
+    ZOCALO_MIN_TOTAL_COUNT_THRESHOLD = 3
 
 
 @dataclass(frozen=True)
@@ -123,6 +122,7 @@ class PlanGroupCheckpointConstants:
     ROTATION_READY_FOR_DC = "rotation_ready_for_data_collection"
     MOVE_GONIO_TO_START = "move_gonio_to_start"
     READY_FOR_OAV = "ready_for_oav"
+    PREPARE_APERTURE = "prepare_aperture"
 
 
 # Eventually replace below with https://github.com/DiamondLightSource/mx-bluesky/issues/798
@@ -134,16 +134,6 @@ class DeviceSettingsConstants:
     )
 
 
-@dataclass(frozen=True)
-class SimConstants:
-    BEAMLINE = "BL03S"
-    INSERTION_PREFIX = "SR03S"
-    # this one is for unit tests
-    ISPYB_CONFIG = "tests/test_data/test_config.cfg"
-    # this one is for system tests
-    DEV_ISPYB_DATABASE_CFG = "/dls_sw/dasc/mariadb/credentials/ispyb-hyperion-dev.cfg"
-
-
 class Actions(Enum):
     START = "start"
     STOP = "stop"

mx_bluesky/common/parameters/gridscan.py

@@ -26,7 +26,11 @@ from mx_bluesky.common.parameters.constants import (
     HardwareConstants,
 )
 
-DETECTOR_SIZE_PER_BEAMLINE = {"i02-1": EIGER2_X_9M_SIZE, "dev": EIGER2_X_16M_SIZE}
+DETECTOR_SIZE_PER_BEAMLINE = {
+    "i02-1": EIGER2_X_9M_SIZE,
+    "dev": EIGER2_X_16M_SIZE,
+    "i03": EIGER2_X_16M_SIZE,
+}
 
 
 class GridCommon(
@@ -50,6 +54,37 @@ class GridCommon(
 
     tip_offset_um: float = Field(default=HardwareConstants.TIP_OFFSET_UM)
 
+    @property
+    def detector_params(self):
+        self.det_dist_to_beam_converter_path = (
+            self.det_dist_to_beam_converter_path
+            or DetectorParamConstants.BEAM_XY_LUT_PATH
+        )
+        optional_args = {}
+        if self.run_number:
+            optional_args["run_number"] = self.run_number
+        assert self.detector_distance_mm is not None, (
+            "Detector distance must be filled before generating DetectorParams"
+        )
+        return DetectorParams(
+            detector_size_constants=DETECTOR_SIZE_PER_BEAMLINE[
+                get_beamline_name("dev")
+            ],
+            expected_energy_ev=self.demand_energy_ev,
+            exposure_time_s=self.exposure_time_s,
+            directory=self.storage_directory,
+            prefix=self.file_name,
+            detector_distance=self.detector_distance_mm,
+            omega_start=self.omega_start_deg or 0,
+            omega_increment=0,
+            num_images_per_trigger=1,
+            num_triggers=self.num_images,
+            use_roi_mode=self.use_roi_mode,
+            det_dist_to_beam_converter_path=self.det_dist_to_beam_converter_path,
+            trigger_mode=self.trigger_mode,
+            **optional_args,
+        )
+
 
 class SpecifiedGrid(XyzStarts, WithScan):
     """A specified grid is one which has defined values for the start position,
@@ -150,34 +185,3 @@ class SpecifiedThreeDGridScan(
     @property
     def num_images(self) -> int:
         return len(self.scan_points["sam_x"])
-
-    @property
-    def detector_params(self):
-        self.det_dist_to_beam_converter_path = (
-            self.det_dist_to_beam_converter_path
-            or DetectorParamConstants.BEAM_XY_LUT_PATH
-        )
-        optional_args = {}
-        if self.run_number:
-            optional_args["run_number"] = self.run_number
-        assert self.detector_distance_mm is not None, (
-            "Detector distance must be filled before generating DetectorParams"
-        )
-        return DetectorParams(
-            detector_size_constants=DETECTOR_SIZE_PER_BEAMLINE[
-                get_beamline_name("dev")
-            ],
-            expected_energy_ev=self.demand_energy_ev,
-            exposure_time=self.exposure_time_s,
-            directory=self.storage_directory,
-            prefix=self.file_name,
-            detector_distance=self.detector_distance_mm,
-            omega_start=self.omega_start_deg or 0,
-            omega_increment=0,
-            num_images_per_trigger=1,
-            num_triggers=self.num_images,
-            use_roi_mode=self.use_roi_mode,
-            det_dist_to_beam_converter_path=self.det_dist_to_beam_converter_path,
-            trigger_mode=self.trigger_mode,
-            **optional_args,
-        )

mx_bluesky/common/plans/common_flyscan_xray_centre_plan.py

@@ -0,0 +1,316 @@
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Callable, Sequence
+from functools import partial
+
+import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
+import numpy as np
+import pydantic
+from bluesky.protocols import Readable
+from bluesky.utils import MsgGenerator
+from dodal.devices.eiger import EigerDetector
+from dodal.devices.fast_grid_scan import (
+    FastGridScanCommon,
+)
+from dodal.devices.smargon import Smargon
+from dodal.devices.synchrotron import Synchrotron
+from dodal.devices.zocalo import ZocaloResults
+from dodal.devices.zocalo.zocalo_results import (
+    XrcResult,
+    get_full_processing_results,
+)
+
+from mx_bluesky.common.parameters.constants import (
+    DocDescriptorNames,
+    GridscanParamConstants,
+    PlanGroupCheckpointConstants,
+    PlanNameConstants,
+)
+from mx_bluesky.common.parameters.gridscan import SpecifiedThreeDGridScan
+from mx_bluesky.common.plans.inner_plans.do_fgs import (
+    ZOCALO_STAGE_GROUP,
+    kickoff_and_complete_gridscan,
+)
+from mx_bluesky.common.plans.read_hardware import (
+    read_hardware_plan,
+)
+from mx_bluesky.common.utils.exceptions import (
+    CrystalNotFoundException,
+    SampleException,
+)
+from mx_bluesky.common.utils.log import LOGGER
+from mx_bluesky.common.utils.tracing import TRACER
+from mx_bluesky.common.xrc_result import XRayCentreResult
+
+
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
+class FlyScanEssentialDevices:
+    eiger: EigerDetector
+    synchrotron: Synchrotron
+    zocalo: ZocaloResults
+    smargon: Smargon
+
+
+@dataclasses.dataclass
+class BeamlineSpecificFGSFeatures:
+    setup_trigger_plan: Callable[..., MsgGenerator]
+    tidy_plan: Callable[..., MsgGenerator]
+    set_flyscan_params_plan: Callable[..., MsgGenerator]
+    fgs_motors: FastGridScanCommon
+    read_pre_flyscan_plan: Callable[
+        ..., MsgGenerator
+    ]  # Eventually replace with https://github.com/DiamondLightSource/mx-bluesky/issues/819
+    read_during_collection_plan: Callable[..., MsgGenerator]
+    get_xrc_results_from_zocalo: bool
+
+
+def construct_beamline_specific_FGS_features(
+    setup_trigger_plan: Callable[..., MsgGenerator],
+    tidy_plan: Callable[..., MsgGenerator],
+    set_flyscan_params_plan: Callable[..., MsgGenerator],
+    fgs_motors: FastGridScanCommon,
+    signals_to_read_pre_flyscan: list[Readable],
+    signals_to_read_during_collection: list[Readable],
+    get_xrc_results_from_zocalo: bool = False,
+) -> BeamlineSpecificFGSFeatures:
+    """Construct the class needed to do beamline-specific parts of the XRC FGS
+
+    Args:
+        setup_trigger_plan (Callable): Configure triggering, for example with the Zebra or PandA device.
+        Ran directly before kicking off the gridscan.
+
+        tidy_plan (Callable): Tidy up states of devices. Ran at the end of the flyscan, regardless of
+        whether or not it finished successfully.
+
+        set_flyscan_params_plan (Callable): Set PV's for the relevant Fast Grid Scan dodal device
+
+        fgs_motors (Callable): Composite device representing the fast grid scan's motion program parameters.
+
+        signals_to_read_pre_flyscan (Callable): Signals which will be read and saved as a bluesky event document
+        after all configuration, but before the gridscan.
+
+        signals_to_read_during_collection (Callable): Signals which will be read and saved as a bluesky event
+        document whilst the gridscan motion is in progress
+
+        get_xrc_results_from_zocalo (bool): If true, fetch grid scan results from zocalo after completion, as well as
+        update the ispyb comment field with information about the results. See _fetch_xrc_results_from_zocalo
+    """
+    read_pre_flyscan_plan = partial(
+        read_hardware_plan,
+        signals_to_read_pre_flyscan,
+        DocDescriptorNames.HARDWARE_READ_PRE,
+    )
+
+    read_during_collection_plan = partial(
+        read_hardware_plan,
+        signals_to_read_during_collection,
+        DocDescriptorNames.HARDWARE_READ_DURING,
+    )
+
+    return BeamlineSpecificFGSFeatures(
+        setup_trigger_plan,
+        tidy_plan,
+        set_flyscan_params_plan,
+        fgs_motors,
+        read_pre_flyscan_plan,
+        read_during_collection_plan,
+        get_xrc_results_from_zocalo,
+    )
+
+
+def common_flyscan_xray_centre(
+    composite: FlyScanEssentialDevices,
+    parameters: SpecifiedThreeDGridScan,
+    beamline_specific: BeamlineSpecificFGSFeatures,
+) -> MsgGenerator:
+    """Main entry point of the MX-Bluesky x-ray centering flyscan
+
+    Args:
+        composite (FlyScanEssentialDevices): Devices required to perform this plan.
+
+        parameters (SpecifiedThreeDGridScan): Parameters required to perform this plan.
+
+        beamline_specific (BeamlineSpecificFGSFeatures): Configure the beamline-specific version
+        of this plan: For example triggering setup and tidy up plans, as well as what to do with the
+        centering results.
+
+    With a minimum set of devices and parameters, prepares for; performs; and tidies up a flyscan
+    x-ray-center plan. This includes: Configuring desired triggering; writing nexus files; triggering zocalo;
+    reading hardware before and during the scan; and tidying up devices after
+    the plan is complete. Optionally fetch results from zocalo after completing the grid scan.
+
+    This plan will also push data to ispyb when used with the ispyb_activation_decorator.
+
+    There are a few other useful decorators to use with this plan, see: verify_undulator_gap_before_run_decorator, transmission_and_xbpm_feedback_for_collection_decorator
+    """
+
+    def _decorated_flyscan():
+        @bpp.set_run_key_decorator(PlanNameConstants.GRIDSCAN_OUTER)
+        @bpp.run_decorator(  # attach experiment metadata to the start document
+            md={
+                "subplan_name": PlanNameConstants.GRIDSCAN_OUTER,
+                "mx_bluesky_parameters": parameters.model_dump_json(),
+                "activate_callbacks": [
+                    "GridscanNexusFileCallback",
+                ],
+            }
+        )
+        @bpp.finalize_decorator(lambda: beamline_specific.tidy_plan(composite))
+        def run_gridscan_and_tidy(
+            fgs_composite: FlyScanEssentialDevices,
+            params: SpecifiedThreeDGridScan,
+            beamline_specific: BeamlineSpecificFGSFeatures,
+        ) -> MsgGenerator:
+            yield from beamline_specific.setup_trigger_plan(fgs_composite, parameters)
+
+            LOGGER.info("Starting grid scan")
+            yield from bps.stage(
+                fgs_composite.zocalo, group=ZOCALO_STAGE_GROUP
+            )  # connect to zocalo and make sure the queue is clear
+            yield from run_gridscan(fgs_composite, params, beamline_specific)
+
+            LOGGER.info("Grid scan finished")
+
+            if beamline_specific.get_xrc_results_from_zocalo:
+                yield from _fetch_xrc_results_from_zocalo(composite.zocalo, parameters)
+
+        yield from run_gridscan_and_tidy(composite, parameters, beamline_specific)
+
+    composite.eiger.set_detector_parameters(parameters.detector_params)
+    yield from _decorated_flyscan()
+
+
+def _fetch_xrc_results_from_zocalo(
+    zocalo_results: ZocaloResults,
+    parameters: SpecifiedThreeDGridScan,
+) -> MsgGenerator:
+    """
+    Get XRC results from the ZocaloResults device which was staged during a grid scan,
+    and store them in XRayCentreEventHandler.xray_centre_results by firing an event.
+
+    The RunEngine must be subscribed to XRayCentreEventHandler for this plan to work.
+    """
+
+    LOGGER.info("Getting X-ray center Zocalo results...")
+
+    yield from bps.trigger(zocalo_results)
+    LOGGER.info("Zocalo triggered and read, interpreting results.")
+    xrc_results = yield from get_full_processing_results(zocalo_results)
+    LOGGER.info(f"Got xray centres, top 5: {xrc_results[:5]}")
+    filtered_results = [
+        result
+        for result in xrc_results
+        if result["total_count"]
+        >= GridscanParamConstants.ZOCALO_MIN_TOTAL_COUNT_THRESHOLD
+    ]
+    discarded_count = len(xrc_results) - len(filtered_results)
+    if discarded_count > 0:
+        LOGGER.info(f"Removed {discarded_count} results because below threshold")
+    if filtered_results:
+        flyscan_results = [
+            _xrc_result_in_boxes_to_result_in_mm(xr, parameters)
+            for xr in filtered_results
+        ]
+    else:
+        LOGGER.warning("No X-ray centre received")
+        raise CrystalNotFoundException()
+    yield from _fire_xray_centre_result_event(flyscan_results)
+
+
+@bpp.set_run_key_decorator(PlanNameConstants.GRIDSCAN_MAIN)
+@bpp.run_decorator(md={"subplan_name": PlanNameConstants.GRIDSCAN_MAIN})
+def run_gridscan(
+    fgs_composite: FlyScanEssentialDevices,
+    parameters: SpecifiedThreeDGridScan,
+    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.smargon.omega, 0)
+
+    with TRACER.start_span("ispyb_hardware_readings"):
+        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)
+
+    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 kickoff_and_complete_gridscan(
+        beamline_specific.fgs_motors,
+        fgs_composite.eiger,
+        fgs_composite.synchrotron,
+        [parameters.scan_points_first_grid, parameters.scan_points_second_grid],
+        parameters.scan_indices,
+        plan_during_collection=beamline_specific.read_during_collection_plan,
+    )
+    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:
+    fgs_params = parameters.FGS_params
+    xray_centre = fgs_params.grid_position_to_motor_position(
+        np.array(xrc_result["centre_of_mass"])
+    )
+    # A correction is applied to the bounding box to map discrete grid coordinates to
+    # the corners of the box in motor-space; we do not apply this correction
+    # to the xray-centre as it is already in continuous space and the conversion has
+    # been performed already
+    # In other words, xrc_result["bounding_box"] contains the position of the box centre,
+    # so we subtract half a box to get the corner of the box
+    return XRayCentreResult(
+        centre_of_mass_mm=xray_centre,
+        bounding_box_mm=(
+            fgs_params.grid_position_to_motor_position(
+                np.array(xrc_result["bounding_box"][0]) - 0.5
+            ),
+            fgs_params.grid_position_to_motor_position(
+                np.array(xrc_result["bounding_box"][1]) - 0.5
+            ),
+        ),
+        max_count=xrc_result["max_count"],
+        total_count=xrc_result["total_count"],
+    )
+
+
+def _fire_xray_centre_result_event(results: Sequence[XRayCentreResult]):
+    def empty_plan():
+        return iter([])
+
+    yield from bpp.set_run_key_wrapper(
+        bpp.run_wrapper(
+            empty_plan(),
+            md={
+                PlanNameConstants.FLYSCAN_RESULTS: [
+                    dataclasses.asdict(r) for r in results
+                ]
+            },
+        ),
+        PlanNameConstants.FLYSCAN_RESULTS,
+    )

mx_bluesky/common/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.dcm import DCM
+from dodal.devices.common_dcm import BaseDCM
 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: DCM,
+    dcm: BaseDCM,
     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: DCM,
+    dcm: BaseDCM,
     detector: EigerDetector,
 ):
     signals_to_read_during_collection = [

mx_bluesky/common/plans/write_sample_status.py

@@ -0,0 +1,46 @@
+from enum import StrEnum
+
+import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
+
+from mx_bluesky.common.external_interaction.callbacks.sample_handling.sample_handling_callback import (
+    SampleHandlingCallback,
+)
+from mx_bluesky.common.utils.exceptions import SampleException
+
+
+class SampleStatusExceptionType(StrEnum):
+    BEAMLINE = "Beamline"
+    SAMPLE = "Sample"
+
+
+@bpp.subs_decorator(SampleHandlingCallback())
+def deposit_sample_error(exception_type: SampleStatusExceptionType, sample_id: int):
+    @bpp.run_decorator(
+        md={
+            "metadata": {"sample_id": sample_id},
+            "activate_callbacks": ["SampleHandlingCallback"],
+        }
+    )
+    def _inner():
+        yield from bps.null()
+        if exception_type == SampleStatusExceptionType.BEAMLINE:
+            raise AssertionError()
+        elif exception_type == SampleStatusExceptionType.SAMPLE:
+            raise SampleException
+
+    yield from _inner()
+
+
+@bpp.subs_decorator(SampleHandlingCallback(record_loaded_on_success=True))
+def deposit_loaded_sample(sample_id: int):
+    @bpp.run_decorator(
+        md={
+            "metadata": {"sample_id": sample_id},
+            "activate_callbacks": ["SampleHandlingCallback"],
+        }
+    )
+    def _inner():
+        yield from bps.null()
+
+    yield from _inner()

mx_bluesky/common/preprocessors/preprocessors.py

@@ -0,0 +1,105 @@
+from bluesky import preprocessors as bpp
+from bluesky.preprocessors import plan_mutator
+from bluesky.utils import Msg, MsgGenerator, make_decorator
+
+from mx_bluesky.common.device_setup_plans.xbpm_feedback import (
+    check_and_pause_feedback,
+    unpause_xbpm_feedback_and_set_transmission_to_1,
+)
+from mx_bluesky.common.parameters.constants import PlanNameConstants
+from mx_bluesky.common.protocols.protocols import (
+    XBPMPauseDevices,
+)
+
+
+def transmission_and_xbpm_feedback_for_collection_wrapper(
+    plan: MsgGenerator,
+    devices: XBPMPauseDevices,
+    desired_transmission_fraction: float,
+    run_key_to_wrap: PlanNameConstants | None = None,
+):
+    """
+    Sets the transmission for the data collection, ensuring the xbpm feedback is valid, then resets it immediately after
+    the run has finished.
+
+    This wrapper should be attached to the entry point of any beamline-specific plan that may disrupt the XBPM feedback,
+    such as a data collection or an x-ray center grid scan.
+    This wrapper will do nothing if no runs are seen.
+
+    XBPM feedback isn't reliable during collections due to:
+     * Objects (e.g. attenuator) crossing the beam can cause large (incorrect) feedback movements
+     * Lower transmissions/higher energies are less reliable for the xbpm
+
+    So we need to keep the transmission at 100% and the feedback on when not collecting
+    and then turn it off and set the correct transmission for collection. The feedback
+    mostly accounts for slow thermal drift so it is safe to assume that the beam is
+    stable during a collection.
+
+    Args:
+        plan: The plan performing the data collection.
+        devices (XBPMPauseDevices): Composite device including The XBPM device that is responsible for keeping
+                                                        the beam in position, and attenuator
+        desired_transmission_fraction (float): The desired transmission for the collection
+        run_key_to_wrap: (str | None): Pausing XBPM and setting transmission is inserted after the 'open_run' message is seen with
+        the matching run key, and unpausing and resetting transmission is inserted after the corresponding 'close_run' message is
+        seen. If not specified, instead wrap the first run encountered.
+    """
+
+    _wrapped_run_name: None | str = None
+
+    def head(msg: Msg):
+        yield from check_and_pause_feedback(
+            devices.xbpm_feedback,
+            devices.attenuator,
+            desired_transmission_fraction,
+        )
+
+        # Allow 'open_run' message to pass through
+        yield msg
+
+    def tail():
+        yield from unpause_xbpm_feedback_and_set_transmission_to_1(
+            devices.xbpm_feedback, devices.attenuator
+        )
+
+    def insert_plans(msg: Msg):
+        # Wrap the specified run, or, if none specified, wrap the first run encountered
+        nonlocal _wrapped_run_name
+
+        match msg.command:
+            case "open_run":
+                # If we specified a run key, did we encounter it
+                # If we didn't specify, then insert the plans and track the name of the run
+                if (
+                    not (run_key_to_wrap or _wrapped_run_name)
+                    or run_key_to_wrap is msg.run
+                ):
+                    _wrapped_run_name = msg.run if msg.run else "unnamed_run"
+                    return head(msg), None
+            case "close_run":
+                # Check if the run tracked from above was closed
+                # An exception is raised in the RunEngine if two unnamed runs are opened
+                # at the same time, so we are safe from unpausing on the wrong run
+                if (_wrapped_run_name == "unnamed_run" and not msg.run) or (
+                    msg.run and _wrapped_run_name and _wrapped_run_name is msg.run
+                ):
+                    return None, tail()
+
+        return None, None
+
+    # Contingency wrapper can cause unpausing to occur on exception and again on close_run.
+    # Not needed after https://github.com/bluesky/bluesky/issues/1891
+    return (
+        yield from bpp.contingency_wrapper(
+            plan_mutator(plan, insert_plans),
+            except_plan=lambda _: unpause_xbpm_feedback_and_set_transmission_to_1(
+                devices.xbpm_feedback,
+                devices.attenuator,
+            ),
+        )
+    )
+
+
+transmission_and_xbpm_feedback_for_collection_decorator = make_decorator(
+    transmission_and_xbpm_feedback_for_collection_wrapper
+)