mx-bluesky 0.0.2__py3-none-any.whl → 1.1.0__py3-none-any.whl

mx_bluesky/hyperion/device_setup_plans/setup_oav.py

@@ -0,0 +1,87 @@
+from functools import partial
+
+import bluesky.plan_stubs as bps
+from dodal.devices.oav.oav_detector import OAV
+from dodal.devices.oav.oav_errors import OAVError_ZoomLevelNotFound
+from dodal.devices.oav.oav_parameters import OAVParameters
+from dodal.devices.oav.pin_image_recognition import PinTipDetection
+from dodal.devices.oav.utils import ColorMode
+
+from mx_bluesky.hyperion.parameters.constants import CONST
+
+# Helper function to make sure we set the waiting groups correctly
+set_using_group = partial(bps.abs_set, group=CONST.WAIT.READY_FOR_OAV)
+
+
+def setup_pin_tip_detection_params(
+    pin_tip_detect_device: PinTipDetection, parameters: OAVParameters
+):
+    # select which blur to apply to image
+    yield from set_using_group(
+        pin_tip_detect_device.preprocess_operation, parameters.preprocess
+    )
+
+    # sets length scale for blurring
+    yield from set_using_group(
+        pin_tip_detect_device.preprocess_ksize, parameters.preprocess_K_size
+    )
+
+    # Canny edge detect - lower
+    yield from set_using_group(
+        pin_tip_detect_device.canny_lower_threshold,
+        parameters.canny_edge_lower_threshold,
+    )
+
+    # Canny edge detect - upper
+    yield from set_using_group(
+        pin_tip_detect_device.canny_upper_threshold,
+        parameters.canny_edge_upper_threshold,
+    )
+
+    # "Close" morphological operation
+    yield from set_using_group(
+        pin_tip_detect_device.close_ksize, parameters.close_ksize
+    )
+
+    # Sample detection direction
+    yield from set_using_group(
+        pin_tip_detect_device.scan_direction, parameters.direction
+    )
+
+    # Minimum height
+    yield from set_using_group(
+        pin_tip_detect_device.min_tip_height,
+        parameters.minimum_height,
+    )
+
+
+def setup_general_oav_params(oav: OAV, parameters: OAVParameters):
+    yield from set_using_group(oav.cam.color_mode, ColorMode.RGB1)
+    yield from set_using_group(oav.cam.acquire_period, parameters.acquire_period)
+    yield from set_using_group(oav.cam.acquire_time, parameters.exposure)
+    yield from set_using_group(oav.cam.gain, parameters.gain)
+
+    zoom_level_str = f"{float(parameters.zoom)}x"
+    if zoom_level_str not in oav.zoom_controller.allowed_zoom_levels:
+        raise OAVError_ZoomLevelNotFound(
+            f"Found {zoom_level_str} as a zoom level but expected one of {oav.zoom_controller.allowed_zoom_levels}"
+        )
+
+    yield from bps.abs_set(
+        oav.zoom_controller,
+        zoom_level_str,
+        wait=True,
+    )
+
+
+def pre_centring_setup_oav(
+    oav: OAV,
+    parameters: OAVParameters,
+    pin_tip_detection_device: PinTipDetection,
+):
+    """
+    Setup OAV PVs with required values.
+    """
+    yield from setup_general_oav_params(oav, parameters)
+    yield from setup_pin_tip_detection_params(pin_tip_detection_device, parameters)
+    yield from bps.wait(CONST.WAIT.READY_FOR_OAV)

mx_bluesky/hyperion/device_setup_plans/setup_panda.py

