From fa1c2d5b2c420cadab9c9e833a8b74c800a0adee Mon Sep 17 00:00:00 2001 From: Oliver Silvester Date: Wed, 18 Feb 2026 17:20:28 +0000 Subject: [PATCH 1/5] Add vmxm FGS entry point --- src/mx_bluesky/beamlines/i02_1/__init__.py | 0 .../i02_1/device_setup_plans/__init__.py | 0 .../i02_1/device_setup_plans/setup_zebra.py | 43 +++++ .../i02_1/i02_1_flyscan_xray_centre_plan.py | 166 ++++++++++++++++++ .../beamlines/i02_1/parameters/gridscan.py | 2 +- .../common_flyscan_xray_centre_plan.py | 15 +- src/mx_bluesky/common/parameters/gridscan.py | 4 +- 7 files changed, 221 insertions(+), 9 deletions(-) create mode 100644 src/mx_bluesky/beamlines/i02_1/__init__.py create mode 100644 src/mx_bluesky/beamlines/i02_1/device_setup_plans/__init__.py create mode 100644 src/mx_bluesky/beamlines/i02_1/device_setup_plans/setup_zebra.py create mode 100644 src/mx_bluesky/beamlines/i02_1/i02_1_flyscan_xray_centre_plan.py diff --git a/src/mx_bluesky/beamlines/i02_1/__init__.py b/src/mx_bluesky/beamlines/i02_1/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/mx_bluesky/beamlines/i02_1/device_setup_plans/__init__.py b/src/mx_bluesky/beamlines/i02_1/device_setup_plans/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/mx_bluesky/beamlines/i02_1/device_setup_plans/setup_zebra.py b/src/mx_bluesky/beamlines/i02_1/device_setup_plans/setup_zebra.py new file mode 100644 index 0000000000..f9a282920d --- /dev/null +++ b/src/mx_bluesky/beamlines/i02_1/device_setup_plans/setup_zebra.py @@ -0,0 +1,43 @@ +import bluesky.plan_stubs as bps +from dodal.devices.zebra.zebra import Zebra + +ZEBRA_STATUS_TIMEOUT = 30 + + +# Control Eiger from motion controller. Fast shutter is configured in GDA +def setup_zebra_for_xrc_flyscan( + zebra: Zebra, + ttl_detector: int | None = None, + group="setup_zebra_for_xrc", + wait=True, +): + """ + Assumes that the motion controller, as part of its gridscan PLC, will send triggers as required to the zebra's + IN1_TTL to control the detector. The fast shutter is configured in GDA, don't need to touch it in Bluesky for now. + """ + ttl_detector = ttl_detector or zebra.mapping.outputs.TTL_EIGER + yield from bps.abs_set( + zebra.output.out_pvs[ttl_detector], + zebra.mapping.sources.IN1_TTL, + ) + if wait: + yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT) + + +def tidy_up_zebra_after_gridscan( + zebra: Zebra, + ttl_detector: int | None = None, + group="tidy_up_vmxm_zebra_after_gridscan", + wait=False, +): + ttl_detector = ttl_detector or zebra.mapping.outputs.TTL_EIGER + + """# Revert zebra to state expected by GDA""" + yield from bps.abs_set( + zebra.output.out_pvs[ttl_detector], + zebra.mapping.sources.OR1, + group=group, + ) + + if wait: + yield from bps.wait(group) diff --git a/src/mx_bluesky/beamlines/i02_1/i02_1_flyscan_xray_centre_plan.py b/src/mx_bluesky/beamlines/i02_1/i02_1_flyscan_xray_centre_plan.py new file mode 100644 index 0000000000..ba029718e0 --- /dev/null +++ b/src/mx_bluesky/beamlines/i02_1/i02_1_flyscan_xray_centre_plan.py @@ -0,0 +1,166 @@ +from functools import partial + +import bluesky.plan_stubs as bps +import bluesky.preprocessors as bpp +import pydantic +from bluesky.utils import MsgGenerator +from dodal.beamlines.i02_1 import SampleMotors, ZebraFastGridScanTwoD +from dodal.common import inject +from dodal.devices.attenuator.attenuator import ReadOnlyAttenuator +from dodal.devices.common_dcm import DoubleCrystalMonochromatorBase +from dodal.devices.eiger import EigerDetector +from dodal.devices.fast_grid_scan import ( + set_fast_grid_scan_params as set_flyscan_params_plan, +) +from dodal.devices.flux import Flux +from dodal.devices.s4_slit_gaps import S4SlitGaps +from dodal.devices.synchrotron import Synchrotron +from dodal.devices.undulator import BaseUndulator +from dodal.devices.zebra.zebra import Zebra + +from mx_bluesky.beamlines.i02_1.device_setup_plans.setup_zebra import ( + setup_zebra_for_xrc_flyscan, + tidy_up_zebra_after_gridscan, +) +from mx_bluesky.beamlines.i02_1.parameters.gridscan import SpecifiedTwoDGridScan +from mx_bluesky.common.experiment_plans.common_flyscan_xray_centre_plan import ( + BeamlineSpecificFGSFeatures, + common_flyscan_xray_centre, + construct_beamline_specific_fast_gridscan_features, +) +from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import ( + ZocaloCallback, +) +from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import ( + GridscanISPyBCallback, + generate_start_info_from_omega_map, +) +from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import ( + GridscanNexusFileCallback, +) +from mx_bluesky.common.parameters.constants import ( + EnvironmentConstants, + PlanNameConstants, +) +from mx_bluesky.common.parameters.device_composites import ( + FlyScanEssentialDevices, + GonioWithOmegaType, +) +from mx_bluesky.common.parameters.gridscan import GenericGrid +from mx_bluesky.common.utils.log import LOGGER + + +def create_gridscan_callbacks() -> tuple[ + GridscanNexusFileCallback, GridscanISPyBCallback +]: + return ( + GridscanNexusFileCallback(param_type=SpecifiedTwoDGridScan), + GridscanISPyBCallback( + param_type=GenericGrid, + emit=ZocaloCallback( + PlanNameConstants.DO_FGS, + EnvironmentConstants.ZOCALO_ENV, + generate_start_info_from_omega_map, + ), + ), + ) + + +@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True}) +class FlyScanXRayCentreComposite(FlyScanEssentialDevices[GonioWithOmegaType]): + """All devices which are directly or indirectly required by this plan""" + + zebra: Zebra + zebra_fast_grid_scan: ZebraFastGridScanTwoD + dcm: DoubleCrystalMonochromatorBase + attenuator: ReadOnlyAttenuator + flux: Flux + undulator: BaseUndulator + s4_slit_gaps: S4SlitGaps + + +def construct_i02_1_specific_features( + fgs_composite: FlyScanXRayCentreComposite, + parameters: SpecifiedTwoDGridScan, +) -> BeamlineSpecificFGSFeatures: + signals_to_read_pre_flyscan = [ + fgs_composite.synchrotron.synchrotron_mode, + fgs_composite.gonio, + fgs_composite.dcm.energy_in_keV, + fgs_composite.undulator.current_gap, + fgs_composite.s4_slit_gaps, + ] + + signals_to_read_during_collection = [ + fgs_composite.attenuator.actual_transmission, + fgs_composite.flux.flux_reading, + fgs_composite.dcm.energy_in_keV, + fgs_composite.eiger.bit_depth, + fgs_composite.eiger.cam.roi_mode, + fgs_composite.eiger.ispyb_detector_id, + ] + + return construct_beamline_specific_fast_gridscan_features( + partial(_zebra_triggering_setup), + partial(_tidy_plan, fgs_composite, group="flyscan_zebra_tidy", wait=True), + partial( + set_flyscan_params_plan, + fgs_composite.zebra_fast_grid_scan, + parameters.fast_gridscan_params, + ), + fgs_composite.zebra_fast_grid_scan, + signals_to_read_pre_flyscan, + signals_to_read_during_collection, # type: ignore # See : https://github.com/bluesky/bluesky/issues/1809 + ) + + +def _zebra_triggering_setup(fgs_composite: FlyScanXRayCentreComposite, _): + yield from setup_zebra_for_xrc_flyscan(fgs_composite.zebra) + + +def _tidy_plan( + fgs_composite: FlyScanXRayCentreComposite, group, wait=True +) -> MsgGenerator: + LOGGER.info("Tidying up Zebra") + yield from tidy_up_zebra_after_gridscan(fgs_composite.zebra) + + +def i02_1_flyscan_xray_centre( + parameters: SpecifiedTwoDGridScan, + eiger: EigerDetector = inject("eiger"), + zebra_fast_grid_scan: ZebraFastGridScanTwoD = inject("ZebraFastGridScanTwoD"), + synchrotron: Synchrotron = inject("synchrotron"), + zebra: Zebra = inject("zebra"), + gonio: SampleMotors = inject("goniometer"), + attenuator: ReadOnlyAttenuator = inject("attenuator"), + dcm: DoubleCrystalMonochromatorBase = inject("dcm"), + flux: Flux = inject("flux"), + undulator: BaseUndulator = inject("undulator"), + s4_slit_gaps: S4SlitGaps = inject("s4_slit_gaps"), +) -> MsgGenerator: + """BlueAPI entry point for XRC grid scans""" + + # Composites have to be made this way until https://github.com/DiamondLightSource/dodal/issues/874 + # is done and we can properly use composite devices in BlueAPI + composite = FlyScanXRayCentreComposite( + eiger, + synchrotron, + gonio, + zebra, + zebra_fast_grid_scan, + dcm, + attenuator, + flux, + undulator, + s4_slit_gaps, + ) + + beamline_specific = construct_i02_1_specific_features(composite, parameters) + callbacks = create_gridscan_callbacks() + + @bpp.subs_decorator(callbacks) + def decorated_flyscan_plan(): + yield from bps.null() + yield from common_flyscan_xray_centre(composite, parameters, beamline_specific) + + yield from decorated_flyscan_plan() diff --git a/src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py b/src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py index 7c31bc5830..34ee4d0212 100644 --- a/src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py +++ b/src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py @@ -24,5 +24,5 @@ def fast_gridscan_params(self) -> ZebraGridScanParamsTwoD: 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, + dwell_time_ms=self.exposure_time_s * 1000, ) diff --git a/src/mx_bluesky/common/experiment_plans/common_flyscan_xray_centre_plan.py b/src/mx_bluesky/common/experiment_plans/common_flyscan_xray_centre_plan.py index 2d2a2a3d34..b176d97045 100644 --- a/src/mx_bluesky/common/experiment_plans/common_flyscan_xray_centre_plan.py +++ b/src/mx_bluesky/common/experiment_plans/common_flyscan_xray_centre_plan.py @@ -30,7 +30,9 @@ FlyScanEssentialDevices, GonioWithOmegaType, ) -from mx_bluesky.common.parameters.gridscan import SpecifiedThreeDGridScan +from mx_bluesky.common.parameters.gridscan import ( + SpecifiedGrids, +) from mx_bluesky.common.utils.exceptions import ( SampleError, ) @@ -114,7 +116,7 @@ def construct_beamline_specific_fast_gridscan_features( def common_flyscan_xray_centre( composite: FlyScanEssentialDevices[GonioWithOmegaType], - parameters: SpecifiedThreeDGridScan, + parameters: SpecifiedGrids, beamline_specific: BeamlineSpecificFGSFeatures, ) -> MsgGenerator: """Main entry point of the MX-Bluesky x-ray centering flyscan @@ -156,7 +158,7 @@ def _decorated_flyscan(): @bpp.finalize_decorator(lambda: _overall_tidy()) def run_gridscan_and_tidy( fgs_composite: FlyScanEssentialDevices[GonioWithOmegaType], - params: SpecifiedThreeDGridScan, + params: SpecifiedGrids, beamline_specific: BeamlineSpecificFGSFeatures, ) -> MsgGenerator: yield from beamline_specific.setup_trigger_plan(fgs_composite, parameters) @@ -174,12 +176,13 @@ def run_gridscan_and_tidy( def run_gridscan( fgs_composite: FlyScanEssentialDevices[GonioWithOmegaType], - parameters: SpecifiedThreeDGridScan, + parameters: SpecifiedGrids, 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.gonio.omega, 0) + yield from bps.abs_set( + fgs_composite.gonio.omega, parameters.omega_starts_deg[0], wait=True + ) with TRACER.start_span("ispyb_hardware_readings"): yield from beamline_specific.read_pre_flyscan_plan() diff --git a/src/mx_bluesky/common/parameters/gridscan.py b/src/mx_bluesky/common/parameters/gridscan.py index 6cef9bdf04..1d2863a490 100644 --- a/src/mx_bluesky/common/parameters/gridscan.py +++ b/src/mx_bluesky/common/parameters/gridscan.py @@ -4,7 +4,7 @@ from typing import Annotated, Generic, TypeVar from dodal.devices.aperturescatterguard import ApertureValue -from dodal.devices.detector.det_dim_constants import EIGER2_X_9M_SIZE, EIGER2_X_16M_SIZE +from dodal.devices.detector.det_dim_constants import EIGER2_X_4M_SIZE, EIGER2_X_16M_SIZE from dodal.devices.detector.detector import DetectorParams from dodal.devices.fast_grid_scan import ( GridScanParamsCommon, @@ -32,7 +32,7 @@ ) DETECTOR_SIZE_PER_BEAMLINE = { - "i02-1": EIGER2_X_9M_SIZE, + "i02-1": EIGER2_X_4M_SIZE, "dev": EIGER2_X_16M_SIZE, "i03": EIGER2_X_16M_SIZE, "i04": EIGER2_X_16M_SIZE, From d0027dd338104e0954907e9e48943bf3cdaff3a8 Mon Sep 17 00:00:00 2001 From: Oliver Silvester Date: Wed, 18 Feb 2026 18:33:46 +0000 Subject: [PATCH 2/5] Fixes and tests --- .../i02_1/device_setup_plans/setup_zebra.py | 10 +- ..._centre_plan.py => i02_1_gridscan_plan.py} | 40 +---- .../beamlines/i02_1/parameters/gridscan.py | 13 ++ src/mx_bluesky/common/parameters/constants.py | 2 + src/mx_bluesky/common/parameters/gridscan.py | 3 + tests/unit_tests/beamlines/i02_1/__init__.py | 0 .../i02_1/test_i02_1_gridscan_plan.py | 151 ++++++++++++++++++ .../beamlines/i02_1/test_setup_zebra.py | 53 ++++++ 8 files changed, 234 insertions(+), 38 deletions(-) rename src/mx_bluesky/beamlines/i02_1/{i02_1_flyscan_xray_centre_plan.py => i02_1_gridscan_plan.py} (77%) create mode 100644 tests/unit_tests/beamlines/i02_1/__init__.py create mode 100644 tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py create mode 100644 tests/unit_tests/beamlines/i02_1/test_setup_zebra.py diff --git a/src/mx_bluesky/beamlines/i02_1/device_setup_plans/setup_zebra.py b/src/mx_bluesky/beamlines/i02_1/device_setup_plans/setup_zebra.py index f9a282920d..be8bd041ff 100644 --- a/src/mx_bluesky/beamlines/i02_1/device_setup_plans/setup_zebra.py +++ b/src/mx_bluesky/beamlines/i02_1/device_setup_plans/setup_zebra.py @@ -1,14 +1,16 @@ import bluesky.plan_stubs as bps from dodal.devices.zebra.zebra import Zebra +from mx_bluesky.common.parameters.constants import PlanGroupCheckpointConstants + ZEBRA_STATUS_TIMEOUT = 30 # Control Eiger from motion controller. Fast shutter is configured in GDA -def setup_zebra_for_xrc_flyscan( +def setup_zebra_for_gridscan( zebra: Zebra, ttl_detector: int | None = None, - group="setup_zebra_for_xrc", + group=PlanGroupCheckpointConstants.SETUP_ZEBRA_FOR_GRIDSCAN, wait=True, ): """ @@ -27,12 +29,12 @@ def setup_zebra_for_xrc_flyscan( def tidy_up_zebra_after_gridscan( zebra: Zebra, ttl_detector: int | None = None, - group="tidy_up_vmxm_zebra_after_gridscan", + group=PlanGroupCheckpointConstants.TIDY_ZEBRA_AFTER_GRIDSCAN, wait=False, ): + """Revert zebra to state expected by GDA""" ttl_detector = ttl_detector or zebra.mapping.outputs.TTL_EIGER - """# Revert zebra to state expected by GDA""" yield from bps.abs_set( zebra.output.out_pvs[ttl_detector], zebra.mapping.sources.OR1, diff --git a/src/mx_bluesky/beamlines/i02_1/i02_1_flyscan_xray_centre_plan.py b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py similarity index 77% rename from src/mx_bluesky/beamlines/i02_1/i02_1_flyscan_xray_centre_plan.py rename to src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py index ba029718e0..fe19a767c5 100644 --- a/src/mx_bluesky/beamlines/i02_1/i02_1_flyscan_xray_centre_plan.py +++ b/src/mx_bluesky/beamlines/i02_1/i02_1_gridscan_plan.py @@ -1,25 +1,22 @@ from functools import partial -import bluesky.plan_stubs as bps import bluesky.preprocessors as bpp import pydantic from bluesky.utils import MsgGenerator -from dodal.beamlines.i02_1 import SampleMotors, ZebraFastGridScanTwoD +from dodal.beamlines.i02_1 import ZebraFastGridScanTwoD from dodal.common import inject from dodal.devices.attenuator.attenuator import ReadOnlyAttenuator from dodal.devices.common_dcm import DoubleCrystalMonochromatorBase -from dodal.devices.eiger import EigerDetector from dodal.devices.fast_grid_scan import ( set_fast_grid_scan_params as set_flyscan_params_plan, ) from dodal.devices.flux import Flux from dodal.devices.s4_slit_gaps import S4SlitGaps -from dodal.devices.synchrotron import Synchrotron from dodal.devices.undulator import BaseUndulator from dodal.devices.zebra.zebra import Zebra from mx_bluesky.beamlines.i02_1.device_setup_plans.setup_zebra import ( - setup_zebra_for_xrc_flyscan, + setup_zebra_for_gridscan, tidy_up_zebra_after_gridscan, ) from mx_bluesky.beamlines.i02_1.parameters.gridscan import SpecifiedTwoDGridScan @@ -115,7 +112,7 @@ def construct_i02_1_specific_features( def _zebra_triggering_setup(fgs_composite: FlyScanXRayCentreComposite, _): - yield from setup_zebra_for_xrc_flyscan(fgs_composite.zebra) + yield from setup_zebra_for_gridscan(fgs_composite.zebra) def _tidy_plan( @@ -125,42 +122,17 @@ def _tidy_plan( yield from tidy_up_zebra_after_gridscan(fgs_composite.zebra) -def i02_1_flyscan_xray_centre( +def i02_1_gridscan_plan( parameters: SpecifiedTwoDGridScan, - eiger: EigerDetector = inject("eiger"), - zebra_fast_grid_scan: ZebraFastGridScanTwoD = inject("ZebraFastGridScanTwoD"), - synchrotron: Synchrotron = inject("synchrotron"), - zebra: Zebra = inject("zebra"), - gonio: SampleMotors = inject("goniometer"), - attenuator: ReadOnlyAttenuator = inject("attenuator"), - dcm: DoubleCrystalMonochromatorBase = inject("dcm"), - flux: Flux = inject("flux"), - undulator: BaseUndulator = inject("undulator"), - s4_slit_gaps: S4SlitGaps = inject("s4_slit_gaps"), + composite: FlyScanXRayCentreComposite = inject(""), ) -> MsgGenerator: - """BlueAPI entry point for XRC grid scans""" - - # Composites have to be made this way until https://github.com/DiamondLightSource/dodal/issues/874 - # is done and we can properly use composite devices in BlueAPI - composite = FlyScanXRayCentreComposite( - eiger, - synchrotron, - gonio, - zebra, - zebra_fast_grid_scan, - dcm, - attenuator, - flux, - undulator, - s4_slit_gaps, - ) + """BlueAPI entry point for i02-1 grid scans""" beamline_specific = construct_i02_1_specific_features(composite, parameters) callbacks = create_gridscan_callbacks() @bpp.subs_decorator(callbacks) def decorated_flyscan_plan(): - yield from bps.null() yield from common_flyscan_xray_centre(composite, parameters, beamline_specific) yield from decorated_flyscan_plan() diff --git a/src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py b/src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py index 34ee4d0212..6bc6ae241b 100644 --- a/src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py +++ b/src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py @@ -1,4 +1,5 @@ from dodal.devices.beamlines.i02_1.fast_grid_scan import ZebraGridScanParamsTwoD +from pydantic import model_validator from mx_bluesky.common.parameters.components import SplitScan, WithOptionalEnergyChange from mx_bluesky.common.parameters.gridscan import SpecifiedGrids @@ -26,3 +27,15 @@ def fast_gridscan_params(self) -> ZebraGridScanParamsTwoD: transmission_fraction=0.5, dwell_time_ms=self.exposure_time_s * 1000, ) + + @model_validator(mode="after") + def validate_y_axes(self): + _err_str = "must be length 1 for 2D scans" + if len(self.y_steps) != 1: + raise ValueError(f"{self.y_steps=} {_err_str}") + if len(self.y_step_sizes_um) != 1: + raise ValueError(f"{self.y_step_sizes_um=} {_err_str}") + if len(self.omega_starts_deg) != 1: + raise ValueError(f"{self.y_step_sizes_um=} {_err_str}") + + return self diff --git a/src/mx_bluesky/common/parameters/constants.py b/src/mx_bluesky/common/parameters/constants.py index d4855219cb..158dbe28af 100644 --- a/src/mx_bluesky/common/parameters/constants.py +++ b/src/mx_bluesky/common/parameters/constants.py @@ -144,6 +144,8 @@ class PlanGroupCheckpointConstants: READY_FOR_OAV = "ready_for_oav" PREPARE_APERTURE = "prepare_aperture" SETUP_ZEBRA_FOR_ROTATION = "setup_zebra_for_rotation" + SETUP_ZEBRA_FOR_GRIDSCAN = "setup_zebra_for_gridscan" + TIDY_ZEBRA_AFTER_GRIDSCAN = "tidy_zebra_after_gridscan" # Eventually replace below with https://github.com/DiamondLightSource/mx-bluesky/issues/798 diff --git a/src/mx_bluesky/common/parameters/gridscan.py b/src/mx_bluesky/common/parameters/gridscan.py index 1d2863a490..4ab7344690 100644 --- a/src/mx_bluesky/common/parameters/gridscan.py +++ b/src/mx_bluesky/common/parameters/gridscan.py @@ -251,6 +251,9 @@ def validate_y_and_z_axes(self): raise ValueError(f"{self.y_steps=} {_err_str}") if len(self.y_step_sizes_um) != 2: raise ValueError(f"{self.y_step_sizes_um=} {_err_str}") + if len(self.omega_starts_deg) != 2: + raise ValueError(f"{self.y_step_sizes_um=} {_err_str}") + return self @property diff --git a/tests/unit_tests/beamlines/i02_1/__init__.py b/tests/unit_tests/beamlines/i02_1/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py b/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py new file mode 100644 index 0000000000..600ec4c6f3 --- /dev/null +++ b/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py @@ -0,0 +1,151 @@ +from unittest.mock import MagicMock, patch + +import pytest +from bluesky.run_engine import RunEngine +from dodal.beamlines import i02_1 +from dodal.beamlines.i02_1 import ZebraFastGridScanTwoD +from dodal.devices.attenuator.attenuator import ReadOnlyAttenuator +from dodal.devices.common_dcm import DoubleCrystalMonochromatorBase +from dodal.devices.eiger import EigerDetector +from dodal.devices.flux import Flux +from dodal.devices.s4_slit_gaps import S4SlitGaps +from dodal.devices.synchrotron import Synchrotron +from dodal.devices.undulator import BaseUndulator +from dodal.devices.zebra.zebra import Zebra +from pydantic import ValidationError + +from mx_bluesky.beamlines.i02_1.i02_1_gridscan_plan import ( + FlyScanXRayCentreComposite, + construct_i02_1_specific_features, + i02_1_gridscan_plan, +) +from mx_bluesky.beamlines.i02_1.parameters.gridscan import SpecifiedTwoDGridScan +from mx_bluesky.common.parameters.components import get_param_version +from mx_bluesky.common.parameters.device_composites import ( + GonioWithOmega, +) + + +@pytest.fixture +def fgs_params_two_d(tmp_path) -> SpecifiedTwoDGridScan: + return SpecifiedTwoDGridScan( + x_start_um=0, + y_starts_um=[0], + z_starts_um=[0], + y_step_sizes_um=[10], + omega_starts_deg=[0], + parameter_model_version=get_param_version(), + sample_id=0, + visit="visit", + file_name="test_file", + storage_directory=str(tmp_path), + x_steps=5, + y_steps=[3], + ) + + +@pytest.fixture +def zebra_fgs_two_d() -> ZebraFastGridScanTwoD: + device = i02_1.zebra_fast_grid_scan.build(connect_immediately=True, mock=True) + + return device + + +@pytest.fixture +def fgs_composite( + eiger: EigerDetector, + synchrotron: Synchrotron, + smargon: GonioWithOmega, + zebra_fgs_two_d: ZebraFastGridScanTwoD, + dcm: DoubleCrystalMonochromatorBase, + attenuator: ReadOnlyAttenuator, + flux: Flux, + undulator: BaseUndulator, + s4_slit_gaps: S4SlitGaps, + zebra: Zebra, +) -> FlyScanXRayCentreComposite: + return FlyScanXRayCentreComposite( + eiger, + synchrotron, + smargon, + zebra, + zebra_fgs_two_d, + dcm, + attenuator, + flux, + undulator, + s4_slit_gaps, + ) + + +@pytest.mark.parametrize( + "y_starts_um, z_starts_um, omega_starts_deg, y_step_sizes_um, y_steps, should_raise", + [ + ([1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], True), + ([1], [1], [1], [1], [1], False), + ], +) +def test_three_d_grid_scan_validation( + y_starts_um: list[float], + z_starts_um: list[float], + omega_starts_deg: list[float], + y_step_sizes_um: list[float], + y_steps: list[int], + should_raise: bool, + tmp_path, +): + if should_raise: + with pytest.raises(ValidationError, match="must be length 1 for 2D scans"): + SpecifiedTwoDGridScan( + x_start_um=0, + y_starts_um=y_starts_um, + z_starts_um=z_starts_um, + y_step_sizes_um=y_step_sizes_um, + omega_starts_deg=omega_starts_deg, + parameter_model_version=get_param_version(), + sample_id=0, + visit="visit", + file_name="test_file", + storage_directory=str(tmp_path), + x_steps=5, + y_steps=y_steps, + ) + else: + SpecifiedTwoDGridScan( + x_start_um=0, + y_starts_um=[0], + z_starts_um=[0], + y_step_sizes_um=[10], + omega_starts_deg=[0], + parameter_model_version=get_param_version(), + sample_id=0, + visit="visit", + file_name="test_file", + storage_directory=str(tmp_path), + x_steps=5, + y_steps=[3], + ) + + +@patch( + "mx_bluesky.beamlines.i02_1.i02_1_gridscan_plan.construct_i02_1_specific_features", +) +@patch( + "mx_bluesky.beamlines.i02_1.i02_1_gridscan_plan.common_flyscan_xray_centre", +) +def test_i02_1_flyscan_xray_centre_in_re( + mock_common_scan: MagicMock, + mock_create_features: MagicMock, + run_engine: RunEngine, + fgs_params_two_d: SpecifiedTwoDGridScan, + fgs_composite: FlyScanXRayCentreComposite, +): + expected_features = construct_i02_1_specific_features( + fgs_composite, fgs_params_two_d + ) + + mock_create_features.return_value = expected_features + run_engine(i02_1_gridscan_plan(fgs_params_two_d, fgs_composite)) + mock_common_scan.assert_called_once_with( + fgs_composite, fgs_params_two_d, expected_features + ) diff --git a/tests/unit_tests/beamlines/i02_1/test_setup_zebra.py b/tests/unit_tests/beamlines/i02_1/test_setup_zebra.py new file mode 100644 index 0000000000..cd9d02c37b --- /dev/null +++ b/tests/unit_tests/beamlines/i02_1/test_setup_zebra.py @@ -0,0 +1,53 @@ +import pytest +from bluesky.simulators import RunEngineSimulator, assert_message_and_return_remaining +from dodal.beamlines import i02_1 +from dodal.devices.zebra.zebra import Zebra + +from mx_bluesky.beamlines.i02_1.device_setup_plans.setup_zebra import ( + setup_zebra_for_gridscan, + tidy_up_zebra_after_gridscan, +) +from mx_bluesky.common.parameters.constants import PlanGroupCheckpointConstants + + +@pytest.fixture +def zebra(): + return i02_1.zebra.build(connect_immediately=True, mock=True) + + +async def test_zebra_set_up_for_gridscan( + sim_run_engine: RunEngineSimulator, + zebra: Zebra, +): + msgs = sim_run_engine.simulate_plan(setup_zebra_for_gridscan(zebra)) + assert_message_and_return_remaining( + msgs, + lambda msg: msg.command == "set" + and msg.obj.name == f"zebra-output-out_pvs-{zebra.mapping.outputs.TTL_EIGER}" + and msg.args[0] == zebra.mapping.sources.IN1_TTL, + ) + assert_message_and_return_remaining( + msgs, + lambda msg: msg.command == "wait" + and msg.kwargs["group"] + == PlanGroupCheckpointConstants.SETUP_ZEBRA_FOR_GRIDSCAN, + ) + + +async def test_tidy_up_zebra_after_gridscan( + sim_run_engine: RunEngineSimulator, + zebra: Zebra, +): + msgs = sim_run_engine.simulate_plan(tidy_up_zebra_after_gridscan(zebra, wait=True)) + assert_message_and_return_remaining( + msgs, + lambda msg: msg.command == "set" + and msg.obj.name == f"zebra-output-out_pvs-{zebra.mapping.outputs.TTL_EIGER}" + and msg.args[0] == zebra.mapping.sources.OR1, + ) + assert_message_and_return_remaining( + msgs, + lambda msg: msg.command == "wait" + and msg.kwargs["group"] + == PlanGroupCheckpointConstants.TIDY_ZEBRA_AFTER_GRIDSCAN, + ) From 6b017924ba254efe96633cb5266ff8549697dbb5 Mon Sep 17 00:00:00 2001 From: Oliver Silvester Date: Wed, 18 Feb 2026 18:37:07 +0000 Subject: [PATCH 3/5] Typo --- src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py b/src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py index 6bc6ae241b..e60a1f534d 100644 --- a/src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py +++ b/src/mx_bluesky/beamlines/i02_1/parameters/gridscan.py @@ -36,6 +36,6 @@ def validate_y_axes(self): if len(self.y_step_sizes_um) != 1: raise ValueError(f"{self.y_step_sizes_um=} {_err_str}") if len(self.omega_starts_deg) != 1: - raise ValueError(f"{self.y_step_sizes_um=} {_err_str}") + raise ValueError(f"{self.omega_starts_deg=} {_err_str}") return self From 4bce4b2f266c08164e6be87f1517fa3e13cdf2d4 Mon Sep 17 00:00:00 2001 From: Oliver Silvester Date: Wed, 18 Feb 2026 18:44:07 +0000 Subject: [PATCH 4/5] make codecov happy --- src/mx_bluesky/common/parameters/gridscan.py | 2 +- tests/unit_tests/common/parameters/test_gridscan.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/mx_bluesky/common/parameters/gridscan.py b/src/mx_bluesky/common/parameters/gridscan.py index 4ab7344690..71469e56eb 100644 --- a/src/mx_bluesky/common/parameters/gridscan.py +++ b/src/mx_bluesky/common/parameters/gridscan.py @@ -252,7 +252,7 @@ def validate_y_and_z_axes(self): if len(self.y_step_sizes_um) != 2: raise ValueError(f"{self.y_step_sizes_um=} {_err_str}") if len(self.omega_starts_deg) != 2: - raise ValueError(f"{self.y_step_sizes_um=} {_err_str}") + raise ValueError(f"{self.omega_starts_deg=} {_err_str}") return self diff --git a/tests/unit_tests/common/parameters/test_gridscan.py b/tests/unit_tests/common/parameters/test_gridscan.py index b619364d64..100413e889 100644 --- a/tests/unit_tests/common/parameters/test_gridscan.py +++ b/tests/unit_tests/common/parameters/test_gridscan.py @@ -78,6 +78,8 @@ def test_specified_grids_validation_error( "y_starts_um, z_starts_um, omega_starts_deg, y_step_sizes_um, y_steps, should_raise", [ ([1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], True), + ([1, 1], [1, 1], [1, 1], [1, 1, 1], [1, 1], False), + ([1, 1], [1, 1], [1, 1, 1], [1, 1], [1, 1], False), ([1, 1], [1, 1], [1, 1], [1, 1], [1, 1], False), ], ) From b817bd7f136ebafcd0c729aee409a4c4f8e29a24 Mon Sep 17 00:00:00 2001 From: Oliver Silvester Date: Thu, 19 Feb 2026 09:45:36 +0000 Subject: [PATCH 5/5] improve tests --- .../i02_1/test_i02_1_gridscan_plan.py | 46 ++++++------ .../common/parameters/test_gridscan.py | 71 ++++++++++--------- 2 files changed, 58 insertions(+), 59 deletions(-) diff --git a/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py b/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py index 600ec4c6f3..49f5cb5513 100644 --- a/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py +++ b/tests/unit_tests/beamlines/i02_1/test_i02_1_gridscan_plan.py @@ -78,14 +78,22 @@ def fgs_composite( ) +class SpecifiedTwoDTest(SpecifiedTwoDGridScan): + # Skip parent validation for easier testing + def _check_lengths_are_same(self): # type: ignore + return self + + @pytest.mark.parametrize( "y_starts_um, z_starts_um, omega_starts_deg, y_step_sizes_um, y_steps, should_raise", [ ([1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], True), + ([1], [1], [1], [1, 1], [1], True), + ([1], [1], [1, 1], [1], [1], True), ([1], [1], [1], [1], [1], False), ], ) -def test_three_d_grid_scan_validation( +def test_two_d_grid_scan_validation( y_starts_um: list[float], z_starts_um: list[float], omega_starts_deg: list[float], @@ -94,38 +102,28 @@ def test_three_d_grid_scan_validation( should_raise: bool, tmp_path, ): - if should_raise: - with pytest.raises(ValidationError, match="must be length 1 for 2D scans"): - SpecifiedTwoDGridScan( - x_start_um=0, - y_starts_um=y_starts_um, - z_starts_um=z_starts_um, - y_step_sizes_um=y_step_sizes_um, - omega_starts_deg=omega_starts_deg, - parameter_model_version=get_param_version(), - sample_id=0, - visit="visit", - file_name="test_file", - storage_directory=str(tmp_path), - x_steps=5, - y_steps=y_steps, - ) - else: - SpecifiedTwoDGridScan( + def create_params(): + SpecifiedTwoDTest( x_start_um=0, - y_starts_um=[0], - z_starts_um=[0], - y_step_sizes_um=[10], - omega_starts_deg=[0], + y_starts_um=y_starts_um, + z_starts_um=z_starts_um, + y_step_sizes_um=y_step_sizes_um, + omega_starts_deg=omega_starts_deg, parameter_model_version=get_param_version(), sample_id=0, visit="visit", file_name="test_file", storage_directory=str(tmp_path), x_steps=5, - y_steps=[3], + y_steps=y_steps, ) + if should_raise: + with pytest.raises(ValidationError, match="must be length 1 for 2D scans"): + create_params() + else: + create_params() + @patch( "mx_bluesky.beamlines.i02_1.i02_1_gridscan_plan.construct_i02_1_specific_features", diff --git a/tests/unit_tests/common/parameters/test_gridscan.py b/tests/unit_tests/common/parameters/test_gridscan.py index 100413e889..60e592e15d 100644 --- a/tests/unit_tests/common/parameters/test_gridscan.py +++ b/tests/unit_tests/common/parameters/test_gridscan.py @@ -39,25 +39,7 @@ def test_specified_grids_validation_error( y_steps: list[int], should_raise: bool, ): - if should_raise: - with pytest.raises( - ValidationError, match="Fields must all have the same length:" - ): - GridParamsTest( - x_start_um=0, - y_starts_um=y_starts_um, - z_starts_um=z_starts_um, - omega_starts_deg=omega_starts_deg, - y_step_sizes_um=y_step_sizes_um, - y_steps=y_steps, - sample_id=0, - visit="/tmp", - parameter_model_version=get_param_version(), - file_name="/tmp", - storage_directory="/tmp", - x_steps=5, - ) - else: + def make_params(): GridParamsTest( x_start_um=0, y_starts_um=y_starts_um, @@ -73,13 +55,27 @@ def test_specified_grids_validation_error( x_steps=5, ) + if should_raise: + with pytest.raises( + ValidationError, match="Fields must all have the same length:" + ): + make_params() + else: + make_params() + + +class SpecifiedThreeDTest(SpecifiedThreeDGridScan): + # Skip parent validation for easier testing + def _check_lengths_are_same(self): # type: ignore + return self + @pytest.mark.parametrize( "y_starts_um, z_starts_um, omega_starts_deg, y_step_sizes_um, y_steps, should_raise", [ ([1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], True), - ([1, 1], [1, 1], [1, 1], [1, 1, 1], [1, 1], False), - ([1, 1], [1, 1], [1, 1, 1], [1, 1], [1, 1], False), + ([1, 1], [1, 1], [1, 1], [1, 1, 1], [1, 1], True), + ([1, 1], [1, 1], [1, 1, 1], [1, 1], [1, 1], True), ([1, 1], [1, 1], [1, 1], [1, 1], [1, 1], False), ], ) @@ -91,19 +87,24 @@ def test_three_d_grid_scan_validation( y_steps: list[int], should_raise: bool, ): + def make_params(): + SpecifiedThreeDTest( + x_start_um=0, + y_starts_um=y_starts_um, + z_starts_um=z_starts_um, + omega_starts_deg=omega_starts_deg, + y_step_sizes_um=y_step_sizes_um, + y_steps=y_steps, + sample_id=0, + visit="/tmp", + parameter_model_version=get_param_version(), + file_name="/tmp", + storage_directory="/tmp", + x_steps=5, + ) + if should_raise: with pytest.raises(ValidationError, match="must be length 2 for 3D scans"): - SpecifiedThreeDGridScan( - x_start_um=0, - y_starts_um=y_starts_um, - z_starts_um=z_starts_um, - omega_starts_deg=omega_starts_deg, - y_step_sizes_um=y_step_sizes_um, - y_steps=y_steps, - sample_id=0, - visit="/tmp", - parameter_model_version=get_param_version(), - file_name="/tmp", - storage_directory="/tmp", - x_steps=5, - ) + make_params() + else: + make_params()