349 Create OAV snapshot plan (#1437)

* (#349) Basic experiment plan for taking snapshots
* Remove xtal_snapshots and xtal_snapshots_omega_end from ispyb_extras
* (#349) Tidy parameter model diagram and commit image
* (#349) tidyups

.gitignore

@@ -8,6 +8,7 @@ __pycache__/
 
 # Output
 *.png
+!docs/*.png
 
 # Distribution / packaging
 .Python

.pre-commit-config.yaml

@@ -14,11 +14,10 @@ repos:
       - id: no-images
         name: Check for image files
         entry: >
-          Images for documentation should go into the documentation repository
-          https://github.com/dials/dials.github.io
+          Images for documentation should go into the docs folder
         language: fail
         files: '.*\.png$'
-
+        exclude: '^docs/'
       - id: ruff
         name: Run ruff
         stages: [commit]

docs/param_hierarchy.puml

@@ -0,0 +1,66 @@
+@startuml
+'https://plantuml.com/class-diagram
+title Hyperion Parameter Model
+
+abstract class BaseModel
+
+package Mixins {
+    class WithSample
+    class WithScan
+    class WithOavCentring
+    class WithSnapshot
+    class OptionalXyzStarts
+    class XyzStarts
+    class OptionalGonioAngleStarts
+    class SplitScan
+}
+
+class HyperionParameters
+note bottom: Base class for all experiment parameter models
+
+package Experiments {
+    class DiffractionExperiment
+    class DiffractionExperimentWithSample
+    class GridCommon
+    class GridScanWithEdgeDetect
+    class PinTipCentreThenXrayCentre
+    class RotationScan
+    class RobotLoadThenCentre
+    class SpecifiedGridScan
+    class ThreeDGridScan
+}
+class TemporaryIspybExtras
+note bottom: To be removed
+
+
+BaseModel <|-- HyperionParameters
+BaseModel <|-- SplitScan
+BaseModel <|-- OptionalGonioAngleStarts
+BaseModel <|-- OptionalXyzStarts
+BaseModel <|-- TemporaryIspybExtras
+BaseModel <|-- WithOavCentring
+BaseModel <|-- WithSnapshot
+BaseModel <|-- WithSample
+BaseModel <|-- WithScan
+BaseModel <|-- XyzStarts
+
+RotationScan *-- TemporaryIspybExtras
+HyperionParameters <|-- DiffractionExperiment
+WithSnapshot <|-- DiffractionExperiment
+DiffractionExperiment <|-- DiffractionExperimentWithSample
+WithSample <|-- DiffractionExperimentWithSample
+DiffractionExperimentWithSample <|-- GridCommon
+GridCommon <|-- GridScanWithEdgeDetect
+GridCommon <|-- PinTipCentreThenXrayCentre
+GridCommon <|-- RobotLoadThenCentre
+GridCommon <|-- SpecifiedGridScan
+WithScan <|-- SpecifiedGridScan
+SpecifiedGridScan <|-- ThreeDGridScan
+SplitScan <|-- ThreeDGridScan
+WithOavCentring <|-- GridCommon
+DiffractionExperimentWithSample <|-- RotationScan
+OptionalXyzStarts <|-- RotationScan
+XyzStarts <|-- SpecifiedGridScan
+OptionalGonioAngleStarts <|-- GridCommon
+OptionalGonioAngleStarts <|-- RotationScan
+@enduml

src/hyperion/device_setup_plans/setup_oav.py

@@ -54,21 +54,12 @@ def setup_pin_tip_detection_params(
     )
 
 
-def pre_centring_setup_oav(
-    oav: OAV,
-    parameters: OAVParameters,
-    pin_tip_detection_device: PinTipDetection,
-):
-    """
-    Setup OAV PVs with required values.
-    """
+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 setup_pin_tip_detection_params(pin_tip_detection_device, parameters)
-
     zoom_level_str = f"{float(parameters.zoom)}x"
     if zoom_level_str not in oav.zoom_controller.allowed_zoom_levels:
         raise OAVError_ZoomLevelNotFound(
@@ -81,6 +72,17 @@ def pre_centring_setup_oav(
         wait=True,
     )
 
+
+def pre_centring_setup_oav(
+    oav: OAV,
+    parameters: OAVParameters,
+    pin_tip_detection_device: PinTipDetection,
+):
+    """
+    Setup OAV PVs with required values.
+    """
+    yield from setup_general_oav_params(oav, parameters)
+    yield from setup_pin_tip_detection_params(pin_tip_detection_device, parameters)
     yield from bps.wait(oav_group)
 
     """

src/hyperion/experiment_plans/oav_grid_detection_plan.py

@@ -159,7 +159,7 @@ def grid_detection_plan(
         yield from bps.abs_set(oav.grid_snapshot.filename, snapshot_filename)
         yield from bps.abs_set(oav.grid_snapshot.directory, snapshot_dir)
         yield from bps.trigger(oav.grid_snapshot, wait=True)
-        yield from bps.create(CONST.DESCRIPTORS.OAV_SNAPSHOT_TRIGGERED)
+        yield from bps.create(CONST.DESCRIPTORS.OAV_GRID_SNAPSHOT_TRIGGERED)
 
         yield from bps.read(oav.grid_snapshot)
         yield from bps.read(smargon)

src/hyperion/experiment_plans/oav_snapshot_plan.py

@@ -0,0 +1,64 @@
+import dataclasses
+from datetime import datetime
+
+from blueapi.core import BlueskyContext, MsgGenerator
+from bluesky import plan_stubs as bps
+from dodal.devices.oav.oav_detector import OAV
+from dodal.devices.oav.oav_parameters import OAVParameters
+from dodal.devices.smargon import Smargon
+
+from hyperion.device_setup_plans.setup_oav import setup_general_oav_params
+from hyperion.parameters.components import WithSnapshot
+from hyperion.parameters.constants import DocDescriptorNames
+from hyperion.utils.context import device_composite_from_context
+
+OAV_SNAPSHOT_GROUP = "oav_snapshot_group"
+
+
+@dataclasses.dataclass
+class OavSnapshotComposite:
+    smargon: Smargon
+    oav: OAV
+
+
+def create_devices(context: BlueskyContext) -> OavSnapshotComposite:
+    return device_composite_from_context(context, OavSnapshotComposite)  # type: ignore
+
+
+def _setup_oav(
+    composite: OavSnapshotComposite,
+    parameters: WithSnapshot,
+    oav_parameters: OAVParameters,
+):
+    yield from setup_general_oav_params(composite.oav, oav_parameters)
+    yield from bps.abs_set(
+        composite.oav.snapshot.directory, str(parameters.snapshot_directory)
+    )
+
+
+def _take_oav_snapshot(
+    composite: OavSnapshotComposite, parameters: WithSnapshot, omega: float
+):
+    yield from bps.abs_set(composite.smargon.omega, omega, group=OAV_SNAPSHOT_GROUP)
+    time_now = datetime.now()
+    filename = f"{time_now.strftime('%H%M%S')}_oav_snapshot_{omega:.0f}"
+    yield from bps.abs_set(composite.oav.snapshot.filename, filename)
+    yield from bps.trigger(composite.oav.snapshot, group=OAV_SNAPSHOT_GROUP)
+    yield from bps.wait(group=OAV_SNAPSHOT_GROUP)
+    yield from bps.create(DocDescriptorNames.OAV_ROTATION_SNAPSHOT_TRIGGERED)
+    yield from bps.read(composite.oav.snapshot)
+    yield from bps.save()
+
+
+def oav_snapshot_plan(
+    composite: OavSnapshotComposite,
+    parameters: WithSnapshot,
+    oav_parameters: OAVParameters,
+    wait: bool = True,
+) -> MsgGenerator:
+    omegas = parameters.snapshot_omegas_deg
+    if not omegas:
+        return
+    yield from _setup_oav(composite, parameters, oav_parameters)
+    for omega in omegas:
+        yield from _take_oav_snapshot(composite, parameters, omega)

src/hyperion/external_interaction/callbacks/common/ispyb_mapping.py

@@ -18,7 +18,6 @@
     VISIT_PATH_REGEX,
     get_current_time_string,
 )
-from hyperion.log import ISPYB_LOGGER
 from hyperion.parameters.components import DiffractionExperimentWithSample
 
 
@@ -87,15 +86,3 @@ def get_visit_string(ispyb_params: IspybParams, detector_params: DetectorParams)
             f"Visit not found from {ispyb_params.visit_path} or {detector_params.directory}"
         )
     return visit_path_match
-
-
-def get_xtal_snapshots(ispyb_params):
-    if ispyb_params.xtal_snapshots_omega_start:
-        xtal_snapshots = ispyb_params.xtal_snapshots_omega_start[:3]
-        ISPYB_LOGGER.info(
-            f"Using rotation scan snapshots {xtal_snapshots} for ISPyB deposition"
-        )
-    else:
-        ISPYB_LOGGER.warning("No xtal snapshot paths sent to ISPyB!")
-        xtal_snapshots = []
-    return xtal_snapshots + [None] * (3 - len(xtal_snapshots))

src/hyperion/external_interaction/callbacks/ispyb_callback_base.py

@@ -93,8 +93,10 @@ def activity_gated_event(self, doc: Event) -> Event:
         match event_descriptor.get("name"):
             case CONST.DESCRIPTORS.ISPYB_HARDWARE_READ:
                 scan_data_infos = self._handle_ispyb_hardware_read(doc)
-            case CONST.DESCRIPTORS.OAV_SNAPSHOT_TRIGGERED:
-                scan_data_infos = self._handle_oav_snapshot_triggered(doc)
+            case CONST.DESCRIPTORS.OAV_ROTATION_SNAPSHOT_TRIGGERED:
+                scan_data_infos = self._handle_oav_rotation_snapshot_triggered(doc)
+            case CONST.DESCRIPTORS.OAV_GRID_SNAPSHOT_TRIGGERED:
+                scan_data_infos = self._handle_oav_grid_snapshot_triggered(doc)
             case CONST.DESCRIPTORS.ISPYB_TRANSMISSION_FLUX_READ:
                 scan_data_infos = self._handle_ispyb_transmission_flux_read(doc)
             case _:
@@ -135,7 +137,25 @@ def _handle_ispyb_hardware_read(self, doc) -> Sequence[ScanDataInfo]:
         ISPYB_LOGGER.info("Updating ispyb data collection after hardware read.")
         return scan_data_infos
 
-    def _handle_oav_snapshot_triggered(self, doc) -> Sequence[ScanDataInfo]:
+    def _handle_oav_rotation_snapshot_triggered(self, doc) -> Sequence[ScanDataInfo]:
+        assert self.ispyb_ids.data_collection_ids, "No current data collection"
+        assert self.params, "ISPyB handler didn't recieve parameters!"
+        data = doc["data"]
+        self._oav_snapshot_event_idx += 1
+        data_collection_info = DataCollectionInfo(
+            **{
+                f"xtal_snapshot{self._oav_snapshot_event_idx}": data.get(
+                    "oav_snapshot_last_saved_path"
+                )
+            }
+        )
+        scan_data_info = ScanDataInfo(
+            data_collection_id=self.ispyb_ids.data_collection_ids[-1],
+            data_collection_info=data_collection_info,
+        )
+        return [scan_data_info]
+
+    def _handle_oav_grid_snapshot_triggered(self, doc) -> Sequence[ScanDataInfo]:
         assert self.ispyb_ids.data_collection_ids, "No current data collection"
         assert self.params, "ISPyB handler didn't recieve parameters!"
         data = doc["data"]

src/hyperion/external_interaction/callbacks/rotation/ispyb_mapping.py

@@ -1,9 +1,7 @@
 from __future__ import annotations
 
-from hyperion.external_interaction.callbacks.common.ispyb_mapping import (
-    get_xtal_snapshots,
-)
 from hyperion.external_interaction.ispyb.data_model import DataCollectionInfo
+from hyperion.log import ISPYB_LOGGER
 from hyperion.parameters.rotation import RotationScan
 
 
@@ -16,7 +14,22 @@ def populate_data_collection_info_for_rotation(params: RotationScan):
         axis_end=(params.omega_start_deg + params.scan_width_deg),
         kappa_start=params.kappa_start_deg,
     )
-    (info.xtal_snapshot1, info.xtal_snapshot2, info.xtal_snapshot3) = (
-        get_xtal_snapshots(params.ispyb_params)
-    )
+    (
+        info.xtal_snapshot1,
+        info.xtal_snapshot2,
+        info.xtal_snapshot3,
+        info.xtal_snapshot4,
+    ) = get_xtal_snapshots(params.ispyb_params)
     return info
+
+
+def get_xtal_snapshots(ispyb_params):
+    if ispyb_params.xtal_snapshots_omega_start:
+        xtal_snapshots = ispyb_params.xtal_snapshots_omega_start[:4]
+        ISPYB_LOGGER.info(
+            f"Using rotation scan snapshots {xtal_snapshots} for ISPyB deposition"
+        )
+    else:
+        ISPYB_LOGGER.warning("No xtal snapshot paths sent to ISPyB!")
+        xtal_snapshots = []
+    return xtal_snapshots + [None] * (4 - len(xtal_snapshots))

src/hyperion/external_interaction/ispyb/data_model.py

@@ -19,6 +19,7 @@ class DataCollectionInfo:
     xtal_snapshot1: Optional[str] = None
     xtal_snapshot2: Optional[str] = None
     xtal_snapshot3: Optional[str] = None
+    xtal_snapshot4: Optional[str] = None
 
     n_images: Optional[int] = None
     axis_range: Optional[float] = None

src/hyperion/external_interaction/ispyb/ispyb_dataclass.py

@@ -8,8 +8,6 @@
     "sample_id": None,
     "visit_path": "",
     "position": None,
-    "xtal_snapshots_omega_start": ["test_1_y", "test_2_y", "test_3_y"],
-    "xtal_snapshots_omega_end": ["test_1_z", "test_2_z", "test_3_z"],
     "comment": "Descriptive comment.",
 }
 
@@ -22,7 +20,6 @@ class IspybParams(BaseModel):
 
     # Optional from GDA as populated by Ophyd
     xtal_snapshots_omega_start: Optional[list[str]] = None
-    xtal_snapshots_omega_end: Optional[list[str]] = None
     ispyb_experiment_type: Optional[str] = None
 
     class Config:

src/hyperion/parameters/components.py

@@ -138,7 +138,12 @@ def from_json(cls, input: str | None, *, allow_extras: bool = False):
         return params
 
 
-class DiffractionExperiment(HyperionParameters):
+class WithSnapshot(BaseModel):
+    snapshot_directory: Path
+    snapshot_omegas_deg: list[float] | None
+
+
+class DiffractionExperiment(HyperionParameters, WithSnapshot):
     """For all experiments which use beam"""
 
     visit: str = Field(min_length=1)
@@ -160,7 +165,6 @@ class DiffractionExperiment(HyperionParameters):
     selected_aperture: AperturePositionGDANames | None = Field(default=None)
     ispyb_experiment_type: IspybExperimentType
     storage_directory: str
-    snapshot_directory: Path
 
     @root_validator(pre=True)
     def validate_snapshot_directory(cls, values):
@@ -260,5 +264,3 @@ class Config:
         extra = Extra.forbid
 
     xtal_snapshots_omega_start: list[str] | None = None
-    xtal_snapshots_omega_end: list[str] | None = None
-    xtal_snapshots: list[str] | None = None

src/hyperion/parameters/constants.py

@@ -39,7 +39,8 @@ class DocDescriptorNames:
     # Robot load event descriptor
     ROBOT_LOAD = "robot_load"
     # For callbacks to use
-    OAV_SNAPSHOT_TRIGGERED = "snapshot_to_ispyb"
+    OAV_ROTATION_SNAPSHOT_TRIGGERED = "rotation_snapshot_triggered"
+    OAV_GRID_SNAPSHOT_TRIGGERED = "snapshot_to_ispyb"
     NEXUS_READ = "nexus_read_plan"
     ISPYB_HARDWARE_READ = "ispyb_reading_hardware"
     ISPYB_TRANSMISSION_FLUX_READ = "ispyb_update_transmission_flux"

src/hyperion/parameters/gridscan.py

@@ -23,7 +23,6 @@
     IspybExperimentType,
     OptionalGonioAngleStarts,
     SplitScan,
-    TemporaryIspybExtras,
     WithOavCentring,
     WithScan,
     XyzStarts,
@@ -50,18 +49,13 @@ class GridCommon(
     selected_aperture: AperturePositionGDANames | None = Field(
         default=AperturePositionGDANames.SMALL_APERTURE
     )
-    # field rather than inherited to make it easier to track when it can be removed:
-    ispyb_extras: TemporaryIspybExtras
 
     @property
     def ispyb_params(self):
         return GridscanIspybParams(
             visit_path=str(self.visit_directory),
             comment=self.comment,
             sample_id=self.sample_id,
-            xtal_snapshots_omega_start=self.ispyb_extras.xtal_snapshots_omega_start
-            or [],
-            xtal_snapshots_omega_end=self.ispyb_extras.xtal_snapshots_omega_end or [],
             ispyb_experiment_type=self.ispyb_experiment_type,
         )