@@ -0,0 +1,210 @@
+from datetime import datetime
+from enum import Enum
+from importlib import resources
+from pathlib import Path
+
+import bluesky.plan_stubs as bps
+from blueapi.core import MsgGenerator
+from dodal.common.beamlines.beamline_utils import get_path_provider
+from dodal.devices.fast_grid_scan import PandAGridScanParams
+from ophyd_async.core import load_device
+from ophyd_async.fastcs.panda import (
+    HDFPanda,
+    SeqTable,
+    SeqTrigger,
+)
+
+import mx_bluesky.hyperion.resources.panda as panda_resource
+from mx_bluesky.hyperion.log import LOGGER
+
+MM_TO_ENCODER_COUNTS = 200000
+GENERAL_TIMEOUT = 60
+TICKS_PER_MS = 1000  # Panda sequencer prescaler will be set to us
+
+
+class Enabled(Enum):
+    ENABLED = "ONE"
+    DISABLED = "ZERO"
+
+
+class PcapArm(Enum):
+    ARMED = "Arm"
+    DISARMED = "Disarm"
+
+
+def _get_seq_table(
+    parameters: PandAGridScanParams, exposure_distance_mm, time_between_steps_ms
+) -> SeqTable:
+    """
+    Generate the sequencer table for the panda.
+
+    - Sending a 'trigger' means trigger PCAP internally and send signal to Eiger via physical panda output
+
+    SEQUENCER TABLE:
+
+        1. Wait for physical trigger from motion script to mark start of scan / change of direction
+        2. Wait for POSA (X2) to be greater than X_START and send x_steps triggers every time_between_steps_ms
+        3. Wait for physical trigger from motion script to mark change of direction
+        4. Wait for POSA (X2) to be less than X_START + X_STEP_SIZE * x_steps + exposure distance, then
+            send x_steps triggers every time_between_steps_ms
+        5. Go back to step one.
+
+        For a more detailed explanation and a diagram, see https://github.com/DiamondLightSource/hyperion/wiki/PandA-constant%E2%80%90motion-scanning
+
+        For documentation on Panda itself, see https://pandablocks.github.io/PandABlocks-FPGA/master/index.html
+
+    Args:
+        exposure_distance_mm: The distance travelled by the sample each time the detector is exposed: exposure time * sample velocity
+        time_between_steps_ms: The time taken to traverse between each grid step.
+        parameters: Parameters for the panda gridscan
+
+    Returns:
+        An instance of SeqTable describing the panda sequencer table
+    """
+
+    start_of_grid_x_counts = int(parameters.x_start * MM_TO_ENCODER_COUNTS)
+
+    # x_start is the first trigger point, so we need to travel to x_steps-1 for the final trigger point
+    end_of_grid_x_counts = int(
+        start_of_grid_x_counts
+        + (parameters.x_step_size * (parameters.x_steps - 1) * MM_TO_ENCODER_COUNTS)
+    )
+
+    exposure_distance_x_counts = int(exposure_distance_mm * MM_TO_ENCODER_COUNTS)
+
+    num_pulses = parameters.x_steps
+
+    delay_between_pulses = time_between_steps_ms * TICKS_PER_MS
+
+    PULSE_WIDTH_US = 1
+
+    assert delay_between_pulses > PULSE_WIDTH_US
+
+    # BITA_1 trigger wired from TTLIN1, this is the trigger input
+
+    # +ve direction scan
+
+    table = (
+        SeqTable.row(trigger=SeqTrigger.BITA_1, time2=1)
+        + SeqTable.row(
+            repeats=num_pulses,
+            trigger=SeqTrigger.POSA_GT,
+            position=start_of_grid_x_counts,
+            time1=PULSE_WIDTH_US,
+            outa1=True,
+            time2=delay_between_pulses - PULSE_WIDTH_US,
+            outa2=False,
+        )
+        +
+        # -ve direction scan
+        SeqTable.row(trigger=SeqTrigger.BITA_1, time2=1)
+        + SeqTable.row(
+            repeats=num_pulses,
+            trigger=SeqTrigger.POSA_LT,
+            position=end_of_grid_x_counts + exposure_distance_x_counts,
+            time1=PULSE_WIDTH_US,
+            outa1=True,
+            time2=delay_between_pulses - PULSE_WIDTH_US,
+            outa2=False,
+        )
+    )
+
+    return table
+
+
+def setup_panda_for_flyscan(
+    panda: HDFPanda,
+    parameters: PandAGridScanParams,
+    initial_x: float,
+    exposure_time_s: float,
+    time_between_x_steps_ms: float,
+    sample_velocity_mm_per_s: float,
+) -> MsgGenerator:
+    """Configures the PandA device for a flyscan.
+    Sets PVs from a yaml file, calibrates the encoder, and
+    adjusts the sequencer table based off the grid parameters. Yaml file can be
+    created using ophyd_async.core.save_device()
+
+    Args:
+        panda (HDFPanda): The PandA Ophyd device
+        parameters (PandAGridScanParams): Grid parameters
+        initial_x (float): Motor positions at time of PandA setup
+        exposure_time_s (float): Detector exposure time per trigger
+        time_between_x_steps_ms (float): Time, in ms, between each trigger. Equal to deadtime + exposure time
+        sample_velocity_mm_per_s (float): Velocity of the sample in mm/s = x_step_size_mm * 1000 /
+            time_between_x_steps_ms
+    Returns:
+        MsgGenerator
+
+    Yields:
+        Iterator[MsgGenerator]
+    """
+    assert parameters.x_steps > 0
+    assert time_between_x_steps_ms * 1000 >= exposure_time_s
+    assert sample_velocity_mm_per_s * exposure_time_s < parameters.x_step_size
+
+    yield from bps.stage(panda, group="panda-config")
+
+    with resources.as_file(
+        resources.files(panda_resource) / "panda-gridscan.yaml"
+    ) as config_yaml_path:
+        yield from load_device(panda, str(config_yaml_path))
+
+    # Home the PandA X encoder using current motor position
+    yield from bps.abs_set(
+        panda.inenc[1].setp,  # type: ignore
+        initial_x * MM_TO_ENCODER_COUNTS,
+        wait=True,
+    )
+
+    yield from bps.abs_set(panda.pulse[1].width, exposure_time_s, group="panda-config")
+
+    exposure_distance_mm = sample_velocity_mm_per_s * exposure_time_s
+
+    table = _get_seq_table(parameters, exposure_distance_mm, time_between_x_steps_ms)
+
+    yield from bps.abs_set(panda.seq[1].table, table, group="panda-config")
+
+    yield from bps.abs_set(
+        panda.pcap.enable,  # type: ignore
+        Enabled.ENABLED.value,
+        group="panda-config",
+    )
+
+    # Values need to be set before blocks are enabled, so wait here
+    yield from bps.wait(group="panda-config", timeout=GENERAL_TIMEOUT)
+
+    LOGGER.info(f"PandA sequencer table has been set to: {str(table)}")
+    table_readback = yield from bps.rd(panda.seq[1].table)
+    LOGGER.debug(f"PandA sequencer table readback is: {str(table_readback)}")
+
+    yield from arm_panda_for_gridscan(panda)
+
+
+def arm_panda_for_gridscan(panda: HDFPanda, group="arm_panda_gridscan"):
+    yield from bps.abs_set(panda.seq[1].enable, Enabled.ENABLED.value, group=group)  # type: ignore
+    yield from bps.abs_set(panda.pulse[1].enable, Enabled.ENABLED.value, group=group)  # type: ignore
+    yield from bps.abs_set(panda.counter[1].enable, Enabled.ENABLED.value, group=group)  # type: ignore
+    yield from bps.abs_set(panda.pcap.arm, PcapArm.ARMED.value, group=group)  # type: ignore
+    yield from bps.wait(group=group, timeout=GENERAL_TIMEOUT)
+    LOGGER.info("PandA has been armed")
+
+
+def disarm_panda_for_gridscan(panda, group="disarm_panda_gridscan") -> MsgGenerator:
+    yield from bps.abs_set(panda.pcap.arm, PcapArm.DISARMED.value, group=group)  # type: ignore
+    yield from bps.abs_set(panda.counter[1].enable, Enabled.DISABLED.value, group=group)  # type: ignore
+    yield from bps.abs_set(panda.seq[1].enable, Enabled.DISABLED.value, group=group)
+    yield from bps.abs_set(panda.pulse[1].enable, Enabled.DISABLED.value, group=group)
+    yield from bps.abs_set(panda.pcap.enable, Enabled.DISABLED.value, group=group)  # type: ignore
+    yield from bps.wait(group=group, timeout=GENERAL_TIMEOUT)
+
+
+def set_panda_directory(panda_directory: Path) -> MsgGenerator:
+    """Updates the root folder which is used by the PandA's PCAP."""
+
+    suffix = datetime.now().strftime("_%Y%m%d%H%M%S")
+
+    async def set_panda_dir():
+        await get_path_provider().update(directory=panda_directory, suffix=suffix)
+
+    yield from bps.wait_for([set_panda_dir])

