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

mx_bluesky/common/{plans → experiment_plans}/common_flyscan_xray_centre_plan.py

@@ -7,35 +7,32 @@ 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.experiment_plans.inner_plans.do_fgs import (
+    ZOCALO_STAGE_GROUP,
+    kickoff_and_complete_gridscan,
+)
+from mx_bluesky.common.experiment_plans.read_hardware import (
+    read_hardware_plan,
+)
 from mx_bluesky.common.parameters.constants import (
     DocDescriptorNames,
     GridscanParamConstants,
     PlanGroupCheckpointConstants,
     PlanNameConstants,
 )
+from mx_bluesky.common.parameters.device_composites import FlyScanEssentialDevices
 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,
@@ -45,14 +42,6 @@ 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]
@@ -251,6 +240,9 @@ def run_gridscan(
         parameters.scan_indices,
         plan_during_collection=beamline_specific.read_during_collection_plan,
     )
+
+    # GDA's gridscans requires Z steps to be at 0, so make sure we leave this device
+    # in a GDA-happy state.
     yield from bps.abs_set(beamline_specific.fgs_motors.z_steps, 0, wait=False)
 
 
@@ -296,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/{hyperion/experiment_plans/grid_detect_then_xray_centre_plan.py → common/experiment_plans/common_grid_detect_then_xray_centre_plan.py}

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 from pathlib import Path
+from typing import Protocol, TypeVar
 
-from blueapi.core import BlueskyContext
 from bluesky import plan_stubs as bps
 from bluesky import preprocessors as bpp
 from bluesky.preprocessors import subs_decorator
@@ -10,20 +10,27 @@ from bluesky.utils import MsgGenerator
 from dodal.devices.backlight import BacklightPosition
 from dodal.devices.eiger import EigerDetector
 from dodal.devices.oav.oav_parameters import OAVParameters
-from dodal.plans.preprocessors.verify_undulator_gap import (
-    verify_undulator_gap_before_run_decorator,
-)
 
 from mx_bluesky.common.device_setup_plans.manipulate_sample import (
     move_aperture_if_required,
 )
+from mx_bluesky.common.device_setup_plans.utils import (
+    start_preparing_data_collection_then_do_plan,
+)
 from mx_bluesky.common.experiment_plans.change_aperture_then_move_plan import (
     change_aperture_then_move_to_xtal,
 )
+from mx_bluesky.common.experiment_plans.common_flyscan_xray_centre_plan import (
+    BeamlineSpecificFGSFeatures,
+    common_flyscan_xray_centre,
+)
 from mx_bluesky.common.experiment_plans.oav_grid_detection_plan import (
     OavGridDetectionComposite,
     grid_detection_plan,
 )
