mx-bluesky 1.2.0__py3-none-any.whl → 1.4.1a0__py3-none-any.whl

mx_bluesky/hyperion/device_setup_plans/dcm_pitch_roll_mirror_adjuster.py

@@ -4,7 +4,7 @@ import bluesky.plan_stubs as bps
 from dodal.devices.focusing_mirror import (
     FocusingMirrorWithStripes,
     MirrorStripe,
-    VFMMirrorVoltages,
+    MirrorVoltages,
 )
 from dodal.devices.undulator_dcm import UndulatorDCM
 from dodal.devices.util.adjuster_plans import lookup_table_adjuster
@@ -13,6 +13,9 @@ from dodal.devices.util.lookup_tables import (
 )
 
 from mx_bluesky.hyperion.log import LOGGER
+from mx_bluesky.hyperion.utils.utils import (
+    energy_to_bragg_angle,
+)
 
 MIRROR_VOLTAGE_GROUP = "MIRROR_VOLTAGE_GROUP"
 DCM_GROUP = "DCM_GROUP"
@@ -20,40 +23,41 @@ DCM_GROUP = "DCM_GROUP"
 
 def _apply_and_wait_for_voltages_to_settle(
     stripe: MirrorStripe,
-    mirror: FocusingMirrorWithStripes,
-    mirror_voltages: VFMMirrorVoltages,
+    mirror_voltages: MirrorVoltages,
 ):
     with open(mirror_voltages.voltage_lookup_table_path) as lut_file:
         json_obj = json.load(lut_file)
 
     # sample mode is the only mode supported
     sample_data = json_obj["sample"]
-    mirror_key = mirror.name.lower()
     if stripe == MirrorStripe.BARE:
         stripe_key = "bare"
     elif stripe == MirrorStripe.RHODIUM:
         stripe_key = "rh"
     elif stripe == MirrorStripe.PLATINUM:
         stripe_key = "pt"
-    else:
-        raise ValueError(f"Unsupported stripe '{stripe}'")
-
-    required_voltages = sample_data[stripe_key][mirror_key]
-    for voltage_channel, required_voltage in zip(
-        mirror_voltages.voltage_channels.values(), required_voltages, strict=False
-    ):
-        LOGGER.debug(
-            f"Applying and waiting for voltage {voltage_channel.name} = {required_voltage}"
-        )
-        yield from bps.abs_set(
-            voltage_channel, required_voltage, group=MIRROR_VOLTAGE_GROUP
-        )
+
+    for mirror_key, channels in {
+        "hfm": mirror_voltages.horizontal_voltages,
+        "vfm": mirror_voltages.vertical_voltages,
+    }.items():
+        required_voltages = sample_data[stripe_key][mirror_key]
+
+        for voltage_channel, required_voltage in zip(
+            channels.values(), required_voltages, strict=True
+        ):
+            LOGGER.debug(
+                f"Applying and waiting for voltage {voltage_channel.name} = {required_voltage}"
+            )
+            yield from bps.abs_set(
+                voltage_channel, required_voltage, group=MIRROR_VOLTAGE_GROUP
+            )
 
     yield from bps.wait(group=MIRROR_VOLTAGE_GROUP)
 
 
 def adjust_mirror_stripe(
-    energy_kev, mirror: FocusingMirrorWithStripes, mirror_voltages: VFMMirrorVoltages
+    energy_kev, mirror: FocusingMirrorWithStripes, mirror_voltages: MirrorVoltages
 ):
     """Feedback should be OFF prior to entry, in order to prevent
     feedback from making unnecessary corrections while beam is being adjusted."""