mx_bluesky/hyperion/device_setup_plans/setup_zebra.py

@@ -0,0 +1,214 @@
+from collections.abc import Callable
+from functools import wraps
+
+import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
+from blueapi.core import MsgGenerator
+from dodal.devices.zebra import (
+    AUTO_SHUTTER_GATE,
+    AUTO_SHUTTER_INPUT,
+    DISCONNECT,
+    IN1_TTL,
+    IN3_TTL,
+    IN4_TTL,
+    PC_GATE,
+    PC_PULSE,
+    TTL_DETECTOR,
+    TTL_PANDA,
+    TTL_XSPRESS3,
+    ArmDemand,
+    EncEnum,
+    I03Axes,
+    RotationDirection,
+    Zebra,
+)
+from dodal.devices.zebra_controlled_shutter import ZebraShutter, ZebraShutterControl
+
+from mx_bluesky.hyperion.log import LOGGER
+
+ZEBRA_STATUS_TIMEOUT = 30
+
+
+def bluesky_retry(func: Callable):
+    """Decorator that will retry the decorated plan if it fails.
+
+    Use this with care as it knows nothing about the state of the world when things fail.
+    If it is possible that your plan fails when the beamline is in a transient state that
+    the plan could not act on do not use this decorator without doing some more intelligent
+    clean up.
+
+    You should avoid using this decorator often in general production as it hides errors,
+    instead it should be used only for debugging these underlying errors.
+    """
+
+    @wraps(func)
+    def newfunc(*args, **kwargs):
+        def log_and_retry(exception):
+            LOGGER.error(f"Function {func.__name__} failed with {exception}, retrying")
+            yield from func(*args, **kwargs)
+
+        yield from bpp.contingency_wrapper(
+            func(*args, **kwargs), except_plan=log_and_retry, auto_raise=False
+        )
+
+    return newfunc
+
+
+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 set_shutter_auto_input(zebra: Zebra, input: int, group="set_shutter_trigger"):
+    """Set the input that the shutter uses when set to auto.
+
+    For more details see the ZebraShutter device."""
+    auto_shutter_control = zebra.logic_gates.and_gates[AUTO_SHUTTER_GATE]
+    yield from bps.abs_set(
+        auto_shutter_control.sources[AUTO_SHUTTER_INPUT], input, group
+    )
+
+
+@bluesky_retry
+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)
+    # Trigger the shutter with the gate
+    yield from bps.abs_set(
+        zebra_shutter.control_mode, ZebraShutterControl.AUTO, group=group
+    )
+    yield from set_shutter_auto_input(zebra, PC_GATE, group=group)
+    # Trigger the detector with a pulse
+    yield from bps.abs_set(zebra.output.out_pvs[TTL_DETECTOR], PC_PULSE, group=group)
+    # Don't use the fluorescence detector
+    yield from bps.abs_set(zebra.output.out_pvs[TTL_XSPRESS3], DISCONNECT, group=group)
+    yield from bps.abs_set(zebra.output.pulse_1.input, 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)
+
+
+@bluesky_retry
+def setup_zebra_for_gridscan(
+    zebra: Zebra,
+    zebra_shutter: ZebraShutter,
+    group="setup_zebra_for_gridscan",
+    wait=True,
+):
+    yield from bps.abs_set(zebra.output.out_pvs[TTL_DETECTOR], IN3_TTL, group=group)
+    yield from bps.abs_set(
+        zebra_shutter.control_mode, ZebraShutterControl.AUTO, group=group
+    )
+    yield from set_shutter_auto_input(zebra, IN4_TTL, group=group)
+    yield from bps.abs_set(zebra.output.out_pvs[TTL_XSPRESS3], DISCONNECT, group=group)
+    yield from bps.abs_set(zebra.output.pulse_1.input, DISCONNECT, group=group)
+
+    if wait:
+        yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT)
+
+
+@bluesky_retry
+def tidy_up_zebra_after_gridscan(
+    zebra: Zebra,
+    zebra_shutter: ZebraShutter,
+    group="tidy_up_zebra_after_gridscan",
+    wait=True,
+) -> MsgGenerator:
+    yield from bps.abs_set(zebra.output.out_pvs[TTL_DETECTOR], PC_PULSE, group=group)
+    yield from bps.abs_set(
+        zebra_shutter.control_mode, ZebraShutterControl.MANUAL, group=group
+    )
+    yield from set_shutter_auto_input(zebra, PC_GATE, group=group)
+
+    if wait:
+        yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT)
+
+
+@bluesky_retry
+def setup_zebra_for_panda_flyscan(
+    zebra: Zebra,
+    zebra_shutter: ZebraShutter,
+    group="setup_zebra_for_panda_flyscan",
+    wait=True,
+):
+    # Forwards eiger trigger signal from panda
+    yield from bps.abs_set(zebra.output.out_pvs[TTL_DETECTOR], IN1_TTL, group=group)
+
+    # Forwards signal from PPMAC to fast shutter. High while panda PLC is running
+    yield from bps.abs_set(
+        zebra_shutter.control_mode, ZebraShutterControl.AUTO, group=group
+    )
+    yield from set_shutter_auto_input(zebra, IN4_TTL, group=group)
+
+    yield from bps.abs_set(zebra.output.out_pvs[TTL_XSPRESS3], DISCONNECT, group=group)
+
+    yield from bps.abs_set(
+        zebra.output.out_pvs[TTL_PANDA], IN3_TTL, group=group
+    )  # Tells panda that motion is beginning/changing direction
+
+    if wait:
+        yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT)

