From 183066109e61ae24890bf80e8023c3179ac26f36 Mon Sep 17 00:00:00 2001 From: Dominic Oram Date: Wed, 4 Sep 2024 17:30:43 +0100 Subject: [PATCH] Change shutter to auto when configuring the zebra --- .../device_setup_plans/setup_zebra.py | 59 +++++++++++++++---- .../flyscan_xray_centre_plan.py | 24 ++++---- .../experiment_plans/rotation_scan_plan.py | 17 ++---- .../device_setup_plans/test_zebra_setup.py | 28 +++++---- 4 files changed, 79 insertions(+), 49 deletions(-) diff --git a/src/mx_bluesky/hyperion/device_setup_plans/setup_zebra.py b/src/mx_bluesky/hyperion/device_setup_plans/setup_zebra.py index 0e59753aa..bf15546bd 100644 --- a/src/mx_bluesky/hyperion/device_setup_plans/setup_zebra.py +++ b/src/mx_bluesky/hyperion/device_setup_plans/setup_zebra.py @@ -22,6 +22,7 @@ RotationDirection, Zebra, ) +from dodal.devices.zebra_controlled_shutter import ZebraShutter, ZebraShutterControl from mx_bluesky.hyperion.log import LOGGER @@ -57,25 +58,34 @@ def arm_zebra(zebra: Zebra): yield from bps.abs_set(zebra.pc.arm, ArmDemand.ARM, wait=True) -def disarm_zebra(zebra: Zebra): - yield from bps.abs_set(zebra.pc.arm, ArmDemand.DISARM, wait=True) +def tidy_up_zebra_after_rotation_scan( + zebra: Zebra, + zebra_shutter: ZebraShutter, + group="tidy_up_zebra_after_rotation", + wait=True, +): + yield from bps.abs_set(zebra.pc.arm, ArmDemand.DISARM, group=group) + yield from bps.abs_set( + zebra_shutter.control_mode, ZebraShutterControl.MANUAL, group=group + ) + if wait: + yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT) -def set_shutter_auto_input( - zebra: Zebra, input: int, group="set_shutter_trigger", wait=True -): +def set_shutter_auto_input(zebra: Zebra, input: int, group="set_shutter_trigger"): """Set the input that the shutter uses when set to auto. For more details see the ZebraShutter device.""" auto_shutter_control = zebra.logic_gates.and_gates[AUTO_SHUTTER_GATE] yield from bps.abs_set( - auto_shutter_control.sources[AUTO_SHUTTER_INPUT], input, group, wait + auto_shutter_control.sources[AUTO_SHUTTER_INPUT], input, group ) @bluesky_retry def setup_zebra_for_rotation( zebra: Zebra, + zebra_shutter: ZebraShutter, axis: EncEnum = I03Axes.OMEGA, start_angle: float = 0, scan_width: float = 360, @@ -128,7 +138,10 @@ def setup_zebra_for_rotation( # Set gate position to be angle of interest yield from bps.abs_set(zebra.pc.gate_trigger, axis.value, group=group) # Trigger the shutter with the gate - yield from set_shutter_auto_input(zebra, PC_GATE) + yield from bps.abs_set( + zebra_shutter.control_mode, ZebraShutterControl.AUTO, group=group + ) + yield from set_shutter_auto_input(zebra, PC_GATE, group=group) # Trigger the detector with a pulse yield from bps.abs_set(zebra.output.out_pvs[TTL_DETECTOR], PC_PULSE, group=group) # Don't use the fluorescence detector @@ -140,9 +153,17 @@ def setup_zebra_for_rotation( @bluesky_retry -def setup_zebra_for_gridscan(zebra: Zebra, group="setup_zebra_for_gridscan", wait=True): +def setup_zebra_for_gridscan( + zebra: Zebra, + zebra_shutter: ZebraShutter, + group="setup_zebra_for_gridscan", + wait=True, +): yield from bps.abs_set(zebra.output.out_pvs[TTL_DETECTOR], IN3_TTL, group=group) - yield from set_shutter_auto_input(zebra, IN4_TTL) + yield from bps.abs_set( + zebra_shutter.control_mode, ZebraShutterControl.AUTO, group=group + ) + yield from set_shutter_auto_input(zebra, IN4_TTL, group=group) yield from bps.abs_set(zebra.output.out_pvs[TTL_XSPRESS3], DISCONNECT, group=group) yield from bps.abs_set(zebra.output.pulse_1.input, DISCONNECT, group=group) @@ -152,10 +173,16 @@ def setup_zebra_for_gridscan(zebra: Zebra, group="setup_zebra_for_gridscan", wai @bluesky_retry def tidy_up_zebra_after_gridscan( - zebra: Zebra, group="tidy_up_zebra_after_gridscan", wait=True + zebra: Zebra, + zebra_shutter: ZebraShutter, + group="tidy_up_zebra_after_gridscan", + wait=True, ) -> MsgGenerator: yield from bps.abs_set(zebra.output.out_pvs[TTL_DETECTOR], PC_PULSE, group=group) - yield from set_shutter_auto_input(zebra, PC_GATE) + yield from bps.abs_set( + zebra_shutter.control_mode, ZebraShutterControl.MANUAL, group=group + ) + yield from set_shutter_auto_input(zebra, PC_GATE, group=group) if wait: yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT) @@ -163,13 +190,19 @@ def tidy_up_zebra_after_gridscan( @bluesky_retry def setup_zebra_for_panda_flyscan( - zebra: Zebra, group="setup_zebra_for_panda_flyscan", wait=True + zebra: Zebra, + zebra_shutter: ZebraShutter, + group="setup_zebra_for_panda_flyscan", + wait=True, ): # Forwards eiger trigger signal from panda yield from bps.abs_set(zebra.output.out_pvs[TTL_DETECTOR], IN1_TTL, group=group) # Forwards signal from PPMAC to fast shutter. High while panda PLC is running - yield from set_shutter_auto_input(zebra, IN4_TTL) + yield from bps.abs_set( + zebra_shutter.control_mode, ZebraShutterControl.AUTO, group=group + ) + yield from set_shutter_auto_input(zebra, IN4_TTL, group=group) yield from bps.abs_set(zebra.output.out_pvs[TTL_XSPRESS3], DISCONNECT, group=group) diff --git a/src/mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py b/src/mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py index fc89df030..17f149241 100755 --- a/src/mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py +++ b/src/mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py @@ -35,7 +35,7 @@ from dodal.devices.undulator import Undulator from dodal.devices.xbpm_feedback import XBPMFeedback from dodal.devices.zebra import Zebra -from dodal.devices.zebra_controlled_shutter import ZebraShutter, ZebraShutterControl +from dodal.devices.zebra_controlled_shutter import ZebraShutter from dodal.devices.zocalo.zocalo_results import ( ZOCALO_READING_PLAN_NAME, ZOCALO_STAGE_GROUP, @@ -275,13 +275,6 @@ def run_gridscan( LOGGER.info("Waiting for gridscan validity check") yield from wait_for_gridscan_valid(feature_controlled.fgs_motors) - LOGGER.info("Set shutter to auto") - yield from bps.abs_set( - fgs_composite.sample_shutter.control_mode, - ZebraShutterControl.AUTO, - group=CONST.WAIT.GRID_READY_FOR_DC, - ) - LOGGER.info("Waiting for arming to finish") yield from bps.wait(CONST.WAIT.GRID_READY_FOR_DC) yield from bps.stage(fgs_composite.eiger) @@ -446,10 +439,9 @@ def _get_feature_controlled( def _generic_tidy( fgs_composite: FlyScanXRayCentreComposite, group, wait=True ) -> MsgGenerator: - LOGGER.info("Tidying up Zebra and shutter") - yield from tidy_up_zebra_after_gridscan(fgs_composite.zebra, group=group, wait=wait) - yield from bps.abs_set( - fgs_composite.sample_shutter.control_mode, ZebraShutterControl.MANUAL + LOGGER.info("Tidying up Zebra") + yield from tidy_up_zebra_after_gridscan( + fgs_composite.zebra, fgs_composite.sample_shutter, group=group, wait=wait ) LOGGER.info("Tidying up Zocalo") # make sure we don't consume any other results @@ -470,7 +462,9 @@ def _zebra_triggering_setup( parameters: ThreeDGridScan, initial_xyz: np.ndarray, ): - yield from setup_zebra_for_gridscan(fgs_composite.zebra, wait=True) + yield from setup_zebra_for_gridscan( + fgs_composite.zebra, fgs_composite.sample_shutter, wait=True + ) def _panda_triggering_setup( @@ -528,4 +522,6 @@ def _panda_triggering_setup( ) LOGGER.info("Setting up Zebra for panda flyscan") - yield from setup_zebra_for_panda_flyscan(fgs_composite.zebra, wait=True) + yield from setup_zebra_for_panda_flyscan( + fgs_composite.zebra, fgs_composite.sample_shutter, wait=True + ) diff --git a/src/mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py b/src/mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py index 1119c5260..ee6d3ab90 100644 --- a/src/mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py +++ b/src/mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py @@ -20,7 +20,7 @@ from dodal.devices.synchrotron import Synchrotron from dodal.devices.undulator import Undulator from dodal.devices.zebra import RotationDirection, Zebra -from dodal.devices.zebra_controlled_shutter import ZebraShutter, ZebraShutterControl +from dodal.devices.zebra_controlled_shutter import ZebraShutter from dodal.plans.check_topup import check_topup_and_wait_if_necessary from mx_bluesky.hyperion.device_setup_plans.manipulate_sample import ( @@ -37,8 +37,8 @@ ) from mx_bluesky.hyperion.device_setup_plans.setup_zebra import ( arm_zebra, - disarm_zebra, setup_zebra_for_rotation, + tidy_up_zebra_after_rotation_scan, ) from mx_bluesky.hyperion.device_setup_plans.utils import ( start_preparing_data_collection_then_do_plan, @@ -218,12 +218,9 @@ def _rotation_scan_plan( group=CONST.WAIT.ROTATION_READY_FOR_DC, ) - yield from bps.abs_set( - composite.sample_shutter.control_mode, ZebraShutterControl.AUTO - ) - yield from setup_zebra_for_rotation( composite.zebra, + composite.sample_shutter, start_angle=motion_values.start_scan_deg, scan_width=motion_values.scan_width_deg, direction=motion_values.direction, @@ -289,12 +286,10 @@ def _cleanup_plan(composite: RotationScanComposite, **kwargs): max_vel = yield from bps.rd(composite.smargon.omega.max_velocity) yield from cleanup_sample_environment(composite.detector_motion, group="cleanup") yield from bps.abs_set(composite.smargon.omega.velocity, max_vel, group="cleanup") - yield from bps.abs_set( - composite.sample_shutter.control_mode, - ZebraShutterControl.MANUAL, - group="cleanup", + yield from tidy_up_zebra_after_rotation_scan( + composite.zebra, composite.sample_shutter, group="cleanup", wait=False ) - yield from bpp.finalize_wrapper(disarm_zebra(composite.zebra), bps.wait("cleanup")) + yield from bps.wait("cleanup") def _move_and_rotation( diff --git a/tests/unit_tests/hyperion/device_setup_plans/test_zebra_setup.py b/tests/unit_tests/hyperion/device_setup_plans/test_zebra_setup.py index 89bcc2fa2..e2ccf0338 100644 --- a/tests/unit_tests/hyperion/device_setup_plans/test_zebra_setup.py +++ b/tests/unit_tests/hyperion/device_setup_plans/test_zebra_setup.py @@ -2,7 +2,6 @@ import pytest from bluesky import plan_stubs as bps -from bluesky.run_engine import RunEngine from dodal.beamlines import i03 from dodal.devices.zebra import ( AUTO_SHUTTER_GATE, @@ -17,6 +16,7 @@ I03Axes, Zebra, ) +from dodal.devices.zebra_controlled_shutter import ZebraShutter from mx_bluesky.hyperion.device_setup_plans.setup_zebra import ( bluesky_retry, @@ -28,11 +28,15 @@ @pytest.fixture -def zebra(): - RunEngine() +def zebra(RE): return i03.zebra(fake_with_ophyd_sim=True) +@pytest.fixture +def zebra_shutter(RE): + return i03.sample_shutter(fake_with_ophyd_sim=True) + + async def _get_shutter_input(zebra: Zebra): return ( await zebra.logic_gates.and_gates[AUTO_SHUTTER_GATE] @@ -41,27 +45,29 @@ async def _get_shutter_input(zebra: Zebra): ) -async def test_zebra_set_up_for_panda_gridscan(RE, zebra: Zebra): - RE(setup_zebra_for_panda_flyscan(zebra, wait=True)) +async def test_zebra_set_up_for_panda_gridscan( + RE, zebra: Zebra, zebra_shutter: ZebraShutter +): + RE(setup_zebra_for_panda_flyscan(zebra, zebra_shutter, wait=True)) assert await zebra.output.out_pvs[TTL_DETECTOR].get_value() == IN1_TTL assert await zebra.output.out_pvs[TTL_PANDA].get_value() == IN3_TTL assert await _get_shutter_input(zebra) == IN4_TTL -async def test_zebra_set_up_for_gridscan(RE, zebra: Zebra): - RE(setup_zebra_for_gridscan(zebra, wait=True)) +async def test_zebra_set_up_for_gridscan(RE, zebra: Zebra, zebra_shutter: ZebraShutter): + RE(setup_zebra_for_gridscan(zebra, zebra_shutter, wait=True)) assert await zebra.output.out_pvs[TTL_DETECTOR].get_value() == IN3_TTL assert await _get_shutter_input(zebra) == IN4_TTL -async def test_zebra_set_up_for_rotation(RE, zebra: Zebra): - RE(setup_zebra_for_rotation(zebra, wait=True)) +async def test_zebra_set_up_for_rotation(RE, zebra: Zebra, zebra_shutter: ZebraShutter): + RE(setup_zebra_for_rotation(zebra, zebra_shutter, wait=True)) assert await zebra.pc.gate_trigger.get_value() == I03Axes.OMEGA.value assert await zebra.pc.gate_width.get_value() == pytest.approx(360, 0.01) -async def test_zebra_cleanup(RE, zebra: Zebra): - RE(tidy_up_zebra_after_gridscan(zebra, wait=True)) +async def test_zebra_cleanup(RE, zebra: Zebra, zebra_shutter: ZebraShutter): + RE(tidy_up_zebra_after_gridscan(zebra, zebra_shutter, wait=True)) assert await zebra.output.out_pvs[TTL_DETECTOR].get_value() == PC_PULSE assert await _get_shutter_input(zebra) == PC_GATE