mx-bluesky 1.5.0__py3-none-any.whl → 1.5.1__py3-none-any.whl

mx_bluesky/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.5.0'
-__version_tuple__ = version_tuple = (1, 5, 0)
+__version__ = version = '1.5.1'
+__version_tuple__ = version_tuple = (1, 5, 1)

mx_bluesky/beamlines/i24/serial/extruder/i24ssx_Extruder_Collect_py3v2.py

@@ -23,8 +23,8 @@ from dodal.devices.i24.beamstop import Beamstop
 from dodal.devices.i24.dcm import DCM
 from dodal.devices.i24.dual_backlight import DualBacklight
 from dodal.devices.i24.focus_mirrors import FocusMirrorsMode
-from dodal.devices.i24.i24_detector_motion import DetectorMotion
 from dodal.devices.i24.pilatus_metadata import PilatusMetadata
+from dodal.devices.motors import YZStage
 from dodal.devices.zebra.zebra import Zebra
 
 from mx_bluesky.beamlines.i24.serial.dcid import (
@@ -70,7 +70,7 @@ def flush_print(text):
 
 @log_on_entry
 def initialise_extruder(
-    detector_stage: DetectorMotion = inject("detector_motion"),
+    detector_stage: YZStage = inject("detector_motion"),
 ) -> MsgGenerator:
     SSX_LOGGER.info("Initialise Parameters for extruder data collection on I24.")
 
@@ -98,7 +98,7 @@ def initialise_extruder(
 def laser_check(
     mode: str,
     zebra: Zebra = inject("zebra"),
-    detector_stage: DetectorMotion = inject("detector_motion"),
+    detector_stage: YZStage = inject("detector_motion"),
 ) -> MsgGenerator:
     """Plan to open the shutter and check the laser beam from the viewer by pressing \
         'Laser On' and 'Laser Off' buttons on the edm.
@@ -138,7 +138,7 @@ def laser_check(
 
 @log_on_entry
 def enter_hutch(
-    detector_stage: DetectorMotion = inject("detector_motion"),
+    detector_stage: YZStage = inject("detector_motion"),
 ) -> MsgGenerator:
     """Move the detector stage before entering hutch."""
     yield from bps.mv(detector_stage.z, SAFE_DET_Z)
@@ -146,12 +146,12 @@ def enter_hutch(
 
 
 @log_on_entry
-def read_parameters(detector_stage: DetectorMotion, attenuator: ReadOnlyAttenuator):
+def read_parameters(detector_stage: YZStage, attenuator: ReadOnlyAttenuator):
     """ Read the parameters from user input and create the parameter model for an \
         extruder collection.
 
     Args:
-        detector_stage (DetectorMotion): The detector stage device.
+        detector_stage (YZStage): The detector stage device.
         attenuator (ReadOnlyAttenuator): A read-only attenuator device to get the \
             transmission value.
 
@@ -208,7 +208,7 @@ def main_extruder_plan(
     aperture: Aperture,
     backlight: DualBacklight,
     beamstop: Beamstop,
-    detector_stage: DetectorMotion,
+    detector_stage: YZStage,
     shutter: HutchShutter,
     dcm: DCM,
     mirrors: FocusMirrorsMode,
@@ -500,7 +500,7 @@ def run_extruder_plan(
     aperture: Aperture = inject("aperture"),
     backlight: DualBacklight = inject("backlight"),
     beamstop: Beamstop = inject("beamstop"),
-    detector_stage: DetectorMotion = inject("detector_motion"),
+    detector_stage: YZStage = inject("detector_motion"),
     shutter: HutchShutter = inject("shutter"),
     dcm: DCM = inject("dcm"),
     mirrors: FocusMirrorsMode = inject("focus_mirrors"),

mx_bluesky/beamlines/i24/serial/fixed_target/i24ssx_Chip_Collect_py3v1.py

@@ -18,9 +18,9 @@ from dodal.devices.i24.beamstop import Beamstop
 from dodal.devices.i24.dcm import DCM
 from dodal.devices.i24.dual_backlight import DualBacklight
 from dodal.devices.i24.focus_mirrors import FocusMirrorsMode
-from dodal.devices.i24.i24_detector_motion import DetectorMotion
 from dodal.devices.i24.pilatus_metadata import PilatusMetadata
 from dodal.devices.i24.pmac import PMAC
+from dodal.devices.motors import YZStage
 from dodal.devices.zebra.zebra import Zebra
 
 from mx_bluesky.beamlines.i24.serial.dcid import (
@@ -269,7 +269,7 @@ def start_i24(
     aperture: Aperture,
     backlight: DualBacklight,
     beamstop: Beamstop,
-    detector_stage: DetectorMotion,
+    detector_stage: YZStage,
     shutter: HutchShutter,
     parameters: FixedTargetParameters,
     dcm: DCM,
@@ -504,7 +504,7 @@ def main_fixed_target_plan(
     aperture: Aperture,
     backlight: DualBacklight,
     beamstop: Beamstop,
-    detector_stage: DetectorMotion,
+    detector_stage: YZStage,
     shutter: HutchShutter,
     dcm: DCM,
     mirrors: FocusMirrorsMode,
@@ -655,7 +655,7 @@ def run_fixed_target_plan(
     aperture: Aperture = inject("aperture"),
     backlight: DualBacklight = inject("backlight"),
     beamstop: Beamstop = inject("beamstop"),
-    detector_stage: DetectorMotion = inject("detector_motion"),
+    detector_stage: YZStage = inject("detector_motion"),
     shutter: HutchShutter = inject("shutter"),
     dcm: DCM = inject("dcm"),
     mirrors: FocusMirrorsMode = inject("focus_mirrors"),
@@ -707,7 +707,7 @@ def run_plan_in_wrapper(
     aperture: Aperture,
     backlight: DualBacklight,
     beamstop: Beamstop,
-    detector_stage: DetectorMotion,
+    detector_stage: YZStage,
     shutter: HutchShutter,
     dcm: DCM,
     mirrors: FocusMirrorsMode,

mx_bluesky/beamlines/i24/serial/fixed_target/i24ssx_Chip_Manager_py3v1.py

@@ -16,8 +16,8 @@ from dodal.common import inject
 from dodal.devices.attenuator.attenuator import ReadOnlyAttenuator
 from dodal.devices.i24.beamstop import Beamstop, BeamstopPositions
 from dodal.devices.i24.dual_backlight import BacklightPositions, DualBacklight
-from dodal.devices.i24.i24_detector_motion import DetectorMotion
 from dodal.devices.i24.pmac import PMAC, EncReset, LaserSettings
+from dodal.devices.motors import YZStage
 
 from mx_bluesky.beamlines.i24.serial.fixed_target.ft_utils import (
     ChipType,
@@ -121,14 +121,14 @@ def _is_checker_pattern() -> bool:
 
 @log_on_entry
 def read_parameters(
-    detector_stage: DetectorMotion,
+    detector_stage: YZStage,
     attenuator: ReadOnlyAttenuator,
 ) -> MsgGenerator:
     """ Read the parameters from user input and create the parameter model for a fixed \
         target collection.
 
     Args:
-        detector_stage (DetectorMotion): The detector stage device.
+        detector_stage (YZStage): The detector stage device.
         attenuator (ReadOnlyAttenuator): A read-only attenuator device to get the \
             transmission value.
 
@@ -566,7 +566,7 @@ def moveto_preset(
     pmac: PMAC = inject("pmac"),
     beamstop: Beamstop = inject("beamstop"),
     backlight: DualBacklight = inject("backlight"),
-    det_stage: DetectorMotion = inject("detector_motion"),
+    det_stage: YZStage = inject("detector_motion"),
 ) -> MsgGenerator:
     # Non Chip Specific Move
     if place == "zero":

mx_bluesky/beamlines/i24/serial/setup_beamline/setup_beamline.py

@@ -7,7 +7,7 @@ from dodal.devices.i24.beam_center import DetectorBeamCenter
 from dodal.devices.i24.beamstop import Beamstop, BeamstopPositions
 from dodal.devices.i24.dcm import DCM
 from dodal.devices.i24.dual_backlight import BacklightPositions, DualBacklight
-from dodal.devices.i24.i24_detector_motion import DetectorMotion
+from dodal.devices.motors import YZStage
 from dodal.devices.util.lookup_tables import (
     linear_interpolation_lut,
     parse_lookup_table,
@@ -69,7 +69,7 @@ def setup_beamline_for_collection_plan(
 
 
 def move_detector_stage_to_position_plan(
-    detector_stage: DetectorMotion,
+    detector_stage: YZStage,
     detector_distance: float,
 ):
     SSX_LOGGER.debug("Setup beamline: moving detector stage.")

mx_bluesky/beamlines/i24/serial/setup_beamline/setup_detector.py

@@ -8,7 +8,7 @@ from enum import IntEnum
 import bluesky.plan_stubs as bps
 from bluesky.utils import Msg, MsgGenerator
 from dodal.common import inject
-from dodal.devices.i24.i24_detector_motion import DetectorMotion
+from dodal.devices.motors import YZStage
 
 from mx_bluesky.beamlines.i24.serial.log import SSX_LOGGER
 from mx_bluesky.beamlines.i24.serial.parameters import SSXType
@@ -38,9 +38,9 @@ class UnknownDetectorType(Exception):
     pass
 
 
-def get_detector_type(detector_stage: DetectorMotion) -> Generator[Msg, None, Detector]:
+def get_detector_type(detector_stage: YZStage) -> Generator[Msg, None, Detector]:
     det_y = yield from bps.rd(detector_stage.y)
-    # DetectorMotion should also be used for this.
+    # YZStage should also be used for this.
     # This should be part of https://github.com/DiamondLightSource/mx_bluesky/issues/51
     if float(det_y) < Eiger.det_y_threshold:
         SSX_LOGGER.info("Eiger detector in use.")
@@ -53,7 +53,7 @@ def get_detector_type(detector_stage: DetectorMotion) -> Generator[Msg, None, De
         raise UnknownDetectorType("Detector not found.")
 
 
-def _move_detector_stage(detector_stage: DetectorMotion, target: float) -> MsgGenerator:
+def _move_detector_stage(detector_stage: YZStage, target: float) -> MsgGenerator:
     SSX_LOGGER.info(f"Moving detector stage to target position: {target}.")
     yield from bps.mv(detector_stage.y, target)
 
@@ -82,7 +82,7 @@ def _get_requested_detector(det_type_pv: str) -> str:
 
 
 def setup_detector_stage(
-    expt_type: SSXType, detector_stage: DetectorMotion = inject("detector_motion")
+    expt_type: SSXType, detector_stage: YZStage = inject("detector_motion")
 ) -> MsgGenerator:
     # Grab the correct PV depending on experiment
     # Its value is set with MUX on edm screen

mx_bluesky/beamlines/i24/serial/web_gui_plans/general_plans.py

@@ -13,9 +13,9 @@ from dodal.devices.i24.beamstop import Beamstop
 from dodal.devices.i24.dcm import DCM
 from dodal.devices.i24.dual_backlight import BacklightPositions, DualBacklight
 from dodal.devices.i24.focus_mirrors import FocusMirrorsMode
-from dodal.devices.i24.i24_detector_motion import DetectorMotion
 from dodal.devices.i24.pilatus_metadata import PilatusMetadata
 from dodal.devices.i24.pmac import PMAC
+from dodal.devices.motors import YZStage
 from dodal.devices.oav.oav_detector import OAVBeamCentreFile
 from dodal.devices.zebra.zebra import Zebra
 
@@ -104,7 +104,7 @@ def gui_sleep(sec: int) -> MsgGenerator:
 @bpp.run_decorator()
 def gui_move_detector(
     det: Literal["eiger", "pilatus"],
-    detector_stage: DetectorMotion = inject("detector_motion"),
+    detector_stage: YZStage = inject("detector_motion"),
 ) -> MsgGenerator:
     det_y_target = Eiger.det_y_target if det == "eiger" else Pilatus.det_y_target
     yield from _move_detector_stage(detector_stage, det_y_target)
@@ -134,7 +134,7 @@ def gui_run_chip_collection(
     aperture: Aperture = inject("aperture"),
     backlight: DualBacklight = inject("backlight"),
     beamstop: Beamstop = inject("beamstop"),
-    detector_stage: DetectorMotion = inject("detector_motion"),
+    detector_stage: YZStage = inject("detector_motion"),
     shutter: HutchShutter = inject("shutter"),
     dcm: DCM = inject("dcm"),
     mirrors: FocusMirrorsMode = inject("focus_mirrors"),

mx_bluesky/common/device_setup_plans/robot_load_unload.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
+from bluesky.utils import MsgGenerator
+from dodal.devices.aperturescatterguard import ApertureScatterguard, ApertureValue
+from dodal.devices.motors import XYZStage
+from dodal.devices.robot import BartRobot
+from dodal.devices.smargon import CombinedMove, Smargon, StubPosition
+from dodal.plan_stubs.motor_utils import MoveTooLarge, home_and_reset_wrapper
+
+from mx_bluesky.common.utils.log import LOGGER
+from mx_bluesky.hyperion.parameters.constants import CONST
+
+
+def wait_for_smargon_not_disabled(smargon: Smargon, timeout=60):
+    """Waits for the smargon disabled flag to go low. The robot hardware is responsible
+    for setting this to low when it is safe to move. It does this through a physical
+    connection between the robot and the smargon.
+    """
+    LOGGER.info("Waiting for smargon enabled")
+    SLEEP_PER_CHECK = 0.1
+    times_to_check = int(timeout / SLEEP_PER_CHECK)
+    for _ in range(times_to_check):
+        smargon_disabled = yield from bps.rd(smargon.disabled)
+        if not smargon_disabled:
+            LOGGER.info("Smargon now enabled")
+            return
+        yield from bps.sleep(SLEEP_PER_CHECK)
+    raise TimeoutError(
+        "Timed out waiting for smargon to become enabled after robot load"
+    )
+
+
+def _raise_exception_if_moved_out_of_cryojet(exception):
+    yield from bps.null()
+    if isinstance(exception, MoveTooLarge):
+        raise Exception(
+            f"Moving {exception.axis} back to {exception.position} after \
+                        robot load would move it out of the cryojet. The max safe \
+                        distance is {exception.maximum_move}"
+        )
+
+
+def do_plan_while_lower_gonio_at_home(plan: MsgGenerator, lower_gonio: XYZStage):
+    """Moves the lower gonio to home then performs the provided plan and moves it back.
+
+    The lower gonio must be in the correct position for the robot load and we
+    want to put it back afterwards. Note we don't need to wait for the move as the robot
+    is interlocked to the lower gonio and the move is quicker than the robot takes to
+    get to the load position.
+
+    Args:
+        plan (MsgGenerator): The plan to run while the lower gonio is at home.
+        lower_gonio (XYZStage): The lower gonio to home.
+    """
+    yield from bpp.contingency_wrapper(
+        home_and_reset_wrapper(
+            plan,
+            lower_gonio,
+            BartRobot.LOAD_TOLERANCE_MM,
+            CONST.HARDWARE.CRYOJET_MARGIN_MM,
+            "lower_gonio",
+            wait_for_all=False,
+        ),
+        except_plan=_raise_exception_if_moved_out_of_cryojet,
+    )
+    return "reset-lower_gonio"
+
+
+def prepare_for_robot_load(
+    aperture_scatterguard: ApertureScatterguard, smargon: Smargon
+):
+    yield from bps.abs_set(
+        aperture_scatterguard.selected_aperture,
+        ApertureValue.OUT_OF_BEAM,
+        group="prepare_robot_load",
+    )
+
+    yield from bps.mv(smargon.stub_offsets, StubPosition.RESET_TO_ROBOT_LOAD)
+
+    yield from bps.mv(smargon, CombinedMove(x=0, y=0, z=0, chi=0, phi=0, omega=0))
+
+    yield from bps.wait("prepare_robot_load")
+
+
+def robot_unload(
+    robot: BartRobot,
+    smargon: Smargon,
+    aperture_scatterguard: ApertureScatterguard,
+    lower_gonio: XYZStage,
+    visit: str,
+):
+    """Unloads the currently mounted pin into the location that it was loaded from. The
+    loaded location is stored on the robot and so need not be provided.
+    """
+    yield from prepare_for_robot_load(aperture_scatterguard, smargon)
+    sample_id = yield from bps.rd(robot.sample_id)
+
+    @bpp.run_decorator(
+        md={
+            "subplan_name": CONST.PLAN.ROBOT_UNLOAD,
+            "metadata": {"visit": visit, "sample_id": sample_id},
+            "activate_callbacks": [
+                "RobotLoadISPyBCallback",
+            ],
+        },
+    )
+    def do_robot_unload_and_send_to_ispyb():
+        yield from bps.create(name=CONST.DESCRIPTORS.ROBOT_UPDATE)
+        yield from bps.read(robot)
+        yield from bps.save()
+
+        def _unload():
+            yield from bps.trigger(robot.unload, wait=True)
+            yield from wait_for_smargon_not_disabled(smargon)
+
+        gonio_finished = yield from do_plan_while_lower_gonio_at_home(
+            _unload(), lower_gonio
+        )
+        yield from bps.wait(gonio_finished)
+
+    yield from do_robot_unload_and_send_to_ispyb()

mx_bluesky/common/experiment_plans/common_flyscan_xray_centre_plan.py

@@ -288,6 +288,7 @@ def _xrc_result_in_boxes_to_result_in_mm(
         ),
         max_count=xrc_result["max_count"],
         total_count=xrc_result["total_count"],
+        sample_id=xrc_result["sample_id"],
     )
 
 

mx_bluesky/common/external_interaction/ispyb/exp_eye_store.py

@@ -1,7 +1,9 @@
 import configparser
 from dataclasses import dataclass
 from enum import StrEnum
+from typing import Any, Literal
 
+from event_model.documents import Event
 from requests import JSONDecodeError, patch, post
 from requests.auth import AuthBase
 
@@ -63,6 +65,19 @@ assert all(len(value) <= 20 for value in BLSampleStatus), (
 )
 
 
+def create_update_data_from_event_doc(
+    mapping: dict[str, str], event: Event
+) -> dict[str, Any]:
+    """Given a mapping between bluesky event data and an event itself this function will
+    create a dict that can be used to update exp-eye."""
+    event_data = event["data"]
+    return {
+        target_key: event_data[source_key]
+        for source_key, target_key in mapping.items()
+        if source_key in event_data
+    }
+
+
 class ExpeyeInteraction:
     """Exposes functionality from the Expeye core API"""
 
@@ -74,24 +89,22 @@ class ExpeyeInteraction:
         self._base_url = url
         self._auth = BearerAuth(token)
 
-    def start_load(
+    def start_robot_action(
         self,
+        action_type: Literal["LOAD", "UNLOAD"],
         proposal_reference: str,
         visit_number: int,
         sample_id: int,
-        dewar_location: int,
-        container_location: int,
     ) -> RobotActionID:
-        """Create a robot load entry in ispyb.
+        """Create a robot action entry in ispyb.
 
         Args:
+            action_type ("LOAD" | "UNLOAD"): The robot action being performed
             proposal_reference (str): The proposal of the experiment e.g. cm37235
             visit_number (int): The visit number for the proposal, usually this can be
                                 found added to the end of the proposal e.g. the data for
                                 visit number 2 of proposal cm37235 is in cm37235-2
             sample_id (int): The id of the sample in the database
-            dewar_location (int): Which puck in the dewar the sample is in
-            container_location (int): Which pin in that puck has the sample
 
         Returns:
             RobotActionID: The id of the robot load action that is created
@@ -102,39 +115,28 @@ class ExpeyeInteraction:
 
         data = {
             "startTimestamp": get_current_time_string(),
+            "actionType": action_type,
             "sampleId": sample_id,
-            "actionType": "LOAD",
-            "containerLocation": container_location,
-            "dewarLocation": dewar_location,
         }
         response = _send_and_get_response(self._auth, url, data, post)
         return response["robotActionId"]
 
-    def update_barcode_and_snapshots(
+    def update_robot_action(
         self,
         action_id: RobotActionID,
-        barcode: str,
-        snapshot_before_path: str,
-        snapshot_after_path: str,
+        data: dict[str, Any],
     ):
-        """Update the barcode and snapshots of an existing robot action.
+        """Update an existing robot action to contain additional info.
 
         Args:
             action_id (RobotActionID): The id of the action to update
-            barcode (str): The barcode to give the action
-            snapshot_before_path (str): Path to the snapshot before robot load
-            snapshot_after_path (str): Path to the snapshot after robot load
+            data (dict): The data to update with, where the keys match those expected
+                         by exp-eye.
         """
         url = self._base_url + self.UPDATE_ROBOT_ACTION.format(action_id=action_id)
-
-        data = {
-            "sampleBarcode": barcode,
-            "xtalSnapshotBefore": snapshot_before_path,
-            "xtalSnapshotAfter": snapshot_after_path,
-        }
         _send_and_get_response(self._auth, url, data, patch)
 
-    def end_load(self, action_id: RobotActionID, status: str, reason: str):
+    def end_robot_action(self, action_id: RobotActionID, status: str, reason: str):
         """Finish an existing robot action, providing final information about how it went
 
         Args:

mx_bluesky/common/external_interaction/ispyb/ispyb_store.py

@@ -41,7 +41,6 @@ class IspybIds(BaseModel):
 class StoreInIspyb:
     def __init__(self, ispyb_config: str) -> None:
         self.ISPYB_CONFIG_PATH: str = ispyb_config
-        self._data_collection_group_id: int | None
 
     def begin_deposition(
         self,

mx_bluesky/common/external_interaction/nexus/write_nexus.py

@@ -20,13 +20,13 @@ from mx_bluesky.common.external_interaction.nexus.nexus_utils import (
     create_goniometer_axes,
     get_start_and_predicted_end_time,
 )
-from mx_bluesky.common.parameters.components import DiffractionExperimentWithSample
+from mx_bluesky.common.parameters.components import DiffractionExperiment
 
 
 class NexusWriter:
     def __init__(
         self,
-        parameters: DiffractionExperimentWithSample,
+        parameters: DiffractionExperiment,
         data_shape: tuple[int, int, int],
         scan_points: AxesPoints,
         *,

mx_bluesky/common/parameters/components.py

@@ -238,9 +238,14 @@ class TopNByMaxCountSelection(MultiXtalSelection):
     n: int
 
 
+class TopNByMaxCountForEachSampleSelection(MultiXtalSelection):
+    name: Literal["TopNByMaxCountForEachSample"] = "TopNByMaxCountForEachSample"  #  pyright: ignore [reportIncompatibleVariableOverride]
+    n: int
+
+
 class WithCentreSelection(BaseModel):
-    select_centres: TopNByMaxCountSelection = Field(
-        discriminator="name", default=TopNByMaxCountSelection(n=1)
+    select_centres: TopNByMaxCountSelection | TopNByMaxCountForEachSampleSelection = (
+        Field(discriminator="name", default=TopNByMaxCountSelection(n=1))
     )
 
     @property

mx_bluesky/common/parameters/constants.py

@@ -15,8 +15,8 @@ TEST_MODE = BEAMLINE == "test"
 
 @dataclass(frozen=True)
 class DocDescriptorNames:
-    # Robot load event descriptor
-    ROBOT_LOAD = "robot_load"
+    # Robot load/unload event descriptor
+    ROBOT_UPDATE = "robot_update"
     # For callbacks to use
     OAV_ROTATION_SNAPSHOT_TRIGGERED = "rotation_snapshot_triggered"
     OAV_GRID_SNAPSHOT_TRIGGERED = "snapshot_to_ispyb"
@@ -41,8 +41,9 @@ class OavConstants:
 @dataclass(frozen=True)
 class PlanNameConstants:
     LOAD_CENTRE_COLLECT = "load_centre_collect"
-    # Robot load subplan
+    # Robot subplans
     ROBOT_LOAD = "robot_load"
+    ROBOT_UNLOAD = "robot_unload"
     # Gridscan
     GRID_DETECT_AND_DO_GRIDSCAN = "grid_detect_and_do_gridscan"
     GRID_DETECT_INNER = "grid_detect"

mx_bluesky/common/xrc_result.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import dataclasses
+from collections import defaultdict
 from collections.abc import Callable, Sequence
 from functools import partial
 
@@ -10,6 +11,7 @@ from event_model import RunStart
 
 from mx_bluesky.common.parameters.components import (
     MultiXtalSelection,
+    TopNByMaxCountForEachSampleSelection,
     TopNByMaxCountSelection,
 )
 
@@ -39,18 +41,21 @@ class XRayCentreResult:
             containing the crystal
         max_count: The maximum spot count encountered in any one grid box in the crystal
         total_count: The total count across all boxes in the crystal.
+        sample_id: The sample id associated with the centre.
     """
 
     centre_of_mass_mm: np.ndarray
     bounding_box_mm: tuple[np.ndarray, np.ndarray]
     max_count: int
     total_count: int
+    sample_id: int | None
 
     def __eq__(self, o):
         return (
             isinstance(o, XRayCentreResult)
             and o.max_count == self.max_count
             and o.total_count == self.total_count
+            and o.sample_id == self.sample_id
             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])
@@ -64,9 +69,27 @@ def top_n_by_max_count(
     return sorted_hits[:n]
 
 
+def top_n_by_max_count_for_each_sample(
+    unfiltered: Sequence[XRayCentreResult], n: int
+) -> Sequence[XRayCentreResult]:
+    xrc_results_by_sample_id: dict[int | None, list[XRayCentreResult]] = defaultdict(
+        list[XRayCentreResult]
+    )
+    for result in unfiltered:
+        xrc_results_by_sample_id[result.sample_id].append(result)
+    return [
+        result
+        for results in xrc_results_by_sample_id.values()
+        for result in sorted(results, key=lambda x: x.max_count, reverse=True)[: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)
+    match params:
+        case TopNByMaxCountSelection():
+            return partial(top_n_by_max_count, n=params.n)
+        case TopNByMaxCountForEachSampleSelection():
+            return partial(top_n_by_max_count_for_each_sample, n=params.n)
     raise ValueError(f"Invalid selection function {params.name}")

mx_bluesky/hyperion/experiment_plans/hyperion_flyscan_xray_centre_plan.py

@@ -102,7 +102,13 @@ def _generic_tidy(
 
     # Turn off dev/shm streaming to avoid filling disk, see https://github.com/DiamondLightSource/hyperion/issues/1395
     LOGGER.info("Turning off Eiger dev/shm streaming")
-    yield from bps.abs_set(xrc_composite.eiger.odin.fan.dev_shm_enable, 0)  # type: ignore # Fix types in ophyd-async (https://github.com/DiamondLightSource/mx-bluesky/issues/855)
+    # Fix types in ophyd-async (https://github.com/DiamondLightSource/mx-bluesky/issues/855)
+    yield from bps.abs_set(
+        xrc_composite.eiger.odin.fan.dev_shm_enable,  # type: ignore
+        0,
+        group=group,
+        wait=wait,
+    )
 
 
 def _panda_tidy(xrc_composite: HyperionFlyScanXRayCentreComposite):