mx_bluesky/hyperion/device_setup_plans/smargon.py

@@ -0,0 +1,25 @@
+import numpy as np
+from bluesky import plan_stubs as bps
+from dodal.devices.smargon import Smargon
+
+from mx_bluesky.hyperion.exceptions import WarningException
+
+
+def move_smargon_warn_on_out_of_range(
+    smargon: Smargon, position: np.ndarray | list[float] | tuple[float, float, float]
+):
+    """Throws a WarningException if the specified position is out of range for the
+    smargon. Otherwise moves to that position."""
+    limits = yield from smargon.get_xyz_limits()
+    if not limits.position_valid(position):
+        raise WarningException(
+            "Pin tip centring failed - pin too long/short/bent and out of range"
+        )
+    yield from bps.mv(
+        smargon.x,
+        position[0],
+        smargon.y,
+        position[1],
+        smargon.z,
+        position[2],
+    )

mx_bluesky/hyperion/device_setup_plans/utils.py

@@ -0,0 +1,44 @@
+from collections.abc import Generator
+
+from bluesky import plan_stubs as bps
+from bluesky import preprocessors as bpp
+from bluesky.utils import Msg
+from dodal.devices.detector.detector_motion import DetectorMotion, ShutterState
+from dodal.devices.eiger import EigerDetector
+
+from mx_bluesky.hyperion.device_setup_plans.position_detector import (
+    set_detector_z_position,
+    set_shutter,
+)
+
+
+def start_preparing_data_collection_then_do_plan(
+    eiger: EigerDetector,
+    detector_motion: DetectorMotion,
+    detector_distance_mm: float | None,
+    plan_to_run: Generator[Msg, None, None],
+    group="ready_for_data_collection",
+) -> Generator[Msg, None, None]:
+    """Starts preparing for the next data collection and then runs the
+    given plan.
+
+     Preparation consists of:
+     * Arming the Eiger
+     * Moving the detector to the specified position
+     * Opening the detect shutter
+     If the plan fails it will disarm the eiger.
+    """
+
+    def wrapped_plan():
+        yield from bps.abs_set(eiger.do_arm, 1, group=group)
+        if detector_distance_mm:
+            yield from set_detector_z_position(
+                detector_motion, detector_distance_mm, group
+            )
+        yield from set_shutter(detector_motion, ShutterState.OPEN, group)
+        yield from plan_to_run
+
+    yield from bpp.contingency_wrapper(
+        wrapped_plan(),
+        except_plan=lambda e: (yield from bps.stop(eiger)),
+    )