@@ -66,26 +70,27 @@ def adjust_mirror_stripe(
     yield from bps.trigger(mirror.apply_stripe)
 
     LOGGER.info("Adjusting mirror voltages...")
-    yield from _apply_and_wait_for_voltages_to_settle(stripe, mirror, mirror_voltages)
+    yield from _apply_and_wait_for_voltages_to_settle(stripe, mirror_voltages)
 
 
 def adjust_dcm_pitch_roll_vfm_from_lut(
     undulator_dcm: UndulatorDCM,
     vfm: FocusingMirrorWithStripes,
-    vfm_mirror_voltages: VFMMirrorVoltages,
+    mirror_voltages: MirrorVoltages,
     energy_kev,
 ):
     """Beamline energy-change post-adjustments : Adjust DCM and VFM directly from lookup tables.
-    Lookups are performed against the Bragg angle which will have been automatically set by EPICS as a side-effect of the
-    energy change prior to calling this function.
+    Lookups are performed against the Bragg angle which is computed directly from the target energy
+    rather than waiting for the EPICS controls PV to reach it.
     Feedback should be OFF prior to entry, in order to prevent
     feedback from making unnecessary corrections while beam is being adjusted."""
 
-    # DCM Pitch
+    # Adjust DCM Pitch
     dcm = undulator_dcm.dcm
     LOGGER.info(f"Adjusting DCM and VFM for {energy_kev} keV")
-    bragg_deg = yield from bps.rd(dcm.bragg_in_degrees.user_readback)
-    LOGGER.info(f"Read Bragg angle = {bragg_deg} degrees")
+    d_spacing_a: float = yield from bps.rd(undulator_dcm.dcm.crystal_metadata_d_spacing)
+    bragg_deg = energy_to_bragg_angle(energy_kev, d_spacing_a)
+    LOGGER.info(f"Target Bragg angle = {bragg_deg} degrees")
     dcm_pitch_adjuster = lookup_table_adjuster(
         linear_interpolation_lut(undulator_dcm.pitch_energy_table_path),
         dcm.pitch_in_mrad,
@@ -119,7 +124,7 @@ def adjust_dcm_pitch_roll_vfm_from_lut(
     # not sure how we check this
 
     # VFM Stripe selection
-    yield from adjust_mirror_stripe(energy_kev, vfm, vfm_mirror_voltages)
+    yield from adjust_mirror_stripe(energy_kev, vfm, mirror_voltages)
     yield from bps.wait(DCM_GROUP)
 
     # VFM Adjust - for I03 this table always returns the same value

mx_bluesky/hyperion/device_setup_plans/read_hardware_for_setup.py

@@ -6,7 +6,6 @@ from dodal.devices.attenuator import Attenuator
 from dodal.devices.dcm import DCM
 from dodal.devices.eiger import EigerDetector
 from dodal.devices.flux import Flux
-from dodal.devices.robot import BartRobot
 from dodal.devices.s4_slit_gaps import S4SlitGaps
 from dodal.devices.smargon import Smargon
 from dodal.devices.synchrotron import Synchrotron
@@ -20,7 +19,7 @@ def read_hardware_pre_collection(
     undulator: Undulator,
     synchrotron: Synchrotron,
     s4_slit_gaps: S4SlitGaps,
-    robot: BartRobot,
+    dcm: DCM,
     smargon: Smargon,
 ):
     LOGGER.info("Reading status of beamline for callbacks, pre collection.")
@@ -29,11 +28,12 @@ def read_hardware_pre_collection(
     )  # gives name to event *descriptor* document
     yield from bps.read(undulator.current_gap)
     yield from bps.read(synchrotron.synchrotron_mode)
-    yield from bps.read(s4_slit_gaps.xgap)
-    yield from bps.read(s4_slit_gaps.ygap)
+    yield from bps.read(s4_slit_gaps.xgap)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+    yield from bps.read(s4_slit_gaps.ygap)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
     yield from bps.read(smargon.x)
     yield from bps.read(smargon.y)
     yield from bps.read(smargon.z)
+    yield from bps.read(dcm.energy_in_kev)
     yield from bps.save()
 
 
@@ -48,13 +48,7 @@ def read_hardware_during_collection(
     yield from bps.create(name=CONST.DESCRIPTORS.HARDWARE_READ_DURING)
     yield from bps.read(aperture_scatterguard)
     yield from bps.read(attenuator.actual_transmission)
-    yield from bps.read(flux.flux_reading)
+    yield from bps.read(flux.flux_reading)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
     yield from bps.read(dcm.energy_in_kev)
-    yield from bps.read(detector.bit_depth)
-    yield from bps.save()
-
-
-def read_hardware_for_zocalo(detector: EigerDetector):
-    yield from bps.create(name=CONST.DESCRIPTORS.ZOCALO_HW_READ)
-    yield from bps.read(detector.odin.file_writer.id)
+    yield from bps.read(detector.bit_depth)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
     yield from bps.save()

mx_bluesky/hyperion/device_setup_plans/setup_oav.py

@@ -1,11 +1,10 @@
 from functools import partial
 
 import bluesky.plan_stubs as bps
+from dodal.devices.areadetector.plugins.CAM import ColorMode
 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
 
@@ -56,19 +55,14 @@ def setup_pin_tip_detection_params(
 
 
 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)
+    yield from set_using_group(oav.cam.color_mode, ColorMode.RGB1)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+    yield from set_using_group(oav.cam.acquire_period, parameters.acquire_period)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+    yield from set_using_group(oav.cam.acquire_time, parameters.exposure)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+    yield from set_using_group(oav.cam.gain, parameters.gain)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
 
     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,
+        oav.zoom_controller,  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
         zoom_level_str,
         wait=True,
     )

mx_bluesky/hyperion/device_setup_plans/setup_panda.py

@@ -4,7 +4,7 @@ from importlib import resources
 from pathlib import Path
 
 import bluesky.plan_stubs as bps
-from blueapi.core import MsgGenerator
+from bluesky.utils 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
@@ -20,6 +20,7 @@ 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
+PULSE_WIDTH_US = 50
 
 
 class Enabled(Enum):
@@ -62,12 +63,12 @@ def _get_seq_table(
         An instance of SeqTable describing the panda sequencer table
     """
 
-    start_of_grid_x_counts = int(parameters.x_start * MM_TO_ENCODER_COUNTS)
+    start_of_grid_x_counts = int(parameters.x_start_mm * 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)
+        + (parameters.x_step_size_mm * (parameters.x_steps - 1) * MM_TO_ENCODER_COUNTS)
     )
 
     exposure_distance_x_counts = int(exposure_distance_mm * MM_TO_ENCODER_COUNTS)
@@ -76,8 +77,6 @@ def _get_seq_table(
 
     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
@@ -141,7 +140,7 @@ def setup_panda_for_flyscan(
     """
     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
+    assert sample_velocity_mm_per_s * exposure_time_s < parameters.x_step_size_mm
 
     yield from bps.stage(panda, group="panda-config")
 

mx_bluesky/hyperion/device_setup_plans/setup_zebra.py

@@ -3,16 +3,18 @@ from functools import wraps
 
 import bluesky.plan_stubs as bps
 import bluesky.preprocessors as bpp
-from blueapi.core import MsgGenerator
+from bluesky.utils import MsgGenerator
 from dodal.devices.zebra import (
     AUTO_SHUTTER_GATE,
-    AUTO_SHUTTER_INPUT,
+    AUTO_SHUTTER_INPUT_1,
+    AUTO_SHUTTER_INPUT_2,
     DISCONNECT,
     IN1_TTL,
     IN3_TTL,
     IN4_TTL,
     PC_GATE,
     PC_PULSE,
+    SOFT_IN1,
     TTL_DETECTOR,
     TTL_PANDA,
     TTL_XSPRESS3,
@@ -55,7 +57,7 @@ def bluesky_retry(func: Callable):
 
 
 def arm_zebra(zebra: Zebra):
-    yield from bps.abs_set(zebra.pc.arm, ArmDemand.ARM, wait=True)
+    yield from bps.abs_set(zebra.pc.arm, ArmDemand.ARM, wait=True)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
 
 
 def tidy_up_zebra_after_rotation_scan(
@@ -64,7 +66,7 @@ def tidy_up_zebra_after_rotation_scan(
     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.pc.arm, ArmDemand.DISARM, group=group)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
     yield from bps.abs_set(
         zebra_shutter.control_mode, ZebraShutterControl.MANUAL, group=group
     )
@@ -73,15 +75,45 @@ def tidy_up_zebra_after_rotation_scan(
 
 
 def set_shutter_auto_input(zebra: Zebra, input: int, group="set_shutter_trigger"):
-    """Set the input that the shutter uses when set to auto.
+    """Set the signal that controls the shutter. We use the second input to the
+    Zebra's AND2 gate for this input. ZebraShutter control mode must be in auto for this input to take control
 
     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
+        auto_shutter_control.sources[AUTO_SHUTTER_INPUT_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 AND2 gate should be used to control the shutter while in auto mode.
+    To do this, we need (AND2 = SOFT_IN1 AND input), where input is the zebra signal we want to control the shutter when in auto mode.
+    """
+    # See https://github.com/DiamondLightSource/dodal/issues/813 for better typing here.
+
+    # Set shutter to auto mode
+    yield from bps.abs_set(
+        zebra_shutter.control_mode, ZebraShutterControl.AUTO, group=group
+    )
+
+    # Set first input of AND2 gate 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_SHUTTER_GATE].sources[AUTO_SHUTTER_INPUT_1],
+        SOFT_IN1,
+        group=group,
+    )
+
+    # Set the second input of AND2 gate to the requested zebra input source
+    yield from set_shutter_auto_input(zebra, input, group=group)
+
+
 @bluesky_retry
 def setup_zebra_for_rotation(
     zebra: Zebra,
@@ -137,11 +169,10 @@ def setup_zebra_for_rotation(
     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
+    # Set shutter to automatic and to trigger via PC_GATE
+    yield from configure_zebra_and_shutter_for_auto_shutter(
+        zebra, zebra_shutter, PC_GATE, 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
@@ -159,11 +190,12 @@ def setup_zebra_for_gridscan(
     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
+    # Set shutter to automatic and to trigger via motion controller GPIO signal (IN4_TTL)
+    yield from configure_zebra_and_shutter_for_auto_shutter(
+        zebra, zebra_shutter, IN4_TTL, group=group
     )
-    yield from set_shutter_auto_input(zebra, IN4_TTL, group=group)
+
+    yield from bps.abs_set(zebra.output.out_pvs[TTL_DETECTOR], IN3_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)
 
@@ -198,11 +230,10 @@ def setup_zebra_for_panda_flyscan(
     # 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
+    # Set shutter to automatic and to trigger via motion controller GPIO signal (IN4_TTL)
+    yield from configure_zebra_and_shutter_for_auto_shutter(
+        zebra, zebra_shutter, IN4_TTL, 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)
 

mx_bluesky/hyperion/device_setup_plans/smargon.py

@@ -16,10 +16,10 @@ def move_smargon_warn_on_out_of_range(
             "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],
+        smargon.x,  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+        position[0],  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+        smargon.y,  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+        position[1],  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+        smargon.z,  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+        position[2],  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
     )

mx_bluesky/hyperion/device_setup_plans/utils.py

@@ -41,7 +41,7 @@ def start_preparing_data_collection_then_do_plan(
     """
 
     def wrapped_plan():
-        yield from bps.abs_set(eiger.do_arm, 1, group=group)
+        yield from bps.abs_set(eiger.do_arm, 1, group=group)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
         if detector_distance_mm:
             yield from set_detector_z_position(
                 detector_motion, detector_distance_mm, group
@@ -51,5 +51,5 @@ def start_preparing_data_collection_then_do_plan(
 
     yield from bpp.contingency_wrapper(
         wrapped_plan(),
-        except_plan=lambda e: (yield from bps.stop(eiger)),
+        except_plan=lambda e: (yield from bps.stop(eiger)),  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
     )

mx_bluesky/hyperion/device_setup_plans/xbpm_feedback.py

@@ -23,14 +23,14 @@ def _check_and_pause_feedback(
                                                turning XBPM feedback off.
 
     """
-    yield from bps.mv(attenuator, 1.0)
+    yield from bps.mv(attenuator, 1.0)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
     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)
+    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
 
 
 def _unpause_xbpm_feedback_and_set_transmission_to_1(
@@ -44,7 +44,7 @@ def _unpause_xbpm_feedback_and_set_transmission_to_1(
                                       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)
+    yield from bps.mv(xbpm_feedback.pause_feedback, Pause.RUN, attenuator, 1.0)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
 
 
 def transmission_and_xbpm_feedback_for_collection_wrapper(

mx_bluesky/hyperion/experiment_plans/__init__.py

@@ -9,6 +9,9 @@ from mx_bluesky.hyperion.experiment_plans.flyscan_xray_centre_plan import (
 from mx_bluesky.hyperion.experiment_plans.grid_detect_then_xray_centre_plan import (
     grid_detect_then_xray_centre,
 )
+from mx_bluesky.hyperion.experiment_plans.load_centre_collect_full_plan import (
+    load_centre_collect_full,
+)
 from mx_bluesky.hyperion.experiment_plans.pin_centre_then_xray_centre_plan import (
     pin_tip_centre_then_xray_centre,
 )
@@ -27,4 +30,5 @@ __all__ = [
     "pin_tip_centre_then_xray_centre",
     "multi_rotation_scan",
     "robot_load_then_centre",
+    "load_centre_collect_full",
 ]

mx_bluesky/hyperion/experiment_plans/change_aperture_then_move_plan.py

@@ -0,0 +1,83 @@
+import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
+import numpy
+from dodal.devices.aperturescatterguard import ApertureScatterguard, ApertureValue
+from dodal.devices.smargon import Smargon, StubPosition
+
+from mx_bluesky.common.utils.tracing import TRACER
+from mx_bluesky.hyperion.device_setup_plans.manipulate_sample import move_x_y_z
+from mx_bluesky.hyperion.experiment_plans.common.xrc_result import XRayCentreResult
+from mx_bluesky.hyperion.log import LOGGER
+from mx_bluesky.hyperion.parameters.gridscan import HyperionThreeDGridScan
+
+
+def change_aperture_then_move_to_xtal(
+    best_hit: XRayCentreResult,
+    smargon: Smargon,
+    aperture_scatterguard: ApertureScatterguard,
+    parameters: HyperionThreeDGridScan | None = None,
+):
+    """For the given x-ray centring result,
+    * Change the aperture so that the beam size is comparable to the crystal size
+    * Centre on the centre-of-mass
+    * Reset the stub offsets if specified by params"""
+    if best_hit.bounding_box_mm is not None:
+        bounding_box_size = numpy.abs(
+            best_hit.bounding_box_mm[1] - best_hit.bounding_box_mm[0]
+        )
+        with TRACER.start_span("change_aperture"):
+            yield from _set_aperture_for_bbox_mm(
+                aperture_scatterguard, bounding_box_size
+            )
+    else:
+        LOGGER.warning("No bounding box size received")
+
+    # once we have the results, go to the appropriate position
+    LOGGER.info("Moving to centre of mass.")
+    with TRACER.start_span("move_to_result"):
+        x, y, z = best_hit.centre_of_mass_mm
+        yield from move_x_y_z(smargon, x, y, z, wait=True)
+
+    # TODO support for setting stub offsets in multipin
+    # https://github.com/DiamondLightSource/mx-bluesky/issues/552
+    if parameters and parameters.FGS_params.set_stub_offsets:
+        LOGGER.info("Recentring smargon co-ordinate system to this point.")
+        yield from bps.mv(
+            # See: https://github.com/bluesky/bluesky/issues/1809
+            smargon.stub_offsets,  # type: ignore
+            StubPosition.CURRENT_AS_CENTER,  # type: ignore
+        )
+
+
+def _set_aperture_for_bbox_mm(
+    aperture_device: ApertureScatterguard, bbox_size_mm: list[float] | numpy.ndarray
+):
+    # TODO confirm correction factor see https://github.com/DiamondLightSource/mx-bluesky/issues/618
+    ASSUMED_BOX_SIZE_MM = 0.020
+    bbox_size_boxes = [round(mm / ASSUMED_BOX_SIZE_MM) for mm in bbox_size_mm]
+    yield from set_aperture_for_bbox_size(aperture_device, bbox_size_boxes)
+
+
+def set_aperture_for_bbox_size(
+    aperture_device: ApertureScatterguard,
+    bbox_size: list[int] | numpy.ndarray,
+):
+    # bbox_size is [x,y,z], for i03 we only care about x
+    new_selected_aperture = (
+        ApertureValue.MEDIUM if bbox_size[0] < 2 else ApertureValue.LARGE
+    )
+    LOGGER.info(
+        f"Setting aperture to {new_selected_aperture} based on bounding box size {bbox_size}."
+    )
+
+    @bpp.set_run_key_decorator("change_aperture")
+    @bpp.run_decorator(
+        md={
+            "subplan_name": "change_aperture",
+            "aperture_size": new_selected_aperture.value,
+        }
+    )
+    def set_aperture():
+        yield from bps.abs_set(aperture_device, new_selected_aperture)
+
+    yield from set_aperture()

mx_bluesky/hyperion/experiment_plans/common/xrc_result.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Callable, Sequence
+from functools import partial
+
+import numpy as np
+
+from mx_bluesky.common.parameters.components import (
+    MultiXtalSelection,
+    TopNByMaxCountSelection,
+)
+
+
+@dataclasses.dataclass
+class XRayCentreResult:
+    """Represents information about a hit from an X-ray centring."""
+
+    centre_of_mass_mm: np.ndarray
+    bounding_box_mm: tuple[np.ndarray, np.ndarray]
+    max_count: int
+    total_count: int
+
+    def __eq__(self, o):
+        return (
+            isinstance(o, XRayCentreResult)
+            and o.max_count == self.max_count
+            and o.total_count == self.total_count
+            and all(o.centre_of_mass_mm == self.centre_of_mass_mm)
+            and all(o.bounding_box_mm[0] == self.bounding_box_mm[0])
+            and all(o.bounding_box_mm[1] == self.bounding_box_mm[1])
+        )
+
+
+def top_n_by_max_count(
+    unfiltered: Sequence[XRayCentreResult], n: int
+) -> Sequence[XRayCentreResult]:
+    sorted_hits = sorted(unfiltered, key=lambda result: result.max_count, reverse=True)
+    return sorted_hits[:n]
+
+
+def resolve_selection_fn(
+    params: MultiXtalSelection,
+) -> Callable[[Sequence[XRayCentreResult]], Sequence[XRayCentreResult]]:
+    if isinstance(params, TopNByMaxCountSelection):
+        return partial(top_n_by_max_count, n=params.n)
+    raise ValueError(f"Invalid selection function {params.name}")

mx_bluesky/hyperion/experiment_plans/experiment_registry.py

@@ -5,6 +5,11 @@ from typing import TypedDict
 
 import mx_bluesky.hyperion.experiment_plans.flyscan_xray_centre_plan as flyscan_xray_centre_plan
 import mx_bluesky.hyperion.experiment_plans.rotation_scan_plan as rotation_scan_plan
+from mx_bluesky.common.parameters.gridscan import (
+    GridScanWithEdgeDetect,
+    PinTipCentreThenXrayCentre,
+    RobotLoadThenCentre,
+)
 from mx_bluesky.hyperion.experiment_plans import (
     grid_detect_then_xray_centre_plan,
     load_centre_collect_full_plan,
@@ -18,12 +23,7 @@ from mx_bluesky.hyperion.external_interaction.callbacks.common.callback_util imp
     create_robot_load_and_centre_callbacks,
     create_rotation_callbacks,
 )
-from mx_bluesky.hyperion.parameters.gridscan import (
-    GridScanWithEdgeDetect,
-    PinTipCentreThenXrayCentre,
-    RobotLoadThenCentre,
-    ThreeDGridScan,
-)
+from mx_bluesky.hyperion.parameters.gridscan import HyperionThreeDGridScan
 from mx_bluesky.hyperion.parameters.load_centre_collect import LoadCentreCollect
 from mx_bluesky.hyperion.parameters.rotation import MultiRotationScan, RotationScan
 
@@ -39,7 +39,7 @@ def do_nothing():
 class ExperimentRegistryEntry(TypedDict):
     setup: Callable
     param_type: type[
-        ThreeDGridScan
+        HyperionThreeDGridScan
         | GridScanWithEdgeDetect
         | RotationScan
         | MultiRotationScan
@@ -53,7 +53,7 @@ class ExperimentRegistryEntry(TypedDict):
 PLAN_REGISTRY: dict[str, ExperimentRegistryEntry] = {
     "flyscan_xray_centre": {
         "setup": flyscan_xray_centre_plan.create_devices,
-        "param_type": ThreeDGridScan,
+        "param_type": HyperionThreeDGridScan,
         "callbacks_factory": create_gridscan_callbacks,
     },
     "grid_detect_then_xray_centre": {
@@ -81,7 +81,7 @@ PLAN_REGISTRY: dict[str, ExperimentRegistryEntry] = {
         "param_type": MultiRotationScan,
         "callbacks_factory": create_rotation_callbacks,
     },
-    "load_centre_collect_full_plan": {
+    "load_centre_collect_full": {
         "setup": load_centre_collect_full_plan.create_devices,
         "param_type": LoadCentreCollect,
         "callbacks_factory": create_load_centre_collect_callbacks,