Allow doing gridscans with generic number of grids

docs/developer/hyperion/reference/gridscan.puml

@@ -1,38 +1,70 @@
 @startuml
 title Gridscan Parameter Relationships
 
-class DiffractionExperiment
-class DiffractionExperimentWithSample
-class GridCommon {
-    grid_width_um
-    exposure_time_s
+package dodal {
+    class AbstractExperimentParameterBase
+    class AbstractExperimentWithBeamParams {
+        transmission_fraction
+    }
+    class GridScanParamsCommon {
+        x_steps
+        y_steps
+        x_step_size_mm
+        y_step_size_mm
+        x_start_mm
+        y1_start_mm
+        z1_start_mm
+        set_stub_offsets
+    }
+    class GridScanParamsThreeD
+    class PandAGridScanParams
+    class ZebraGridScanParamsThreeD
+    class ZebraGridScanParamsTwoD
 }
-class GridScanWithEdgeDetect {
-    box_size_um
-}
-class HyperionGridCommon {
-    enable_dev_shm
+
+rectangle "Internal Experiment Parameter Model" {
+    class MxBlueSkyParameters
+    class DiffractionExperiment
+    class DiffractionExperimentWithSample
+    class GenericGrid {
+        grid_width_um
+        exposure_time_s
+    }
+    class GenericGridWithHyperionDetectorParams
+    class GridScanWithEdgeDetect {
+        box_size_um
+    }
+    class PinTipCentreThenXrayCentre
+    class RobotLoadThenCentre
 }
-class HyperionThreeDGridScan {
-    x_step_size_um
-    y_step_size_um
-    z_step_size_um
-    y2_start_um
-    z2_start_um
-    --
-    grid_1_spec
-    grid_2_spec
-    scan_indices
-    scan_spec
-    scan_points
-    scan_points_first_grid
-    scan_points_second_grid
-    num_images
-    FGS_Params
-    panda_FGS_Params
+
+
+
+rectangle "Specified Grids" {
+    class SpecifiedGrids<GenericParamType> {
+        omega_starts_deg: float[]
+        x_step_size_um
+        y_step_sizes_um
+        x_steps
+        y_steps: int[]
+        --
+        fast_grid_scan_params: GenericParamType
+        grid_specs
+        scan_indices
+        scan_points
+        scan_spec
+        num_images
+    }
+    class "SpecifiedGrids<ZebraGridScanParamsThreeD>" as SpecifiedGridsZebraGridScanParamsThreeD
+    class "SpecifiedGrids<ZebraGridScanParamsTwoD>" as SpecifiedGridsZebraGridScanParamsTwoD
+    class SpecifiedTwoDGridScan
+    class SpecifiedThreeDGridScan
+    class HyperionSpecifiedThreeDGridScan {
+        --
+        panda_fast_gridscan_params: PandAGridScanParams
+    }
 }
-class MxBlueSkyParameters
-class SpecifiedGrid
+
 class XyzStarts {
     x_start_um
     y_start_um
@@ -43,67 +75,51 @@ class OptionalXYZStarts {
     y_start_um
     z_start_um
 }
-class RotationScanPerSweep
 
 MxBlueSkyParameters <|-- DiffractionExperiment
 DiffractionExperiment <|-- DiffractionExperimentWithSample
-DiffractionExperimentWithSample <|-- GridCommon
-GridCommon <|-- GridScanWithEdgeDetect
-GridCommon <|-- HyperionGridCommon
-HyperionGridCommon <|-- HyperionThreeDGridScan
-SpecifiedGrid <|-- HyperionThreeDGridScan
-XyzStarts <|-- SpecifiedGrid
+DiffractionExperimentWithSample <|-- GenericGrid
+GenericGrid <|-- GenericGridWithHyperionDetectorParams
+GenericGridWithHyperionDetectorParams <|-- GridScanWithEdgeDetect
+GenericGridWithHyperionDetectorParams <|-- PinTipCentreThenXrayCentre
+GenericGridWithHyperionDetectorParams <|-- RobotLoadThenCentre
+GenericGrid <|-- SpecifiedGrids
+SpecifiedGridsZebraGridScanParamsTwoD <|-- SpecifiedTwoDGridScan
+SpecifiedGridsZebraGridScanParamsThreeD <|-- SpecifiedThreeDGridScan
+SpecifiedThreeDGridScan <|-- HyperionSpecifiedThreeDGridScan
+XyzStarts <|-- SpecifiedGrids
 OptionalXYZStarts <|-- RotationScanPerSweep
 class GridParamUpdate {
-    x_start_um
-    y_start_um
-    y2_start_um
-    z_start_um
-    z2_start_um
-    x_steps
-    y_steps
-    z_steps
-    x_step_size_um
-    y_step_size_um
-    z_step_size_um
+    x_start_um: float
+    y_starts_um: float[]
+    z_starts_um: float[]
+    x_steps: float
+    y_steps: float[]
+    x_step_size_um: float
+    y_step_sizes_um: float[]
 }
 
 class GridDetectionCallback {
     get_grid_parameters() -> GridParamUpdate
 }
 
 GridDetectionCallback --> GridParamUpdate : generates from event. Adds 0.5 to get box-centres
-GridParamUpdate --> HyperionThreeDGridScan : combines with GridScanWithEdgeDetect
+GridParamUpdate --> HyperionSpecifiedThreeDGridScan : combines with GridScanWithEdgeDetect
 
 class experiment_plans {
     grid_detect_then_xray_centre()
     common_flyscan_xray_centre()
-    create_parameters_for_flyscan_xray_centre(GridScanWithEdgeDetect, GridParamUpdate) -> HyperionThreeDGridScan
-}
-
-class AbstractExperimentBase
-class AbstractExperimentWithBeamParams
-class GridScanParamsCommon {
-    x_steps
-    y_steps
-    z_steps
-    x_step_size_mm
-    y_step_size_mm
-    z_step_size_mm
-    x_start_mm
-    y1_start_mm
-    y2_start_mm
-    z1_start_mm
-    z2_start_mm
+    create_parameters_for_flyscan_xray_centre(GridScanWithEdgeDetect, GridParamUpdate) -> HyperionSpecifiedThreeDGridScan
 }
-class PandAGridScanParams
-class ZebraGridScanParamsThreeD
 
-AbstractExperimentBase <|-- AbstractExperimentWithBeamParams
+AbstractExperimentParameterBase <|-- AbstractExperimentWithBeamParams
 AbstractExperimentWithBeamParams <|-- GridScanParamsCommon
-GridScanParamsCommon <|-- PandAGridScanParams
-GridScanParamsCommon <|-- ZebraGridScanParamsThreeD
+GridScanParamsCommon <|-- ZebraGridScanParamsTwoD
+GridScanParamsCommon <|-- GridScanParamsThreeD
+GridScanParamsThreeD <|-- PandAGridScanParams
+GridScanParamsThreeD <|-- ZebraGridScanParamsThreeD
 
-HyperionThreeDGridScan --> ZebraGridScanParamsThreeD : generates
-HyperionThreeDGridScan --> PandAGridScanParams : generates
+HyperionSpecifiedThreeDGridScan --> ZebraGridScanParamsThreeD : generates
+HyperionSpecifiedThreeDGridScan --> PandAGridScanParams : generates
+SpecifiedTwoDGridScan --> ZebraGridScanParamsTwoD : generates
 @enduml

src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py

@@ -1,34 +1,27 @@
 from dodal.devices.beamlines.i02_1.fast_grid_scan import ZebraGridScanParamsTwoD
-from scanspec.specs import Product
 
 from mx_bluesky.common.parameters.components import SplitScan, WithOptionalEnergyChange
-from mx_bluesky.common.parameters.gridscan import SpecifiedGrid
+from mx_bluesky.common.parameters.gridscan import SpecifiedGrids
 
 
 class SpecifiedTwoDGridScan(
-    SpecifiedGrid[ZebraGridScanParamsTwoD],
+    SpecifiedGrids[ZebraGridScanParamsTwoD],
     SplitScan,
     WithOptionalEnergyChange,
 ):
     """Parameters representing a so-called 2D grid scan, which consists of doing a
     gridscan in X and Y."""
 
-    @property
-    def scan_spec(self) -> Product[str]:
-        """A fully specified ScanSpec object representing the grid, with x, y, z and
-        omega positions."""
-        return self.grid_1_spec
-
     @property
     def fast_gridscan_params(self) -> ZebraGridScanParamsTwoD:
         return ZebraGridScanParamsTwoD(
             x_steps=self.x_steps,
-            y_steps=self.y_steps,
+            y_steps=self.y_steps[0],
             x_step_size_mm=self.x_step_size_um / 1000,
-            y_step_size_mm=self.y_step_size_um / 1000,
+            y_step_size_mm=self.y_step_sizes_um[0] / 1000,
             x_start_mm=self.x_start_um / 1000,
-            y1_start_mm=self.y_start_um / 1000,
-            z1_start_mm=self.z_start_um / 1000,
+            y1_start_mm=self.y_starts_um[0] / 1000,
+            z1_start_mm=self.z_starts_um[0] / 1000,
             set_stub_offsets=self._set_stub_offsets,
             transmission_fraction=0.5,
             dwell_time_ms=self.exposure_time_s,

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

@@ -80,7 +80,7 @@
     GridDetectThenXRayCentreComposite,
 )
 from mx_bluesky.common.parameters.gridscan import (
-    GridCommon,
+    GenericGrid,
     SpecifiedThreeDGridScan,
 )
 from mx_bluesky.common.preprocessors.preprocessors import (
@@ -105,7 +105,7 @@ class I04AutoXrcParams(BaseModel):
 
 
 def _change_beamsize(
-    transfocator: Transfocator, beamsize: float, parameters: GridCommon
+    transfocator: Transfocator, beamsize: float, parameters: GenericGrid
 ):
     """i04 always uses the large aperture and changes beamsize with the transfocator.
 
@@ -188,7 +188,7 @@ def i04_default_grid_detect_and_xray_centre(
     initial_z = yield from bps.rd(smargon.z.user_readback)
 
     _current_wavelength_a = yield from bps.rd(composite.dcm.wavelength_in_a)
-    grid_common_params = _get_grid_common_params(_current_wavelength_a, parameters)
+    grid_common_params = _get_generic_grid_params(_current_wavelength_a, parameters)
 
     def tidy_beamline():
         yield from bps.mv(transfocator, initial_beamsize)
@@ -275,7 +275,7 @@ def create_gridscan_callbacks() -> tuple[
     return (
         GridscanNexusFileCallback(param_type=SpecifiedThreeDGridScan),
         GridscanISPyBCallback(
-            param_type=GridCommon,
+            param_type=GenericGrid,
             emit=ZocaloCallback(
                 PlanNameConstants.DO_FGS,
                 EnvironmentConstants.ZOCALO_ENV,
@@ -336,9 +336,9 @@ def construct_i04_specific_features(
     )
 
 
-def _get_grid_common_params(
+def _get_generic_grid_params(
     _current_wavelength_a: float, parameters: I04AutoXrcParams
-) -> GridCommon:
+) -> GenericGrid:
     """Calculate scaled transmission and exposure by comparing current beamline energy to default energy"""
     feature_settings = get_i04_feature_settings()
     _assumed_wavelength_a = feature_settings.ASSUMED_WAVELENGTH_IN_A
@@ -353,7 +353,7 @@ def _get_grid_common_params(
         )
     )
 
-    return GridCommon(
+    return GenericGrid(
         sample_id=parameters.sample_id,
         file_name=parameters.file_name,
         visit=parameters.visit,

src/mx_bluesky/common/experiment_plans/common_flyscan_xray_centre_plan.py

@@ -204,7 +204,8 @@ def run_gridscan(
         beamline_specific.fgs_motors,
         fgs_composite.eiger,
         fgs_composite.synchrotron,
-        [parameters.scan_points_first_grid, parameters.scan_points_second_grid],
+        parameters.scan_points,
+        parameters.omega_starts_deg,
         plan_during_collection=beamline_specific.read_during_collection_plan,
     )
 

src/mx_bluesky/common/experiment_plans/common_grid_detect_then_xray_centre_plan.py

@@ -43,7 +43,7 @@
     FlyScanEssentialDevices,
     GridDetectThenXRayCentreComposite,
 )
-from mx_bluesky.common.parameters.gridscan import GridCommon, SpecifiedThreeDGridScan
+from mx_bluesky.common.parameters.gridscan import GenericGrid, SpecifiedThreeDGridScan
 from mx_bluesky.common.utils.log import LOGGER
 
 TFlyScanEssentialDevices = TypeVar(
@@ -56,7 +56,7 @@
 
 def grid_detect_then_xray_centre(
     composite: GridDetectThenXRayCentreComposite,
-    parameters: GridCommon,
+    parameters: GenericGrid,
     xrc_params_type: type[SpecifiedThreeDGridScan],
     construct_beamline_specific: ConstructBeamlineSpecificFeatures,
     oav_config: str = OavConstants.OAV_CONFIG_JSON,
@@ -97,7 +97,7 @@ def plan_to_perform():
 # This function should be private but is currently called by Hyperion, see https://github.com/DiamondLightSource/mx-bluesky/issues/1148
 def detect_grid_and_do_gridscan(
     composite: GridDetectThenXRayCentreComposite,
-    parameters: GridCommon,
+    parameters: GenericGrid,
     oav_params: OAVParameters,
     xrc_params_type: type[SpecifiedThreeDGridScan],
     construct_beamline_specific: ConstructBeamlineSpecificFeatures,
@@ -181,7 +181,7 @@ def __call__(
 
 
 def create_parameters_for_flyscan_xray_centre(
-    parameters: GridCommon,
+    parameters: GenericGrid,
     grid_parameters: GridParamUpdate,
     xrc_params_type: type[SpecifiedThreeDGridScan],
 ) -> SpecifiedThreeDGridScan:

src/mx_bluesky/common/experiment_plans/inner_plans/do_fgs.py

@@ -17,9 +17,6 @@
 from mx_bluesky.common.experiment_plans.inner_plans.read_hardware import (
     read_hardware_for_zocalo,
 )
-from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import (
-    GridscanPlane,
-)
 from mx_bluesky.common.parameters.constants import (
     PlanNameConstants,
 )
@@ -68,6 +65,7 @@ def kickoff_and_complete_gridscan(
     detector: EigerDetector,  # Once Eiger inherits from StandardDetector, use that type instead
     synchrotron: Synchrotron,
     scan_points: list[AxesPoints[Axis]],
+    omega_starts_deg: list[float],
     plan_during_collection: Callable[[], MsgGenerator] | None = None,
 ):
     """Triggers a grid scan motion program and waits for completion, accounting for synchrotron topup.
@@ -95,8 +93,9 @@ def kickoff_and_complete_gridscan(
             "omega_to_scan_spec": {
                 # These have to be cast to strings due to a bug in orsjon. See
                 # https://github.com/ijl/orjson/issues/414
-                str(GridscanPlane.OMEGA_XY): scan_points[0],
-                str(GridscanPlane.OMEGA_XZ): scan_points[1],
+                # See https://github.com/DiamondLightSource/mx-bluesky/issues/1631 regarding integer cast
+                str(int(omega_starts_deg[i])): scan_points[i]
+                for i in range(len(omega_starts_deg))
             },
         }
     )

src/mx_bluesky/common/experiment_plans/inner_plans/xrc_results_utils.py

@@ -70,7 +70,23 @@ def fetch_xrc_results_from_zocalo(
 
 
 def _generate_dummy_xrc_result(params: SpecifiedThreeDGridScan) -> XRayCentreResult:
-    com = [params.x_steps / 2, params.y_steps / 2, params.z_steps / 2]
+    coms = []
+    assert params.num_grids % 2 == 0, (
+        "XRC results in commissioning mode currently only works for an even number of grids"
+    )
+
+    for grid in range(int(params.num_grids / 2)):
+        # For even number of grids, Z steps are actually the even indexed y steps
+        coms.append(
+            [
+                params.x_steps / 2,
+                params.y_steps[2 * grid] / 2,
+                params.y_steps[2 * grid + 1] / 2,
+            ]
+        )
+
+    com = [sum(x) / len(x) for x in zip(*coms, strict=True)]  # Get average
+
     max_voxel = [round(p) for p in com]
     return _xrc_result_in_boxes_to_result_in_mm(
         XrcResult(