mx_bluesky/hyperion/device_setup_plans/xbpm_feedback.py

@@ -0,0 +1,93 @@
+from bluesky import plan_stubs as bps
+from bluesky.preprocessors import finalize_wrapper
+from bluesky.utils import make_decorator
+from dodal.devices.attenuator import Attenuator
+from dodal.devices.xbpm_feedback import Pause, XBPMFeedback
+
+from mx_bluesky.hyperion.log import LOGGER
+
+
+def _check_and_pause_feedback(
+    xbpm_feedback: XBPMFeedback,
+    attenuator: Attenuator,
+    desired_transmission_fraction: float,
+):
+    """Checks that the xbpm is in position before then turning it off and setting a new
+    transmission.
+
+    Args:
+        xbpm_feedback (XBPMFeedback): The XBPM device that is responsible for keeping
+                                      the beam in position
+        attenuator (Attenuator): The attenuator used to set transmission
+        desired_transmission_fraction (float): The desired transmission to set after
+                                               turning XBPM feedback off.
+
+    """
+    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)
+    yield from bps.mv(attenuator, desired_transmission_fraction)
+
+
+def _unpause_xbpm_feedback_and_set_transmission_to_1(
+    xbpm_feedback: XBPMFeedback, attenuator: Attenuator
+):
+    """Turns the XBPM feedback back on and sets transmission to 1 so that it keeps the
+    beam aligned whilst not collecting.
+
+    Args:
+        xbpm_feedback (XBPMFeedback): The XBPM device that is responsible for keeping
+                                      the beam in position
+        attenuator (Attenuator): The attenuator used to set transmission
+    """
+    yield from bps.mv(xbpm_feedback.pause_feedback, Pause.RUN, attenuator, 1.0)
+
+
+def transmission_and_xbpm_feedback_for_collection_wrapper(
+    plan,
+    xbpm_feedback: XBPMFeedback,
+    attenuator: Attenuator,
+    desired_transmission_fraction: float,
+):
+    """Sets the transmission for the data collection, ensuring the xbpm feedback is valid
+    this wrapper should be run around every data collection or movement that may disrupt
+    the XBPM feedback.
+
+    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
+        xbpm_feedback (XBPMFeedback): The XBPM device that is responsible for keeping
+                                      the beam in position
+        attenuator (Attenuator): The attenuator used to set transmission
+        desired_transmission_fraction (float): The desired transmission for the collection
+    """
+
+    def _inner_plan():
+        yield from _check_and_pause_feedback(
+            xbpm_feedback, attenuator, desired_transmission_fraction
+        )
+        return (yield from plan)
+
+    return (
+        yield from finalize_wrapper(
+            _inner_plan(),
+            _unpause_xbpm_feedback_and_set_transmission_to_1(xbpm_feedback, attenuator),
+        )
+    )
+
+
+transmission_and_xbpm_feedback_for_collection_decorator = make_decorator(
+    transmission_and_xbpm_feedback_for_collection_wrapper
+)