Use protocols instead of smargon type for FGS composite, rename smargon to gonio

pyproject.toml

@@ -4,9 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [tool.uv]
 # Restrict lockfile to a sane subset of platforms
-environments = [
-    "sys_platform == 'linux' and platform_machine == 'x86_64'"
-]
+environments = ["sys_platform == 'linux' and platform_machine == 'x86_64'"]
 
 [project]
 name = "mx-bluesky"
@@ -31,7 +29,7 @@ dependencies = [
     "matplotlib",
     "nexgen >= 0.11.0",
     "numpy >= 2.3.5",
-    "opencv-python",                  # Needed for I24 ssx moveonclick. To be changed to headless once this is moved to separate ui.
+    "opencv-python",                   # Needed for I24 ssx moveonclick. To be changed to headless once this is moved to separate ui.
     "opentelemetry-distro",
     "opentelemetry-exporter-otlp",
     "pydantic",
@@ -46,7 +44,7 @@ dependencies = [
     "matplotlib",
     "cachetools",
     "bluesky-stomp >= 0.2.0",
-    "mysql-connector-python == 9.5.0",  # Can unpin once https://github.com/DiamondLightSource/ispyb-api/pull/244 is merged and released
+    "mysql-connector-python == 9.5.0", # Can unpin once https://github.com/DiamondLightSource/ispyb-api/pull/244 is merged and released
 
     #
     # These dependencies may be issued as pre-release versions and should have a pin constraint
@@ -57,7 +55,7 @@ dependencies = [
     "ophyd >= 1.10.5",
     "ophyd-async >= 0.14.0",
     "bluesky >= 1.14.6",
-    "dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git@main",
+    "dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git@28c3f20f6108d2769e7995f1008e7d76446fb38b",
 ]
 
 
@@ -269,10 +267,10 @@ env_dir = ".venv"
 skip_install = true
 commands_pre = [
     [
-        "python", 
-        "{toxinidir}/utility_scripts/generate_plantuml.py", 
-        "docs/developer/hyperion/reference/param_hierarchy.puml"
-    ]
+        "python",
+        "{toxinidir}/utility_scripts/generate_plantuml.py",
+        "docs/developer/hyperion/reference/param_hierarchy.puml",
+    ],
 ]
 commands = [
     [

src/mx_bluesky/beamlines/i04/callbacks/murko_callback.py

@@ -95,15 +95,15 @@ def event(self, doc: Event) -> Event:
                     }
                 )
 
-        if (latest_omega := data.get("smargon-omega")) is not None:
+        if (latest_omega := data.get("gonio-omega")) is not None:
             if len(self.previous_omegas) <= 2 and self.last_uuid:
                 # For the first few images there's not enough data to extrapolate so we
                 # match them one to one
                 self.call_murko(self.last_uuid, latest_omega)
             self.previous_omegas.append(
                 OmegaReading(
                     value=latest_omega,
-                    timestamp=doc["timestamps"]["smargon-omega"],
+                    timestamp=doc["timestamps"]["gonio-omega"],
                 )
             )
         elif (uuid := doc["data"].get("oav_to_redis_forwarder-uuid")) is not None:

src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py

@@ -140,7 +140,7 @@ def i04_default_grid_detect_and_xray_centre(
     eiger: EigerDetector = inject("eiger"),
     synchrotron: Synchrotron = inject("synchrotron"),
     zocalo: ZocaloResults = inject("zocalo"),
-    smargon: Smargon = inject("smargon"),
+    smargon: Smargon = inject("gonio"),
     detector_motion: DetectorMotion = inject("detector_motion"),
     transfocator: Transfocator = inject("transfocator"),
     oav_config: str = OavConstants.OAV_CONFIG_JSON,
@@ -195,7 +195,7 @@ def tidy_beamline():
 
         if not udc:
             yield from get_ready_for_oav_and_close_shutter(
-                composite.smargon,
+                composite.gonio,
                 composite.backlight,
                 composite.aperture_scatterguard,
                 composite.detector_motion,
@@ -296,7 +296,7 @@ def construct_i04_specific_features(
         xrc_composite.undulator.current_gap,
         xrc_composite.synchrotron.synchrotron_mode,
         xrc_composite.s4_slit_gaps,
-        xrc_composite.smargon,
+        xrc_composite.gonio,
         xrc_composite.dcm.energy_in_keV,
     ]
 

src/mx_bluesky/beamlines/i04/thawing_plan.py

@@ -18,7 +18,7 @@ def thaw(
     time_to_thaw: float,
     rotation: float = 360,
     thawer: Thawer = inject("thawer"),
-    smargon: Smargon = inject("smargon"),
+    smargon: Smargon = inject("gonio"),
 ) -> MsgGenerator:
     """Turns on the thawer and rotates the sample by {rotation} degrees to thaw it, then
     rotates {rotation} degrees back and turns the thawer off. The speed of the goniometer
@@ -54,7 +54,7 @@ def thaw_and_murko_centre(
     rotation: float = 360,
     robot: BartRobot = inject("robot"),
     thawer: Thawer = inject("thawer"),
-    smargon: Smargon = inject("smargon"),
+    smargon: Smargon = inject("gonio"),
     murko_results: MurkoResultsDevice = inject("murko_results"),
     oav_to_redis_forwarder: OAVToRedisForwarder = inject("oav_to_redis_forwarder"),
 ) -> MsgGenerator:
@@ -157,7 +157,7 @@ def thaw_and_stream_to_redis(
     rotation: float = 360,
     robot: BartRobot = inject("robot"),
     thawer: Thawer = inject("thawer"),
-    smargon: Smargon = inject("smargon"),
+    smargon: Smargon = inject("gonio"),
     oav_to_redis_forwarder: OAVToRedisForwarder = inject("oav_to_redis_forwarder"),
 ) -> MsgGenerator:
     """Turns on the thawer and rotates the sample by {rotation} degrees to thaw it, then
@@ -243,7 +243,7 @@ def get_metadata_from_current_oav():
     )
 
     yield from get_metadata_from_current_oav()
-    yield from bps.monitor(smargon.omega.user_readback, name="smargon")
+    yield from bps.monitor(smargon.omega.user_readback, name="gonio")
     yield from bps.monitor(oav_to_redis_forwarder.uuid, name="oav")
 
     yield from bps.kickoff(oav_to_redis_forwarder, wait=True)

src/mx_bluesky/common/experiment_plans/change_aperture_then_move_plan.py

@@ -28,7 +28,9 @@ def get_results_then_change_aperture_and_move_to_xtal(
         "Flyscan result event not received or no crystal found and exception not raised"
     )
     yield from change_aperture_then_move_to_xtal(
-        flyscan_results[0], composite.smargon, composite.aperture_scatterguard
+        flyscan_results[0],
+        composite.gonio,
+        composite.aperture_scatterguard,
     )
 
 

src/mx_bluesky/common/experiment_plans/common_flyscan_xray_centre_plan.py

@@ -26,7 +26,10 @@
     PlanGroupCheckpointConstants,
     PlanNameConstants,
 )
-from mx_bluesky.common.parameters.device_composites import FlyScanEssentialDevices
+from mx_bluesky.common.parameters.device_composites import (
+    FlyScanEssentialDevices,
+    GonioWithOmegaType,
+)
 from mx_bluesky.common.parameters.gridscan import SpecifiedThreeDGridScan
 from mx_bluesky.common.utils.exceptions import (
     SampleError,
@@ -110,7 +113,7 @@ def construct_beamline_specific_fast_gridscan_features(
 
 
 def common_flyscan_xray_centre(
-    composite: FlyScanEssentialDevices,
+    composite: FlyScanEssentialDevices[GonioWithOmegaType],
     parameters: SpecifiedThreeDGridScan,
     beamline_specific: BeamlineSpecificFGSFeatures,
 ) -> MsgGenerator:
@@ -152,7 +155,7 @@ def _decorated_flyscan():
         )
         @bpp.finalize_decorator(lambda: _overall_tidy())
         def run_gridscan_and_tidy(
-            fgs_composite: FlyScanEssentialDevices,
+            fgs_composite: FlyScanEssentialDevices[GonioWithOmegaType],
             params: SpecifiedThreeDGridScan,
             beamline_specific: BeamlineSpecificFGSFeatures,
         ) -> MsgGenerator:
@@ -170,13 +173,13 @@ def run_gridscan_and_tidy(
 
 
 def run_gridscan(
-    fgs_composite: FlyScanEssentialDevices,
+    fgs_composite: FlyScanEssentialDevices[GonioWithOmegaType],
     parameters: SpecifiedThreeDGridScan,
     beamline_specific: BeamlineSpecificFGSFeatures,
 ):
     # Currently gridscan only works for omega 0, see https://github.com/DiamondLightSource/mx-bluesky/issues/410
     with TRACER.start_span("moving_omega_to_0"):
-        yield from bps.abs_set(fgs_composite.smargon.omega, 0)
+        yield from bps.abs_set(fgs_composite.gonio.omega, 0)
 
     with TRACER.start_span("ispyb_hardware_readings"):
         yield from beamline_specific.read_pre_flyscan_plan()

src/mx_bluesky/common/experiment_plans/common_grid_detect_then_xray_centre_plan.py

@@ -106,7 +106,7 @@ def detect_grid_and_do_gridscan(
     grid_params_callback = GridDetectionCallback()
 
     yield from setup_beamline_for_oav(
-        composite.smargon,
+        composite.gonio,
         composite.backlight,
         composite.aperture_scatterguard,
         wait=True,
@@ -121,7 +121,7 @@ def run_grid_detection_plan(
         grid_detect_composite = OavGridDetectionComposite(
             backlight=composite.backlight,
             oav=composite.oav,
-            smargon=composite.smargon,
+            gonio=composite.gonio,
             pin_tip_detection=composite.pin_tip_detection,
         )
 

src/mx_bluesky/common/experiment_plans/oav_grid_detection_plan.py

@@ -81,7 +81,7 @@ def grid_detection_plan(
         box_size_um (float): The size of each box of the grid in microns
     """
     oav: OAV = composite.oav
-    smargon: Smargon = composite.smargon
+    smargon: Smargon = composite.gonio
     pin_tip_detection: PinTipDetection = composite.pin_tip_detection
 
     LOGGER.info("OAV Centring: Starting grid detection centring")

src/mx_bluesky/common/experiment_plans/oav_snapshot_plan.py

@@ -21,7 +21,7 @@
 
 
 class OavSnapshotComposite(Protocol):
-    smargon: Smargon
+    gonio: Smargon
     oav: OAV
     aperture_scatterguard: ApertureScatterguard
 
@@ -78,15 +78,13 @@ def _generate_oav_snapshots(composite: OavSnapshotComposite, params: WithSnapsho
     for _ in 0, 270:
         yield from bps.create(DocDescriptorNames.OAV_ROTATION_SNAPSHOT_TRIGGERED)
         yield from bps.read(composite.oav)
-        yield from bps.read(composite.smargon)
+        yield from bps.read(composite.gonio)
         yield from bps.save()
 
 
 def _take_oav_snapshot(composite: OavSnapshotComposite, omega: float):
     """Create new snapshots by triggering the OAV"""
-    yield from bps.abs_set(
-        composite.smargon.omega, omega, group=OAV_SNAPSHOT_SETUP_SHOT
-    )
+    yield from bps.abs_set(composite.gonio.omega, omega, group=OAV_SNAPSHOT_SETUP_SHOT)
     filename = _snapshot_filename(omega)
     yield from bps.abs_set(
         composite.oav.snapshot.filename,

src/mx_bluesky/common/external_interaction/callbacks/common/grid_detection_callback.py

@@ -68,10 +68,8 @@ def event(self, doc: Event):
         top_left_y_px = data["oav-grid_snapshot-top_left_y"]
         y_of_centre_of_first_box_px = top_left_y_px + box_width_px / 2
 
-        smargon_omega = data["smargon-omega"]
-        current_xyz = np.array(
-            [data["smargon-x"], data["smargon-y"], data["smargon-z"]]
-        )
+        gonio_omega = data["gonio-omega"]
+        current_xyz = np.array([data["gonio-x"], data["gonio-y"], data["gonio-z"]])
 
         centre_of_first_box = (
             x_of_centre_of_first_box_px,
@@ -89,7 +87,7 @@ def event(self, doc: Event):
 
         position_grid_start_mm = calculate_x_y_z_of_pixel(
             current_xyz,
-            smargon_omega,
+            gonio_omega,
             centre_of_first_box,
             (beam_x, beam_y),
             (microns_per_pixel_x, microns_per_pixel_y),
@@ -98,19 +96,19 @@ def event(self, doc: Event):
         LOGGER.info(f"Calculated start position {position_grid_start_mm}")
 
         # If data is taken at omega=~0 then it gives us x-y info, at omega=~-90 it is x-z
-        if abs(smargon_omega) < self.OMEGA_TOLERANCE:
+        if abs(gonio_omega) < self.OMEGA_TOLERANCE:
             self.start_positions_um["x"] = position_grid_start_mm[0] * 1000
             self.start_positions_um["y"] = position_grid_start_mm[1] * 1000
             self.box_numbers["x"] = data["oav-grid_snapshot-num_boxes_x"]
             self.box_numbers["y"] = data["oav-grid_snapshot-num_boxes_y"]
-        elif abs(smargon_omega + 90) < self.OMEGA_TOLERANCE:
+        elif abs(gonio_omega + 90) < self.OMEGA_TOLERANCE:
             self.start_positions_um["x"] = position_grid_start_mm[0] * 1000
             self.start_positions_um["z"] = position_grid_start_mm[2] * 1000
             self.box_numbers["x"] = data["oav-grid_snapshot-num_boxes_x"]
             self.box_numbers["z"] = data["oav-grid_snapshot-num_boxes_y"]
         else:
             raise ValueError(
-                f"Grid detection only works at omegas of 0 or -90, omega of {smargon_omega} given."
+                f"Grid detection only works at omegas of 0 or -90, omega of {gonio_omega} given."
             )
 
         self.x_step_size_um = box_width_px * microns_per_pixel_x

src/mx_bluesky/common/external_interaction/callbacks/common/ispyb_callback_base.py

@@ -145,9 +145,9 @@ def _handle_ispyb_hardware_read(self, doc) -> Sequence[ScanDataInfo]:
             doc, self.params.detector_params, hwscan_data_collection_info
         )
         hwscan_position_info = DataCollectionPositionInfo(
-            pos_x=float(doc["data"]["smargon-x"]),
-            pos_y=float(doc["data"]["smargon-y"]),
-            pos_z=float(doc["data"]["smargon-z"]),
+            pos_x=float(doc["data"]["gonio-x"]),
+            pos_y=float(doc["data"]["gonio-y"]),
+            pos_z=float(doc["data"]["gonio-z"]),
         )
         scan_data_infos = self.populate_info_for_update(
             hwscan_data_collection_info, hwscan_position_info, self.params

src/mx_bluesky/common/external_interaction/callbacks/xray_centre/ispyb_callback.py

@@ -187,7 +187,7 @@ def _handle_oav_grid_snapshot_triggered(self, doc) -> Sequence[ScanDataInfo]:
         assert self.params, "ISPyB handler didn't receive parameters!"
         assert self.data_collection_group_info, "No data collection group"
         data = doc["data"]
-        omega = doc["data"]["smargon-omega"]
+        omega = doc["data"]["gonio-omega"]
         grid_plane = _smargon_omega_to_xyxz_plane(omega)
         ISPYB_ZOCALO_CALLBACK_LOGGER.info(
             f"Generating dc info for gridplane {grid_plane}, omega {omega}"
@@ -255,13 +255,13 @@ def _handle_oav_grid_snapshot_triggered(self, doc) -> Sequence[ScanDataInfo]:
         return [scan_data_info]
 
     def _populate_axis_info(self, data_collection_info: DataCollectionInfo, doc: dict):
-        if (omega_start := doc.get("smargon-omega")) is not None:
+        if (omega_start := doc.get("gonio-omega")) is not None:
             omega_in_gda_space = -omega_start
             data_collection_info.omega_start = omega_in_gda_space
             data_collection_info.axis_start = omega_in_gda_space
             data_collection_info.axis_end = omega_in_gda_space
             data_collection_info.axis_range = 0
-        if (chi_start := doc.get("smargon-chi")) is not None:
+        if (chi_start := doc.get("gonio-chi")) is not None:
             data_collection_info.chi_start = chi_start
 
     def populate_info_for_update(

src/mx_bluesky/common/parameters/device_composites.py

@@ -1,3 +1,5 @@
+from typing import Generic, Protocol, TypeVar, runtime_checkable
+
 import pydantic
 from dodal.devices.aperturescatterguard import (
     ApertureScatterguard,
@@ -24,13 +26,23 @@
 from dodal.devices.zebra.zebra import Zebra
 from dodal.devices.zebra.zebra_controlled_shutter import ZebraShutter
 from dodal.devices.zocalo import ZocaloResults
+from ophyd_async.epics.motor import Motor
+
+
+# MX gridscans only uses the gonio to set omega to 0. Other motors are only accessed in the motion program
+@runtime_checkable
+class GonioWithOmega(Protocol):
+    omega: Motor
+
+
+GonioWithOmegaType = TypeVar("GonioWithOmegaType", bound=GonioWithOmega)
 
 
 @pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
-class FlyScanEssentialDevices:
+class FlyScanEssentialDevices(Generic[GonioWithOmegaType]):
     eiger: EigerDetector
     synchrotron: Synchrotron
-    smargon: Smargon
+    gonio: GonioWithOmegaType
 
 
 @pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
@@ -39,12 +51,12 @@ class OavGridDetectionComposite:
 
     backlight: Backlight
     oav: OAV
-    smargon: Smargon
+    gonio: Smargon
     pin_tip_detection: PinTipDetection
 
 
 @pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
-class GridDetectThenXRayCentreComposite(FlyScanEssentialDevices):
+class GridDetectThenXRayCentreComposite(FlyScanEssentialDevices[Smargon]):
     """All devices which are directly or indirectly required by this plan"""
 
     aperture_scatterguard: ApertureScatterguard

src/mx_bluesky/hyperion/blueapi_plans/__init__.py

@@ -60,7 +60,7 @@ def load_centre_collect(
 def robot_unload(
     visit: str,
     robot: BartRobot = inject("robot"),
-    smargon: Smargon = inject("smargon"),
+    smargon: Smargon = inject("gonio"),
     aperture_scatterguard: ApertureScatterguard = inject("aperture_scatterguard"),
     lower_gonio: XYZStage = inject("lower_gonio"),
 ) -> MsgGenerator:
@@ -75,7 +75,7 @@ def clean_up_udc(
     visit: str,
     cleanup_group: str = "cleanup",
     robot: BartRobot = inject("robot"),
-    smargon: Smargon = inject("smargon"),
+    smargon: Smargon = inject("gonio"),
     aperture_scatterguard: ApertureScatterguard = inject("aperture_scatterguard"),
     lower_gonio: XYZStage = inject("lower_gonio"),
     detector_motion: DetectorMotion = inject("detector_motion"),