+from mx_bluesky.common.experiment_plans.oav_snapshot_plan import (
+    setup_beamline_for_OAV,
+)
 from mx_bluesky.common.external_interaction.callbacks.common.grid_detection_callback import (
     GridDetectionCallback,
     GridParamUpdate,
@@ -31,62 +38,97 @@ from mx_bluesky.common.external_interaction.callbacks.common.grid_detection_call
 from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import (
     ispyb_activation_wrapper,
 )
-from mx_bluesky.common.parameters.constants import OavConstants
-from mx_bluesky.common.parameters.gridscan import GridCommon
-from mx_bluesky.common.plans.common_flyscan_xray_centre_plan import (
-    BeamlineSpecificFGSFeatures,
-    common_flyscan_xray_centre,
+from mx_bluesky.common.parameters.constants import (
+    OavConstants,
+    PlanGroupCheckpointConstants,
 )
-from mx_bluesky.common.preprocessors.preprocessors import (
-    transmission_and_xbpm_feedback_for_collection_decorator,
+from mx_bluesky.common.parameters.device_composites import (
+    FlyScanEssentialDevices,
+    GridDetectThenXRayCentreComposite,
 )
-from mx_bluesky.common.utils.context import device_composite_from_context
+from mx_bluesky.common.parameters.gridscan import GridCommon, SpecifiedThreeDGridScan
 from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.common.xrc_result import XRayCentreEventHandler
-from mx_bluesky.hyperion.device_setup_plans.utils import (
-    start_preparing_data_collection_then_do_plan,
-)
-from mx_bluesky.hyperion.experiment_plans.hyperion_flyscan_xray_centre_plan import (
-    construct_hyperion_specific_features,
-)
-from mx_bluesky.hyperion.experiment_plans.oav_snapshot_plan import (
-    setup_beamline_for_OAV,
-)
-from mx_bluesky.hyperion.parameters.constants import CONST
-from mx_bluesky.hyperion.parameters.device_composites import (
-    GridDetectThenXRayCentreComposite,
-    HyperionFlyScanXRayCentreComposite,
+
+TFlyScanEssentialDevices = TypeVar(
+    "TFlyScanEssentialDevices", bound=FlyScanEssentialDevices, contravariant=True
 )
-from mx_bluesky.hyperion.parameters.gridscan import (
-    GridScanWithEdgeDetect,
-    HyperionSpecifiedThreeDGridScan,
+TSpecifiedThreeDGridScan = TypeVar(
+    "TSpecifiedThreeDGridScan", bound=SpecifiedThreeDGridScan, contravariant=True
 )
 
 
-def create_devices(context: BlueskyContext) -> GridDetectThenXRayCentreComposite:
-    return device_composite_from_context(context, GridDetectThenXRayCentreComposite)
+def grid_detect_then_xray_centre(
+    composite: GridDetectThenXRayCentreComposite,
+    parameters: GridCommon,
+    xrc_params_type: type[SpecifiedThreeDGridScan],
+    construct_beamline_specific: ConstructBeamlineSpecificFeatures,
+    oav_config: str = OavConstants.OAV_CONFIG_JSON,
+) -> MsgGenerator:
+    """
+    A plan which combines the collection of snapshots from the OAV and the determination
+    of the grid dimensions to use for the following grid scan.
+    """
 
+    eiger: EigerDetector = composite.eiger
 
-def create_parameters_for_flyscan_xray_centre(
-    grid_scan_with_edge_params: GridCommon,
-    grid_parameters: GridParamUpdate,
-) -> HyperionSpecifiedThreeDGridScan:
-    params_json = grid_scan_with_edge_params.model_dump()
-    params_json.update(grid_parameters)
-    flyscan_xray_centre_parameters = HyperionSpecifiedThreeDGridScan(**params_json)
-    LOGGER.info(f"Parameters for FGS: {flyscan_xray_centre_parameters}")
-    return flyscan_xray_centre_parameters
+    eiger.set_detector_parameters(parameters.detector_params)
+
+    oav_params = OAVParameters("xrayCentring", oav_config)
+
+    flyscan_event_handler = XRayCentreEventHandler()
+
+    @subs_decorator(flyscan_event_handler)
+    def plan_to_perform():
+        yield from ispyb_activation_wrapper(
+            detect_grid_and_do_gridscan(
+                composite,
+                parameters,
+                oav_params,
+                xrc_params_type,
+                construct_beamline_specific,
+            ),
+            parameters,
+        )
+
+    yield from start_preparing_data_collection_then_do_plan(
+        composite.beamstop,
+        eiger,
+        composite.detector_motion,
+        parameters.detector_params.detector_distance,
+        plan_to_perform(),
+        group=PlanGroupCheckpointConstants.GRID_READY_FOR_DC,
+    )
+
+    assert flyscan_event_handler.xray_centre_results, (
+        "Flyscan result event not received or no crystal found and exception not raised"
+    )
+
+    yield from change_aperture_then_move_to_xtal(
+        flyscan_event_handler.xray_centre_results[0],
+        composite.smargon,
+        composite.aperture_scatterguard,
+    )
 
 
 def detect_grid_and_do_gridscan(
     composite: GridDetectThenXRayCentreComposite,
     parameters: GridCommon,
     oav_params: OAVParameters,
+    xrc_params_type: type[SpecifiedThreeDGridScan],
+    construct_beamline_specific: ConstructBeamlineSpecificFeatures,
 ):
     snapshot_template = f"{parameters.detector_params.prefix}_{parameters.detector_params.run_number}_{{angle}}"
 
     grid_params_callback = GridDetectionCallback()
 
+    yield from setup_beamline_for_OAV(
+        composite.smargon,
+        composite.backlight,
+        composite.aperture_scatterguard,
+        wait=True,
+    )
+
     @bpp.subs_decorator([grid_params_callback])
     def run_grid_detection_plan(
         oav_params,
@@ -109,19 +151,12 @@ def detect_grid_and_do_gridscan(
             parameters.box_size_um,
         )
 
-    yield from setup_beamline_for_OAV(
-        composite.smargon,
-        composite.backlight,
-        composite.aperture_scatterguard,
-        wait=True,
-    )
-
     if parameters.selected_aperture:
         # Start moving the aperture/scatterguard into position without moving it in
         yield from bps.prepare(
             composite.aperture_scatterguard,
             parameters.selected_aperture,
-            group=CONST.WAIT.PREPARE_APERTURE,
+            group=PlanGroupCheckpointConstants.PREPARE_APERTURE,
         )
 
     yield from run_grid_detection_plan(
@@ -131,105 +166,42 @@ def detect_grid_and_do_gridscan(
     )
 
     yield from bps.abs_set(
-        composite.backlight, BacklightPosition.OUT, group=CONST.WAIT.GRID_READY_FOR_DC
+        composite.backlight,
+        BacklightPosition.OUT,
+        group=PlanGroupCheckpointConstants.GRID_READY_FOR_DC,
     )
 
-    yield from bps.wait(CONST.WAIT.PREPARE_APERTURE)
+    yield from bps.wait(PlanGroupCheckpointConstants.PREPARE_APERTURE)
     yield from move_aperture_if_required(
         composite.aperture_scatterguard,
         parameters.selected_aperture,
-        group=CONST.WAIT.GRID_READY_FOR_DC,
+        group=PlanGroupCheckpointConstants.GRID_READY_FOR_DC,
     )
-
-    xrc_composite = HyperionFlyScanXRayCentreComposite(
-        aperture_scatterguard=composite.aperture_scatterguard,
-        attenuator=composite.attenuator,
-        backlight=composite.backlight,
-        eiger=composite.eiger,
-        panda_fast_grid_scan=composite.panda_fast_grid_scan,
-        flux=composite.flux,
-        s4_slit_gaps=composite.s4_slit_gaps,
-        smargon=composite.smargon,
-        undulator=composite.undulator,
-        synchrotron=composite.synchrotron,
-        xbpm_feedback=composite.xbpm_feedback,
-        zebra=composite.zebra,
-        zocalo=composite.zocalo,
-        panda=composite.panda,
-        zebra_fast_grid_scan=composite.zebra_fast_grid_scan,
-        dcm=composite.dcm,
-        robot=composite.robot,
-        sample_shutter=composite.sample_shutter,
+    xrc_params = create_parameters_for_flyscan_xray_centre(
+        parameters, grid_params_callback.get_grid_parameters(), xrc_params_type
     )
+    beamline_specific = construct_beamline_specific(composite, xrc_params)
 
-    params = create_parameters_for_flyscan_xray_centre(
-        parameters, grid_params_callback.get_grid_parameters()
-    )
+    yield from common_flyscan_xray_centre(composite, xrc_params, beamline_specific)
 
-    beamline_specific = construct_hyperion_specific_features(xrc_composite, params)
 
-    yield from _gridscan_with_undulator_checks(xrc_composite, params, beamline_specific)
-
-
-def _gridscan_with_undulator_checks(
-    composite: HyperionFlyScanXRayCentreComposite,
-    params: HyperionSpecifiedThreeDGridScan,
-    beamline_specific: BeamlineSpecificFGSFeatures,
+class ConstructBeamlineSpecificFeatures(
+    Protocol[TFlyScanEssentialDevices, TSpecifiedThreeDGridScan]
 ):
-    @transmission_and_xbpm_feedback_for_collection_decorator(
-        composite, params.transmission_frac
-    )
-    @verify_undulator_gap_before_run_decorator(composite)
-    def _inner():
-        yield from common_flyscan_xray_centre(composite, params, beamline_specific)
-
-    yield from _inner()
+    def __call__(
+        self,
+        xrc_composite: TFlyScanEssentialDevices,
+        xrc_parameters: TSpecifiedThreeDGridScan,
+    ) -> BeamlineSpecificFGSFeatures: ...
 
 
-def grid_detect_then_xray_centre(
-    composite: GridDetectThenXRayCentreComposite,
-    parameters: GridScanWithEdgeDetect,
-    oav_config: str = OavConstants.OAV_CONFIG_JSON,
-) -> MsgGenerator:
-    """
-    A plan which combines the collection of snapshots from the OAV and the determination
-    of the grid dimensions to use for the following grid scan.
-    """
-
-    eiger: EigerDetector = composite.eiger
-
-    eiger.set_detector_parameters(parameters.detector_params)
-
-    oav_params = OAVParameters("xrayCentring", oav_config)
-
-    flyscan_event_handler = XRayCentreEventHandler()
-
-    @subs_decorator(flyscan_event_handler)
-    def plan_to_perform():
-        yield from ispyb_activation_wrapper(
-            detect_grid_and_do_gridscan(
-                composite,
-                parameters,
-                oav_params,
-            ),
-            parameters,
-        )
-
-    yield from start_preparing_data_collection_then_do_plan(
-        composite.beamstop,
-        eiger,
-        composite.detector_motion,
-        parameters.detector_params.detector_distance,
-        plan_to_perform(),
-        group=CONST.WAIT.GRID_READY_FOR_DC,
-    )
-
-    assert flyscan_event_handler.xray_centre_results, (
-        "Flyscan result event not received or no crystal found and exception not raised"
-    )
-
-    yield from change_aperture_then_move_to_xtal(
-        flyscan_event_handler.xray_centre_results[0],
-        composite.smargon,
-        composite.aperture_scatterguard,
-    )
+def create_parameters_for_flyscan_xray_centre(
+    parameters: GridCommon,
+    grid_parameters: GridParamUpdate,
+    xrc_params_type: type[SpecifiedThreeDGridScan],
+) -> SpecifiedThreeDGridScan:
+    params_json = parameters.model_dump()
+    params_json.update(grid_parameters)
+    flyscan_xray_centre_parameters = xrc_params_type(**params_json)
+    LOGGER.info(f"Parameters for FGS: {flyscan_xray_centre_parameters}")
+    return flyscan_xray_centre_parameters

mx_bluesky/common/{plans → experiment_plans}/inner_plans/do_fgs.py

@@ -14,10 +14,10 @@ from dodal.log import LOGGER
 from dodal.plan_stubs.check_topup import check_topup_and_wait_if_necessary
 from scanspec.core import AxesPoints, Axis
 
+from mx_bluesky.common.experiment_plans.read_hardware import read_hardware_for_zocalo
 from mx_bluesky.common.parameters.constants import (
     PlanNameConstants,
 )
-from mx_bluesky.common.plans.read_hardware import read_hardware_for_zocalo
 from mx_bluesky.common.utils.tracing import TRACER
 
 

mx_bluesky/common/experiment_plans/oav_grid_detection_plan.py

@@ -5,9 +5,7 @@ from typing import TYPE_CHECKING
 
 import bluesky.plan_stubs as bps
 import numpy as np
-import pydantic
 from blueapi.core import BlueskyContext
-from dodal.devices.backlight import Backlight
 from dodal.devices.oav.oav_detector import OAV
 from dodal.devices.oav.pin_image_recognition import PinTipDetection
 from dodal.devices.oav.pin_image_recognition.utils import NONE_VALUE
@@ -17,7 +15,11 @@ from dodal.devices.smargon import Smargon
 from mx_bluesky.common.device_setup_plans.setup_oav import (
     pre_centring_setup_oav,
 )
-from mx_bluesky.common.parameters.constants import DocDescriptorNames, HardwareConstants
+from mx_bluesky.common.parameters.constants import (
+    DocDescriptorNames,
+    HardwareConstants,
+)
+from mx_bluesky.common.parameters.device_composites import OavGridDetectionComposite
 from mx_bluesky.common.utils.context import device_composite_from_context
 from mx_bluesky.common.utils.exceptions import catch_exception_and_warn
 from mx_bluesky.common.utils.log import LOGGER
@@ -26,16 +28,6 @@ if TYPE_CHECKING:
     from dodal.devices.oav.oav_parameters import OAVParameters
 
 
-@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
-class OavGridDetectionComposite:
-    """All devices which are directly or indirectly required by this plan"""
-
-    backlight: Backlight
-    oav: OAV
-    smargon: Smargon
-    pin_tip_detection: PinTipDetection
-
-
 def create_devices(context: BlueskyContext) -> OavGridDetectionComposite:
     return device_composite_from_context(context, OavGridDetectionComposite)
 

mx_bluesky/{hyperion → common}/experiment_plans/oav_snapshot_plan.py

@@ -11,7 +11,10 @@ from dodal.devices.smargon import Smargon
 
 from mx_bluesky.common.device_setup_plans.setup_oav import setup_general_oav_params
 from mx_bluesky.common.parameters.components import WithSnapshot
-from mx_bluesky.hyperion.parameters.constants import CONST, DocDescriptorNames
+from mx_bluesky.common.parameters.constants import (
+    DocDescriptorNames,
+    PlanGroupCheckpointConstants,
+)
 
 OAV_SNAPSHOT_SETUP_SHOT = "oav_snapshot_setup_shot"
 OAV_SNAPSHOT_GROUP = "oav_snapshot_group"
@@ -28,7 +31,7 @@ def setup_beamline_for_OAV(
     smargon: Smargon,
     backlight: Backlight,
     aperture_scatterguard: ApertureScatterguard,
-    group=CONST.WAIT.READY_FOR_OAV,
+    group=PlanGroupCheckpointConstants.READY_FOR_OAV,
     wait=False,
 ):
     max_vel = yield from bps.rd(smargon.omega.max_velocity)

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

@@ -136,7 +136,7 @@ class WithSnapshot(BaseModel):
 
     @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, (
+        assert not self.use_grid_snapshots or not self.snapshot_omegas_deg, (
             "snapshot_omegas may not be specified with use_grid_snapshots"
         )
         return self
@@ -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"