From 50908fd4611390668c2911957fefbb38c38874ea Mon Sep 17 00:00:00 2001 From: Tom Kane Date: Tue, 3 Nov 2026 11:50:47 +0000 Subject: [PATCH 01/12] remove redundant terms from function arguments, update with changes from dodal (XThetaStage for tmography and sample stage rename --- ....yaml => spectroscopy_stage_baseline.yaml} | 0 src/test_rig_bluesky/plans.py | 2 +- src/test_rig_bluesky/spectroscopy_plans.py | 66 ++++++++-------- src/test_rig_bluesky/tomography_plans.py | 50 ++++++------ tests/system_tests/test_plans_system.py | 20 +++-- tests/unit_tests/test_plans.py | 77 ++++++++++--------- 6 files changed, 107 insertions(+), 108 deletions(-) rename src/test_rig_bluesky/baselines/{sample_stage_baseline.yaml => spectroscopy_stage_baseline.yaml} (100%) diff --git a/src/test_rig_bluesky/baselines/sample_stage_baseline.yaml b/src/test_rig_bluesky/baselines/spectroscopy_stage_baseline.yaml similarity index 100% rename from src/test_rig_bluesky/baselines/sample_stage_baseline.yaml rename to src/test_rig_bluesky/baselines/spectroscopy_stage_baseline.yaml diff --git a/src/test_rig_bluesky/plans.py b/src/test_rig_bluesky/plans.py index e75548e..92095b1 100644 --- a/src/test_rig_bluesky/plans.py +++ b/src/test_rig_bluesky/plans.py @@ -106,7 +106,7 @@ def serialize_spec(spec: Spec[Any]) -> Any: """Render a Spec as something JSON-serializable, for scan metadata. Axes that are devices are written out by their ``name`` (e.g. - ``"sample_stage-x"``), so consumers of the run's start document can tell + ``"spectroscopy_stage-x"``), so consumers of the run's start document can tell which axis is which. """ return spec.serialize() diff --git a/src/test_rig_bluesky/spectroscopy_plans.py b/src/test_rig_bluesky/spectroscopy_plans.py index 1fb855c..b4b5fd2 100644 --- a/src/test_rig_bluesky/spectroscopy_plans.py +++ b/src/test_rig_bluesky/spectroscopy_plans.py @@ -35,10 +35,10 @@ LOGGER = logging.getLogger(__name__) -spectroscopy_detector = inject("spectroscopy_detector") -spectroscopy_diodes = inject("spectroscopy_diodes") +detector = inject("spectroscopy_detector") +diodes = inject("spectroscopy_diodes") pandabox = inject("pandabox") -sample_stage = inject("spectroscopy_stage") +stage = inject("spectroscopy_stage") # The encoder entries map to specific axes, so despite the PandA being shared # hardware this list is spectroscopy's. @@ -62,12 +62,12 @@ def spectroscopy( - spectroscopy_detector: AravisDetector = spectroscopy_detector, - sample_stage: XYZStage = sample_stage, + detector: AravisDetector = detector, + stage: XYZStage = stage, pmac: PmacIO = pmac, pandabrick: HDFPanda = pandabrick, pandabox: HDFPanda = pandabox, - spectroscopy_diodes: RGBDiodes = spectroscopy_diodes, + diodes: RGBDiodes = diodes, spec: Spec[Movable] | None = None, exposure_time: float = 0.1, amplifier_range: EnzLocumRange | None = None, @@ -78,19 +78,15 @@ def spectroscopy( """Do a spectroscopy scan, acquiring from the camera and the RGB diodes.""" # The range is only recorded once per scan, as configuration, so it must not # change between points - scan_range = amplifier_range or ( - yield from bps.rd(spectroscopy_diodes.amplifier.range) - ) + scan_range = amplifier_range or (yield from bps.rd(diodes.amplifier.range)) if scan_range == EnzLocumRange.AUTO: raise ValueError("The diode amplifier must be on a fixed range, not Auto") if amplifier_range is not None: - yield from bps.abs_set( - spectroscopy_diodes.amplifier.range, amplifier_range, wait=True - ) + yield from bps.abs_set(diodes.amplifier.range, amplifier_range, wait=True) yield from load_settings( - device=spectroscopy_detector, + device=detector, design_name="spectroscopy_detector_baseline", whitelist_pvs=[ "hdf-nd_array_port", @@ -126,17 +122,17 @@ def spectroscopy( # outside of a run. # See: https://github.com/DiamondLightSource/blueapi/issues/1211 yield from bps.mv( - *(spectroscopy_detector.driver.acquire_time, exposure_time), + *(detector.driver.acquire_time, exposure_time), *( - spectroscopy_detector.driver.acquire_period, + detector.driver.acquire_period, exposure_time + MANTA_ACQUIRE_PERIOD_PAD, ), - group="spectroscopy_detector_aquire", + group="detector_aquire", ) params: list[NDAttributeParam] = [] - for channel in list(spectroscopy_detector.roistat.channels.keys()): # type: ignore - roistatn = spectroscopy_detector.roistat.channels[channel] # type: ignore + for channel in list(detector.roistat.channels.keys()): # type: ignore + roistatn = detector.roistat.channels[channel] # type: ignore assert isinstance(roistatn, NDROIStatNIO) channel_name = yield from bps.rd(roistatn.name_) @@ -151,11 +147,11 @@ def spectroscopy( ) ) - yield from setup_ndattributes(spectroscopy_detector.roistat, params) # type: ignore + yield from setup_ndattributes(detector.roistat, params) # type: ignore yield from load_settings( - device=sample_stage, - design_name="sample_stage_baseline", + device=stage, + design_name="spectroscopy_stage_baseline", whitelist_pvs=[ "x-acceleration_time", "x-velocity", @@ -164,7 +160,7 @@ def spectroscopy( ], ) - spec = spec or Line(sample_stage.x, 0, 5, 5) # type: ignore + spec = spec or Line(stage.x, 0, 5, 5) # type: ignore # write json-serializable form of spec into metadata if metadata is None: @@ -181,7 +177,7 @@ def spectroscopy( ) yield from fly_scan( spec, - spectroscopy_detector, + detector, pmac=pmac, pandabrick=pandabrick, exposure_time=exposure_time, @@ -189,24 +185,24 @@ def spectroscopy( panda_whitelist=PANDA_WHITELIST, metadata=metadata, pandabox=pandabox, - diodes=spectroscopy_diodes, + diodes=diodes, ) elif scan_type == ScanType.STEP: LOGGER.info("Performing a step scan.") yield from spec_scan( - {spectroscopy_detector, spectroscopy_diodes, sample_stage}, + {detector, diodes, stage}, spec, # type: ignore metadata=metadata, ) def demo_spectroscopy( - spectroscopy_detector: AravisDetector = spectroscopy_detector, - sample_stage: XYZStage = sample_stage, + detector: AravisDetector = detector, + stage: XYZStage = stage, pmac: PmacIO = pmac, pandabrick: HDFPanda = pandabrick, pandabox: HDFPanda = pandabox, - spectroscopy_diodes: RGBDiodes = spectroscopy_diodes, + diodes: RGBDiodes = diodes, total_number_of_scan_points: int = 25, grid_size: float = 5.0, grid_origin_x: float = 0.0, @@ -239,23 +235,21 @@ def demo_spectroscopy( # Move to the start point # pmac should handle premove - yield from bps.mv( - *(sample_stage.x, xmin), *(sample_stage.y, ymin), group="initial_move" - ) + yield from bps.mv(*(stage.x, xmin), *(stage.y, ymin), group="initial_move") - grid = Line(sample_stage.y, ymin, ymax, ysteps) * ~Line( # type: ignore - sample_stage.x, # type: ignore + grid = Line(stage.y, ymin, ymax, ysteps) * ~Line( # type: ignore + stage.x, # type: ignore xmin, xmax, xsteps, ) yield from spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=detector, + stage=stage, pmac=pmac, pandabrick=pandabrick, pandabox=pandabox, - spectroscopy_diodes=spectroscopy_diodes, + diodes=diodes, spec=grid, exposure_time=exposure_time, metadata=metadata, diff --git a/src/test_rig_bluesky/tomography_plans.py b/src/test_rig_bluesky/tomography_plans.py index 31c5b05..6b3da22 100644 --- a/src/test_rig_bluesky/tomography_plans.py +++ b/src/test_rig_bluesky/tomography_plans.py @@ -10,10 +10,10 @@ from bluesky.protocols import Movable from bluesky.utils import MsgGenerator from dodal.common import inject +from dodal.devices.motors import XThetaStage from dodal.plans import spec_scan from ophyd_async.core import DetectorTrigger from ophyd_async.epics.adaravis import AravisDetector -from ophyd_async.epics.motor import Motor from ophyd_async.epics.pmac import PmacIO from ophyd_async.fastcs.panda import HDFPanda from pydantic import Field @@ -24,8 +24,8 @@ LOGGER = logging.getLogger(__name__) # pmac and pandabrick come from plans.py -tomography_detector = inject("tomography_detector") -tomography_stage = inject("tomography_stage") +detector = inject("tomography_detector") +stage = inject("tomography_stage") class LightSource(StrEnum): @@ -57,7 +57,7 @@ class LightSource(StrEnum): MAKO_ACQUIRE_PERIOD_PAD = 25e-3 MAKO_DETECTOR_DEADTIME = 25e-3 -TOMOGRAPHY_DETECTOR_TRIGGER = DetectorTrigger.EXTERNAL_EDGE +DETECTOR_TRIGGER = DetectorTrigger.EXTERNAL_EDGE DETECTOR_WHITELIST = [ "hdf-nd_array_port", @@ -67,7 +67,7 @@ class LightSource(StrEnum): "driver-trigger_source", ] -TOMOGRAPHY_STAGE_WHITELIST = [ +STAGE_WHITELIST = [ "acceleration_time", "velocity", ] @@ -95,14 +95,14 @@ def _setup_detector( detector.driver.acquire_period, exposure_time + MAKO_ACQUIRE_PERIOD_PAD, ), - group="tomography_detector_acquire", + group="detector_acquire", ) # @attach_data_session_metadata_decorator() def darks_flats( light_source: LightSource, - tomography_detector: AravisDetector = tomography_detector, + detector: AravisDetector = detector, num_images: int = 20, exposure_time: float = 0.1, metadata: dict[str, Any] | None = None, @@ -128,15 +128,15 @@ def darks_flats( light_source.value, ) - yield from _setup_detector(tomography_detector, exposure_time) + yield from _setup_detector(detector, exposure_time) - yield from count([tomography_detector], num=num_images, md=run_metadata) + yield from count([detector], num=num_images, md=run_metadata) def tomography( light_source: LightSource, - tomography_detector: AravisDetector = tomography_detector, - tomography_stage: Motor = tomography_stage, + detector: AravisDetector = detector, + stage: XThetaStage = stage, pmac: PmacIO = pmac, pandabrick: HDFPanda = pandabrick, exposure_time: float = 0.1, @@ -146,16 +146,16 @@ def tomography( scan_type: ScanType, ) -> MsgGenerator[None]: """Do a tomography scan.""" - yield from _setup_detector(tomography_detector, exposure_time) + yield from _setup_detector(detector, exposure_time) yield from load_settings( - device=tomography_stage, + device=stage, design_name="tomography_stage_baseline", - whitelist_pvs=TOMOGRAPHY_STAGE_WHITELIST, + whitelist_pvs=STAGE_WHITELIST, ) # One projection per degree over a full rotation - spec = spec or Line(tomography_stage, 0, 359, 360) # type: ignore + spec = spec or Line(stage, 0, 359, 360) # type: ignore if metadata is None: metadata = {} @@ -167,12 +167,12 @@ def tomography( LOGGER.info("Performing a tomography fly scan.") yield from fly_scan( spec, - tomography_detector, + detector, pmac=pmac, pandabrick=pandabrick, exposure_time=exposure_time, detector_deadtime=MAKO_DETECTOR_DEADTIME, - detector_trigger=TOMOGRAPHY_DETECTOR_TRIGGER, + detector_trigger=DETECTOR_TRIGGER, panda_whitelist=PANDA_WHITELIST, plan_name="tomography_fly_scan", metadata=metadata, @@ -180,7 +180,7 @@ def tomography( elif scan_type == ScanType.STEP: LOGGER.info("Performing a tomography step scan.") yield from spec_scan( - {tomography_detector, tomography_stage}, + {detector, stage}, spec, # type: ignore metadata=metadata, ) @@ -188,8 +188,8 @@ def tomography( def demo_tomography( light_source: LightSource, - tomography_detector: AravisDetector = tomography_detector, - tomography_stage: Motor = tomography_stage, + detector: AravisDetector = detector, + stage: XThetaStage = stage, pmac: PmacIO = pmac, pandabrick: HDFPanda = pandabrick, number_of_projections: NumProjections = 360, @@ -206,24 +206,24 @@ def demo_tomography( """ angular_step = angular_range / number_of_projections last_angle = start_angle + angular_range - angular_step - spec = Line(tomography_stage, start_angle, last_angle, number_of_projections) # type: ignore + spec = Line(stage, start_angle, last_angle, number_of_projections) # type: ignore # PMAC handles the run-up, but move to the nominal start first so a failure # here is reported before the trajectory is built. - yield from bps.mv(tomography_stage, start_angle, group="initial_move") + yield from bps.mv(stage, start_angle, group="initial_move") def wind_back() -> MsgGenerator[None]: # unwind to leave the stage where the next scan expects it. # currently this winds back the full rotation but should look # into changing it so that we just advance one step to the # beginning again and reset the angle to 0? - yield from bps.mv(tomography_stage, start_angle, group="wind_back") + yield from bps.mv(stage, start_angle, group="wind_back") yield from bpp.finalize_wrapper( tomography( light_source, - tomography_detector=tomography_detector, - tomography_stage=tomography_stage, + detector=detector, + stage=stage, pmac=pmac, pandabrick=pandabrick, exposure_time=exposure_time, diff --git a/tests/system_tests/test_plans_system.py b/tests/system_tests/test_plans_system.py index f0417dc..dd9711d 100644 --- a/tests/system_tests/test_plans_system.py +++ b/tests/system_tests/test_plans_system.py @@ -117,14 +117,16 @@ def test_generic_scan( def test_spectroscopy_re(): run_engine = RunEngine() spectroscopy_detector = b01_1.spectroscopy_detector.build(connect_immediately=True) - sample_stage = b01_1.spectroscopy_stage.build(connect_immediately=True) + spectroscopy_stage = b01_1.spectroscopy_stage.build(connect_immediately=True) - scan_spec = Line(sample_stage.y, 0, 5, 50) * Line(sample_stage.x, 2, 5, 30) + scan_spec = Line(spectroscopy_stage.y, 0, 5, 50) * Line( + spectroscopy_stage.x, 2, 5, 30 + ) run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, spec=scan_spec, scan_type=ScanType.FLY, ) @@ -168,14 +170,16 @@ def test_spectroscopy_re_run_engine(): run_engine = RunEngine() spectroscopy_detector = b01_1.spectroscopy_detector.build(connect_immediately=True) - sample_stage = b01_1.spectroscopy_stage.build(connect_immediately=True) + spectroscopy_stage = b01_1.spectroscopy_stage.build(connect_immediately=True) - scan_spec = Line(sample_stage.y, 0, 5, 50) * Line(sample_stage.x, 2, 5, 30) + scan_spec = Line(spectroscopy_stage.y, 0, 5, 50) * Line( + spectroscopy_stage.x, 2, 5, 30 + ) run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, spec=scan_spec, scan_type=ScanType.STEP, ) diff --git a/tests/unit_tests/test_plans.py b/tests/unit_tests/test_plans.py index 36260d8..37df9d9 100644 --- a/tests/unit_tests/test_plans.py +++ b/tests/unit_tests/test_plans.py @@ -46,7 +46,7 @@ def spectroscopy_detector(run_engine: RunEngine) -> AravisDetector: @pytest.fixture -def sample_stage(run_engine: RunEngine) -> XYZStage: +def spectroscopy_stage(run_engine: RunEngine) -> XYZStage: stage = b01_1.spectroscopy_stage.build(connect_immediately=True, mock=True) set_mock_value(stage.x.low_limit_travel, -10.0) @@ -143,12 +143,12 @@ async def test_load_settings( def test_snapshot( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, ): docs = defaultdict(list) run_engine.subscribe(lambda name, doc: docs[name].append(doc)) - run_engine(snapshot([spectroscopy_detector, sample_stage])) + run_engine(snapshot([spectroscopy_detector, spectroscopy_stage])) assert_emitted( docs, start=1, descriptor=1, stream_resource=1, stream_datum=1, event=1, stop=1 @@ -164,7 +164,7 @@ def test_snapshot( async def test_spectroscopy( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, pandabrick: HDFPanda, spectroscopy_diodes: RGBDiodes, ): @@ -173,11 +173,12 @@ async def test_spectroscopy( run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, pandabrick=pandabrick, - spectroscopy_diodes=spectroscopy_diodes, - spec=Line(sample_stage.y, 4.2, 6, 3) * Line(sample_stage.x, 0, 5, 10), + diodes=spectroscopy_diodes, + spec=Line(spectroscopy_stage.y, 4.2, 6, 3) + * Line(spectroscopy_stage.x, 0, 5, 10), exposure_time=0.2, scan_type=ScanType.STEP, ) @@ -199,7 +200,7 @@ async def test_spectroscopy( async def test_spectroscopy_defaults( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, pandabrick: HDFPanda, spectroscopy_diodes: RGBDiodes, ): @@ -208,10 +209,10 @@ async def test_spectroscopy_defaults( run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, pandabrick=pandabrick, - spectroscopy_diodes=spectroscopy_diodes, + diodes=spectroscopy_diodes, scan_type=ScanType.STEP, ) ) @@ -232,7 +233,7 @@ async def test_spectroscopy_defaults( def test_spectroscopy_datasets( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, pandabrick: HDFPanda, spectroscopy_diodes: RGBDiodes, ): @@ -241,10 +242,10 @@ def test_spectroscopy_datasets( run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, pandabrick=pandabrick, - spectroscopy_diodes=spectroscopy_diodes, + diodes=spectroscopy_diodes, scan_type=ScanType.STEP, ) ) @@ -264,16 +265,16 @@ def test_spectroscopy_datasets( async def test_spectroscopy_sets_exposure_time_and_acquire_period( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, pandabrick: HDFPanda, spectroscopy_diodes: RGBDiodes, ): run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, pandabrick=pandabrick, - spectroscopy_diodes=spectroscopy_diodes, + diodes=spectroscopy_diodes, exposure_time=1.0, scan_type=ScanType.STEP, ) @@ -292,8 +293,8 @@ def test_demo_spectroscopy(): ) as mock_spec: # Call the generator function and exhaust it generator = demo_spectroscopy( - spectroscopy_detector=fake_detector, - sample_stage=fake_stage, + detector=fake_detector, + stage=fake_stage, total_number_of_scan_points=25, scan_type=ScanType.FLY, ) @@ -304,31 +305,31 @@ def test_demo_spectroscopy(): mock_spec.assert_called_once() called_kwargs = mock_spec.call_args.kwargs - assert called_kwargs["spectroscopy_detector"] is fake_detector - assert called_kwargs["sample_stage"] is fake_stage + assert called_kwargs["detector"] is fake_detector + assert called_kwargs["stage"] is fake_stage assert called_kwargs["spec"] == Line(fake_stage.y, 0.0, 5.0, 5) * ~Line( fake_stage.x, 0.0, 5.0, 5 ) assert called_kwargs["scan_type"] == ScanType.FLY -def test_serialize_spec_names_device_axes(sample_stage: XYZStage): - spec = Line(sample_stage.y, 4.2, 6, 3) * Line(sample_stage.x, 0, 5, 10) +def test_serialize_spec_names_device_axes(spectroscopy_stage: XYZStage): + spec = Line(spectroscopy_stage.y, 4.2, 6, 3) * Line(spectroscopy_stage.x, 0, 5, 10) serialized = serialize_spec(spec) - assert serialized["outer"]["axis"] == sample_stage.y.name - assert serialized["inner"]["axis"] == sample_stage.x.name + assert serialized["outer"]["axis"] == spectroscopy_stage.y.name + assert serialized["inner"]["axis"] == spectroscopy_stage.x.name assert Spec.deserialize(serialized).axes() == [ - sample_stage.y.name, - sample_stage.x.name, + spectroscopy_stage.y.name, + spectroscopy_stage.x.name, ] async def test_spectroscopy_sets_amplifier_range( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, pandabrick: HDFPanda, spectroscopy_diodes: RGBDiodes, ): @@ -337,10 +338,10 @@ async def test_spectroscopy_sets_amplifier_range( run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, pandabrick=pandabrick, - spectroscopy_diodes=spectroscopy_diodes, + diodes=spectroscopy_diodes, amplifier_range=EnzLocumRange.RANGE_10NA, scan_type=ScanType.STEP, ) @@ -356,7 +357,7 @@ async def test_spectroscopy_sets_amplifier_range( def test_spectroscopy_rejects_auto_amplifier_range( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, spectroscopy_diodes: RGBDiodes, ): set_mock_value(spectroscopy_diodes.amplifier.range, EnzLocumRange.AUTO) @@ -364,9 +365,9 @@ def test_spectroscopy_rejects_auto_amplifier_range( with pytest.raises(ValueError, match="Auto"): run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, - spectroscopy_diodes=spectroscopy_diodes, + detector=spectroscopy_detector, + stage=spectroscopy_stage, + diodes=spectroscopy_diodes, scan_type=ScanType.STEP, ) ) From fee59c1c40e06427680e820f076942917006fbf0 Mon Sep 17 00:00:00 2001 From: Tom Kane Date: Tue, 3 Nov 2026 16:37:28 +0000 Subject: [PATCH 02/12] incorporate darks and flats inside tomography plan --- .../baselines/tomography_stage_baseline.yaml | 16 +- src/test_rig_bluesky/tomography_plans.py | 231 +++++++++++------- tests/unit_tests/test_plans.py | 154 +++++++++++- 3 files changed, 300 insertions(+), 101 deletions(-) diff --git a/src/test_rig_bluesky/baselines/tomography_stage_baseline.yaml b/src/test_rig_bluesky/baselines/tomography_stage_baseline.yaml index 6b9b570..f16bb5f 100644 --- a/src/test_rig_bluesky/baselines/tomography_stage_baseline.yaml +++ b/src/test_rig_bluesky/baselines/tomography_stage_baseline.yaml @@ -1,8 +1,8 @@ -acceleration_time: 0.001 -high_limit_travel: 1000 -low_limit_travel: -1000 -offset: 0.0 -offset_freeze_switch: Variable -set_use_switch: Use -user_setpoint: 5.00004 -velocity: 30 +theta.acceleration_time: 0.001 +theta.high_limit_travel: 1000 +theta.low_limit_travel: -1000 +theta.offset: 0.0 +theta.offset_freeze_switch: Variable +theta.set_use_switch: Use +theta.user_setpoint: 5.00004 +theta.velocity: 30 diff --git a/src/test_rig_bluesky/tomography_plans.py b/src/test_rig_bluesky/tomography_plans.py index 6b3da22..a35526c 100644 --- a/src/test_rig_bluesky/tomography_plans.py +++ b/src/test_rig_bluesky/tomography_plans.py @@ -2,7 +2,9 @@ import logging from enum import StrEnum +from itertools import chain from typing import Annotated, Any +from uuid import uuid4 import bluesky.plan_stubs as bps import bluesky.preprocessors as bpp @@ -10,6 +12,7 @@ from bluesky.protocols import Movable from bluesky.utils import MsgGenerator from dodal.common import inject +from dodal.devices.beamlines.b01_1.flip_mirror import FlipMirror from dodal.devices.motors import XThetaStage from dodal.plans import spec_scan from ophyd_async.core import DetectorTrigger @@ -26,6 +29,8 @@ # pmac and pandabrick come from plans.py detector = inject("tomography_detector") stage = inject("tomography_stage") +tomography_flip_mirror = inject("tomography_flip_mirror") +spectroscopy_flip_mirror = inject("spectroscopy_flip_mirror") class LightSource(StrEnum): @@ -39,6 +44,17 @@ class LightSource(StrEnum): # Metadata keys the reconstruction workflow reads to pair projections with the # calibration images taken under the same illumination LIGHT_SOURCE_KEY = "light_source" +# "dark", "flat" or "projection" +IMAGE_KEY = "image_key" +# Shared by the darks, flats and projections runs of one tomography scan +GROUP_KEY = "tomography_group" +# Recorded on the projections run, pointing at its calibration runs +DARKS_UID_KEY = "darks_uid" +FLATS_UID_KEY = "flats_uid" + +# Tomography stage x positions, in mm +SAMPLE_OUT_X = 0.0 # TODO: placeholder, sample out of the field of view for flats +SAMPLE_IN_X = 0.0 # TODO: placeholder, sample back in the field of view MIN_PROJECTIONS = 30 MAX_PROJECTIONS = 1440 @@ -68,85 +84,130 @@ class LightSource(StrEnum): ] STAGE_WHITELIST = [ - "acceleration_time", - "velocity", + "theta-acceleration_time", + "theta-velocity", ] # TODO: confirm which PandA encoder input the rotation axis is wired to. PANDA_WHITELIST = ["incenc-8-val_dataset"] -def _setup_detector( +def _darks_flats_projections( detector: AravisDetector, + stage: XThetaStage, + tomography_flip_mirror: FlipMirror, + spectroscopy_flip_mirror: FlipMirror, + pmac: PmacIO, + pandabrick: HDFPanda, exposure_time: float, + spec: Spec[Movable], + num_darks: int, + num_flats: int, + metadata: dict[str, Any], + scan_type: ScanType, ) -> MsgGenerator[None]: - """Put the detector in its scanning state.""" - yield from load_settings( - device=detector, - design_name="tomography_detector_baseline", - whitelist_pvs=DETECTOR_WHITELIST, + """Take darks, then flats with the sample out of the way, then projections.""" + LOGGER.info("Collecting %d dark images.", num_darks) + yield from bps.mv( + *(stage.x, SAMPLE_OUT_X), + *(spectroscopy_flip_mirror, False), + *(tomography_flip_mirror, False), + group="darks", + ) + darks_uid = yield from count( + [detector], num=num_darks, md={IMAGE_KEY: "dark", **metadata} ) - # mv rather than prepare, because prepare cannot be used outside a run. - # See: https://github.com/DiamondLightSource/blueapi/issues/1211 - yield from bps.mv( - *(detector.driver.acquire_time, exposure_time), - *( - detector.driver.acquire_period, - exposure_time + MAKO_ACQUIRE_PERIOD_PAD, - ), - group="detector_acquire", + LOGGER.info("Collecting %d flat images.", num_flats) + yield from bps.mv(tomography_flip_mirror, True, group="flats") + flats_uid = yield from count( + [detector], num=num_flats, md={IMAGE_KEY: "flat", **metadata} ) + yield from bps.mv(stage.x, SAMPLE_IN_X, group="projections") + projections_metadata = { + IMAGE_KEY: "projection", + DARKS_UID_KEY: darks_uid, + FLATS_UID_KEY: flats_uid, + "spec": serialize_spec(spec), + **metadata, + } -# @attach_data_session_metadata_decorator() -def darks_flats( - light_source: LightSource, - detector: AravisDetector = detector, - num_images: int = 20, - exposure_time: float = 0.1, - metadata: dict[str, Any] | None = None, -) -> MsgGenerator[None]: - """Collect flat or dark field images for the reconstruction to normalise with. + if scan_type == ScanType.FLY: + LOGGER.info("Performing a tomography fly scan.") + yield from fly_scan( + spec, + detector, + pmac=pmac, + pandabrick=pandabrick, + exposure_time=exposure_time, + detector_deadtime=MAKO_DETECTOR_DEADTIME, + detector_trigger=DETECTOR_TRIGGER, + panda_whitelist=PANDA_WHITELIST, + plan_name="tomography_fly_scan", + metadata=projections_metadata, + ) + elif scan_type == ScanType.STEP: + LOGGER.info("Performing a tomography step scan.") + yield from spec_scan( + {detector, stage}, + spec, # type: ignore + metadata=projections_metadata, + ) - Recorded in the run's metadata under "light_source" so the workflow can find - the right images for a given scan. LightSource.NONE means the beam is off, - i.e. dark fields, and one such set serves every illumination. - :param light_source: illumination the images are taken under, or - LightSource.NONE for dark fields. - """ - run_metadata: dict[str, Any] = {LIGHT_SOURCE_KEY: light_source.value} - run_metadata.update(metadata or {}) - - if light_source is LightSource.DARK: - LOGGER.info("Collecting %d dark images.", num_images) - else: - LOGGER.info( - "Collecting %d flat images under %s illumination.", - num_images, - light_source.value, - ) +def _restore_mirrors(positions: list[tuple[FlipMirror, bool]]) -> MsgGenerator[None]: + yield from bps.mv(*chain.from_iterable(positions)) - yield from _setup_detector(detector, exposure_time) - yield from count([detector], num=num_images, md=run_metadata) +def _wind_back(stage: XThetaStage, start_angle: float) -> MsgGenerator[None]: + # unwind to leave the stage where the next scan expects it. + # currently this winds back the full rotation but should look + # into changing it so that we just advance one step to the + # beginning again and reset the angle to 0? + yield from bps.mv(stage.theta, start_angle, group="wind_back") def tomography( light_source: LightSource, detector: AravisDetector = detector, stage: XThetaStage = stage, + tomography_flip_mirror: FlipMirror = tomography_flip_mirror, + spectroscopy_flip_mirror: FlipMirror = spectroscopy_flip_mirror, pmac: PmacIO = pmac, pandabrick: HDFPanda = pandabrick, exposure_time: float = 0.1, spec: Spec[Movable] | None = None, + num_darks: int = 20, + num_flats: int = 20, metadata: dict[str, Any] | None = None, *, scan_type: ScanType, ) -> MsgGenerator[None]: - """Do a tomography scan.""" - yield from _setup_detector(detector, exposure_time) + """Do a tomography scan, taking darks and flats first. + + The darks, flats and projections are each their own run, sharing a + "tomography_group" id, and the projections run records the uids of the other + two. The flip mirrors are put back as they were afterwards. + + :param light_source: illumination the flats and projections are taken under. + """ + yield from load_settings( + device=detector, + design_name="tomography_detector_baseline", + whitelist_pvs=DETECTOR_WHITELIST, + ) + + # mv rather than prepare, because prepare cannot be used outside a run. + # See: https://github.com/DiamondLightSource/blueapi/issues/1211 + yield from bps.mv( + *(detector.driver.acquire_time, exposure_time), + *( + detector.driver.acquire_period, + exposure_time + MAKO_ACQUIRE_PERIOD_PAD, + ), + group="detector_acquire", + ) yield from load_settings( device=stage, @@ -155,41 +216,48 @@ def tomography( ) # One projection per degree over a full rotation - spec = spec or Line(stage, 0, 359, 360) # type: ignore + spec = spec or Line(stage.theta, 0, 359, 360) # type: ignore - if metadata is None: - metadata = {} - metadata["spec"] = serialize_spec(spec) # Tells the workflow which calibration images belong with these projections. - metadata[LIGHT_SOURCE_KEY] = light_source.value + shared_metadata: dict[str, Any] = { + GROUP_KEY: str(uuid4()), + LIGHT_SOURCE_KEY: light_source.value, + **(metadata or {}), + } + + initial_mirror_positions = [ + (tomography_flip_mirror, (yield from bps.rd(tomography_flip_mirror.position))), + ( + spectroscopy_flip_mirror, + (yield from bps.rd(spectroscopy_flip_mirror.position)), + ), + ] - if scan_type == ScanType.FLY: - LOGGER.info("Performing a tomography fly scan.") - yield from fly_scan( - spec, + yield from bpp.finalize_wrapper( + _darks_flats_projections( detector, - pmac=pmac, - pandabrick=pandabrick, - exposure_time=exposure_time, - detector_deadtime=MAKO_DETECTOR_DEADTIME, - detector_trigger=DETECTOR_TRIGGER, - panda_whitelist=PANDA_WHITELIST, - plan_name="tomography_fly_scan", - metadata=metadata, - ) - elif scan_type == ScanType.STEP: - LOGGER.info("Performing a tomography step scan.") - yield from spec_scan( - {detector, stage}, - spec, # type: ignore - metadata=metadata, - ) + stage, + tomography_flip_mirror, + spectroscopy_flip_mirror, + pmac, + pandabrick, + exposure_time, + spec, + num_darks, + num_flats, + shared_metadata, + scan_type, + ), + _restore_mirrors(initial_mirror_positions), + ) def demo_tomography( light_source: LightSource, detector: AravisDetector = detector, stage: XThetaStage = stage, + tomography_flip_mirror: FlipMirror = tomography_flip_mirror, + spectroscopy_flip_mirror: FlipMirror = spectroscopy_flip_mirror, pmac: PmacIO = pmac, pandabrick: HDFPanda = pandabrick, number_of_projections: NumProjections = 360, @@ -201,29 +269,24 @@ def demo_tomography( scan_type: ScanType, ) -> MsgGenerator[None]: """Tomography plan intended for use in Visr demonstrations to visitors. - Takes number_of_projections evenly spaced over angular_range from start_angle, - then winds the stage back to start_angle. + Takes darks and flats, then number_of_projections evenly spaced over + angular_range from start_angle, then winds the stage back to start_angle. """ angular_step = angular_range / number_of_projections last_angle = start_angle + angular_range - angular_step - spec = Line(stage, start_angle, last_angle, number_of_projections) # type: ignore + spec = Line(stage.theta, start_angle, last_angle, number_of_projections) # type: ignore # PMAC handles the run-up, but move to the nominal start first so a failure # here is reported before the trajectory is built. - yield from bps.mv(stage, start_angle, group="initial_move") - - def wind_back() -> MsgGenerator[None]: - # unwind to leave the stage where the next scan expects it. - # currently this winds back the full rotation but should look - # into changing it so that we just advance one step to the - # beginning again and reset the angle to 0? - yield from bps.mv(stage, start_angle, group="wind_back") + yield from bps.mv(stage.theta, start_angle, group="initial_move") yield from bpp.finalize_wrapper( tomography( light_source, detector=detector, stage=stage, + tomography_flip_mirror=tomography_flip_mirror, + spectroscopy_flip_mirror=spectroscopy_flip_mirror, pmac=pmac, pandabrick=pandabrick, exposure_time=exposure_time, @@ -231,5 +294,5 @@ def wind_back() -> MsgGenerator[None]: metadata=metadata, scan_type=scan_type, ), - wind_back, + _wind_back(stage, start_angle), ) diff --git a/tests/unit_tests/test_plans.py b/tests/unit_tests/test_plans.py index 37df9d9..15456d8 100644 --- a/tests/unit_tests/test_plans.py +++ b/tests/unit_tests/test_plans.py @@ -9,8 +9,9 @@ import yaml from bluesky import RunEngine from dodal.devices.beamlines.b01_1.diodes import RGBDiodes +from dodal.devices.beamlines.b01_1.flip_mirror import FlipMirror from dodal.devices.current_amplifiers import EnzLocumRange -from dodal.devices.motors import XYZStage +from dodal.devices.motors import XThetaStage, XYZStage from ophyd_async.core import ( callback_on_mock_put, get_mock_put, @@ -36,6 +37,12 @@ demo_spectroscopy, spectroscopy, ) +from test_rig_bluesky.tomography_plans import ( + SAMPLE_IN_X, + SAMPLE_OUT_X, + LightSource, + tomography, +) @pytest.fixture @@ -92,8 +99,14 @@ async def on_acquire(acquire: bool) -> None: if acquire: asyncio.create_task(mock_acquisition()) + def on_capture(capture: bool) -> None: + # Like the IOC, start counting from zero for each new file + if capture: + set_mock_value(hdf.num_captured, 0) + set_mock_value(hdf.file_path_exists, True) callback_on_mock_put(detector.driver.acquire, on_acquire) + callback_on_mock_put(hdf.capture, on_capture) @patch("test_rig_bluesky.plans.YamlSettingsProvider") @@ -376,13 +389,13 @@ def test_spectroscopy_rejects_auto_amplifier_range( def test_fly_scan_rejects_pandabox_without_diodes( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, pandabox: HDFPanda, ): with pytest.raises(ValueError, match="pandabox was given without diodes"): run_engine( fly_scan( - Line(sample_stage.x, 0, 5, 5), + Line(spectroscopy_stage.x, 0, 5, 5), spectroscopy_detector, pandabox=pandabox, ) @@ -392,13 +405,13 @@ def test_fly_scan_rejects_pandabox_without_diodes( def test_fly_scan_rejects_diodes_without_pandabox( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, spectroscopy_diodes: RGBDiodes, ): with pytest.raises(ValueError, match="diodes were given without pandabox"): run_engine( fly_scan( - Line(sample_stage.x, 0, 5, 5), + Line(spectroscopy_stage.x, 0, 5, 5), spectroscopy_detector, diodes=spectroscopy_diodes, ) @@ -408,7 +421,7 @@ def test_fly_scan_rejects_diodes_without_pandabox( def test_spectroscopy_fly_scan_uses_pandabox_and_diodes( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, pandabrick: HDFPanda, pandabox: HDFPanda, spectroscopy_diodes: RGBDiodes, @@ -424,11 +437,11 @@ def test_spectroscopy_fly_scan_uses_pandabox_and_diodes( ): run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, pandabrick=pandabrick, pandabox=pandabox, - spectroscopy_diodes=spectroscopy_diodes, + diodes=spectroscopy_diodes, scan_type=ScanType.FLY, ) ) @@ -455,3 +468,126 @@ def test_pandabox_baseline_captures_diodes_under_their_names( assert datasets["fmc_in.val5_dataset"] == spectroscopy_diodes.red.name assert datasets["fmc_in.val6_dataset"] == spectroscopy_diodes.green.name assert datasets["fmc_in.val7_dataset"] == spectroscopy_diodes.blue.name + + +@pytest.fixture +def tomography_detector(run_engine: RunEngine) -> AravisDetector: + det = b01_1.tomography_detector.build(connect_immediately=True, mock=True) + _mock_detector_behavior(det) + return det + + +@pytest.fixture +def tomography_stage(run_engine: RunEngine) -> XThetaStage: + stage = b01_1.tomography_stage.build(connect_immediately=True, mock=True) + + set_mock_value(stage.x.low_limit_travel, -10.0) + set_mock_value(stage.x.high_limit_travel, 10.0) + set_mock_value(stage.x.velocity, 1.0) + set_mock_value(stage.x.user_readback, 5.0) + set_mock_value(stage.theta.low_limit_travel, -400.0) + set_mock_value(stage.theta.high_limit_travel, 400.0) + set_mock_value(stage.theta.velocity, 10.0) + + return stage + + +@pytest.fixture +def tomography_flip_mirror(run_engine: RunEngine) -> FlipMirror: + mirror = b01_1.tomography_flip_mirror.build(connect_immediately=True, mock=True) + set_mock_value(mirror.position, True) + return mirror + + +@pytest.fixture +def spectroscopy_flip_mirror(run_engine: RunEngine) -> FlipMirror: + mirror = b01_1.spectroscopy_flip_mirror.build(connect_immediately=True, mock=True) + set_mock_value(mirror.position, True) + return mirror + + +def test_tomography_takes_darks_then_flats_then_projections( + run_engine: RunEngine, + tomography_detector: AravisDetector, + tomography_stage: XThetaStage, + tomography_flip_mirror: FlipMirror, + spectroscopy_flip_mirror: FlipMirror, +): + messages = [] + run_engine.msg_hook = messages.append + starts = [] + run_engine.subscribe(lambda name, doc: starts.append(doc), "start") + + run_engine( + tomography( + LightSource.LED, + detector=tomography_detector, + stage=tomography_stage, + tomography_flip_mirror=tomography_flip_mirror, + spectroscopy_flip_mirror=spectroscopy_flip_mirror, + spec=Line(tomography_stage.theta, 0, 3, 4), + num_darks=2, + num_flats=3, + scan_type=ScanType.STEP, + ) + ) + + darks, flats, projections = starts + assert [start["image_key"] for start in starts] == ["dark", "flat", "projection"] + assert [start["num_points"] for start in starts] == [2, 3, 4] + assert len({start["tomography_group"] for start in starts}) == 1 + assert all(start["light_source"] == "led" for start in starts) + assert projections["darks_uid"] == darks["uid"] + assert projections["flats_uid"] == flats["uid"] + + # Each run is preceded by the moves that set it up + sets_and_opens = [ + (msg.command, msg.obj, msg.args) + for msg in messages + if msg.command == "open_run" + or ( + msg.command == "set" + and msg.obj + in (tomography_stage.x, tomography_flip_mirror, spectroscopy_flip_mirror) + ) + ] + assert sets_and_opens == [ + ("set", tomography_stage.x, (SAMPLE_OUT_X,)), + ("set", spectroscopy_flip_mirror, (False,)), + ("set", tomography_flip_mirror, (False,)), + ("open_run", None, ()), + ("set", tomography_flip_mirror, (True,)), + ("open_run", None, ()), + ("set", tomography_stage.x, (SAMPLE_IN_X,)), + ("open_run", None, ()), + # Mirrors put back as they were + ("set", tomography_flip_mirror, (True,)), + ("set", spectroscopy_flip_mirror, (True,)), + ] + + +async def test_tomography_restores_mirrors_on_failure( + run_engine: RunEngine, + tomography_detector: AravisDetector, + tomography_stage: XThetaStage, + tomography_flip_mirror: FlipMirror, + spectroscopy_flip_mirror: FlipMirror, +): + with patch( + "test_rig_bluesky.tomography_plans.count", side_effect=RuntimeError("boom") + ): + with pytest.raises(RuntimeError, match="boom"): + run_engine( + tomography( + LightSource.LED, + detector=tomography_detector, + stage=tomography_stage, + tomography_flip_mirror=tomography_flip_mirror, + spectroscopy_flip_mirror=spectroscopy_flip_mirror, + spec=Line(tomography_stage.theta, 0, 3, 4), + scan_type=ScanType.STEP, + ) + ) + + assert await tomography_flip_mirror.position.get_value() is True + assert await spectroscopy_flip_mirror.position.get_value() is True From 2448525605c5aa7a64668228294f5511cb73aa44 Mon Sep 17 00:00:00 2001 From: Tom Kane Date: Wed, 4 Nov 2026 14:30:47 +0000 Subject: [PATCH 03/12] refactor tomography --- src/test_rig_bluesky/tomography_plans.py | 248 +++++++++++++++-------- tests/unit_tests/test_plans.py | 11 +- 2 files changed, 166 insertions(+), 93 deletions(-) diff --git a/src/test_rig_bluesky/tomography_plans.py b/src/test_rig_bluesky/tomography_plans.py index a35526c..2fef75a 100644 --- a/src/test_rig_bluesky/tomography_plans.py +++ b/src/test_rig_bluesky/tomography_plans.py @@ -2,7 +2,6 @@ import logging from enum import StrEnum -from itertools import chain from typing import Annotated, Any from uuid import uuid4 @@ -34,17 +33,22 @@ class LightSource(StrEnum): - """The illumination in use.""" + """Deprecated: only kept so the demo_tomography API doesn't change. Ignored.""" LED = "led" SR = "sr" DARK = "dark" -# Metadata keys the reconstruction workflow reads to pair projections with the -# calibration images taken under the same illumination -LIGHT_SOURCE_KEY = "light_source" -# "dark", "flat" or "projection" +class ImageType(StrEnum): + """What a run of images is, so the reconstruction workflow can tell them apart.""" + + DARK = "dark" + FLAT = "flat" + PROJECTION = "projection" + + +# Metadata key holding the run's ImageType IMAGE_KEY = "image_key" # Shared by the darks, flats and projections runs of one tomography scan GROUP_KEY = "tomography_group" @@ -56,6 +60,14 @@ class LightSource(StrEnum): SAMPLE_OUT_X = 0.0 # TODO: placeholder, sample out of the field of view for flats SAMPLE_IN_X = 0.0 # TODO: placeholder, sample back in the field of view +# Where the sample sits and whether the tomography light is on for each image type. +# The spectroscopy mirror is always flipped out of the way. +SAMPLE_X_AND_LIGHT: dict[ImageType, tuple[float, bool]] = { + ImageType.DARK: (SAMPLE_OUT_X, False), + ImageType.FLAT: (SAMPLE_OUT_X, True), + ImageType.PROJECTION: (SAMPLE_IN_X, True), +} + MIN_PROJECTIONS = 30 MAX_PROJECTIONS = 1440 @@ -70,8 +82,8 @@ class LightSource(StrEnum): # TODO: conservative estimates for now -MAKO_ACQUIRE_PERIOD_PAD = 25e-3 -MAKO_DETECTOR_DEADTIME = 25e-3 +ALVIUM_ACQUIRE_PERIOD_PAD = 25e-3 +ALVIUM_DETECTOR_DEADTIME = 25e-3 DETECTOR_TRIGGER = DetectorTrigger.EXTERNAL_EDGE @@ -92,7 +104,67 @@ class LightSource(StrEnum): PANDA_WHITELIST = ["incenc-8-val_dataset"] -def _darks_flats_projections( +def _setup_detector(detector: AravisDetector, exposure_time: float) -> MsgGenerator: + """Put the detector in its scanning state.""" + # ---- load baseline settings ---------------------------------------------- + yield from load_settings( + device=detector, + design_name="tomography_detector_baseline", + whitelist_pvs=DETECTOR_WHITELIST, + ) + + # ---- set exposure -------------------------------------------------------- + # mv rather than prepare, because prepare cannot be used outside a run. + # See: https://github.com/DiamondLightSource/blueapi/issues/1211 + yield from bps.mv( + *(detector.driver.acquire_time, exposure_time), + *( + detector.driver.acquire_period, + exposure_time + ALVIUM_ACQUIRE_PERIOD_PAD, + ), + group="detector_acquire", + ) + + +def _move_for( + image_type: ImageType, + stage: XThetaStage, + tomography_flip_mirror: FlipMirror, + spectroscopy_flip_mirror: FlipMirror, +) -> MsgGenerator: + """Move the sample and mirrors into place for this type of image.""" + sample_x, light_on = SAMPLE_X_AND_LIGHT[image_type] + yield from bps.mv( + *(stage.x, sample_x), + *(spectroscopy_flip_mirror, False), + *(tomography_flip_mirror, light_on), + group=f"move_for_{image_type.value}", + ) + + +def _collect_calibration( + image_type: ImageType, + num_images: int, + detector: AravisDetector, + stage: XThetaStage, + tomography_flip_mirror: FlipMirror, + spectroscopy_flip_mirror: FlipMirror, + metadata: dict[str, Any], +) -> MsgGenerator[str]: + """Collect dark or flat images, returning the run's uid.""" + LOGGER.info("Collecting %d %s images.", num_images, image_type.value) + yield from _move_for( + image_type, stage, tomography_flip_mirror, spectroscopy_flip_mirror + ) + return ( + yield from count( + [detector], num=num_images, md={IMAGE_KEY: image_type.value, **metadata} + ) + ) + + +def _collect_projections( + spec: Spec[Movable], detector: AravisDetector, stage: XThetaStage, tomography_flip_mirror: FlipMirror, @@ -100,39 +172,23 @@ def _darks_flats_projections( pmac: PmacIO, pandabrick: HDFPanda, exposure_time: float, - spec: Spec[Movable], - num_darks: int, - num_flats: int, metadata: dict[str, Any], scan_type: ScanType, -) -> MsgGenerator[None]: - """Take darks, then flats with the sample out of the way, then projections.""" - LOGGER.info("Collecting %d dark images.", num_darks) - yield from bps.mv( - *(stage.x, SAMPLE_OUT_X), - *(spectroscopy_flip_mirror, False), - *(tomography_flip_mirror, False), - group="darks", - ) - darks_uid = yield from count( - [detector], num=num_darks, md={IMAGE_KEY: "dark", **metadata} - ) - - LOGGER.info("Collecting %d flat images.", num_flats) - yield from bps.mv(tomography_flip_mirror, True, group="flats") - flats_uid = yield from count( - [detector], num=num_flats, md={IMAGE_KEY: "flat", **metadata} +) -> MsgGenerator: + """Collect projections over spec, by fly or step scan.""" + # ---- move sample in, light on -------------------------------------------- + yield from _move_for( + ImageType.PROJECTION, stage, tomography_flip_mirror, spectroscopy_flip_mirror ) - yield from bps.mv(stage.x, SAMPLE_IN_X, group="projections") - projections_metadata = { - IMAGE_KEY: "projection", - DARKS_UID_KEY: darks_uid, - FLATS_UID_KEY: flats_uid, + # ---- build metadata ------------------------------------------------------ + metadata = { + IMAGE_KEY: ImageType.PROJECTION.value, "spec": serialize_spec(spec), **metadata, } + # ---- scan ---------------------------------------------------------------- if scan_type == ScanType.FLY: LOGGER.info("Performing a tomography fly scan.") yield from fly_scan( @@ -141,23 +197,68 @@ def _darks_flats_projections( pmac=pmac, pandabrick=pandabrick, exposure_time=exposure_time, - detector_deadtime=MAKO_DETECTOR_DEADTIME, + detector_deadtime=ALVIUM_DETECTOR_DEADTIME, detector_trigger=DETECTOR_TRIGGER, panda_whitelist=PANDA_WHITELIST, plan_name="tomography_fly_scan", - metadata=projections_metadata, + metadata=metadata, ) elif scan_type == ScanType.STEP: LOGGER.info("Performing a tomography step scan.") yield from spec_scan( {detector, stage}, spec, # type: ignore - metadata=projections_metadata, + metadata=metadata, ) -def _restore_mirrors(positions: list[tuple[FlipMirror, bool]]) -> MsgGenerator[None]: - yield from bps.mv(*chain.from_iterable(positions)) +def _collect_images( + spec: Spec[Movable], + detector: AravisDetector, + stage: XThetaStage, + tomography_flip_mirror: FlipMirror, + spectroscopy_flip_mirror: FlipMirror, + pmac: PmacIO, + pandabrick: HDFPanda, + exposure_time: float, + num_darks: int, + num_flats: int, + metadata: dict[str, Any], + scan_type: ScanType, +) -> MsgGenerator: + """Darks, then flats, then projections linked to both.""" + mirrors = (tomography_flip_mirror, spectroscopy_flip_mirror) + + # ---- collect darks ------------------------------------------------------- + darks_uid = yield from _collect_calibration( + ImageType.DARK, num_darks, detector, stage, *mirrors, metadata + ) + + # ---- collect flats ------------------------------------------------------- + flats_uid = yield from _collect_calibration( + ImageType.FLAT, num_flats, detector, stage, *mirrors, metadata + ) + + # ---- collect projections ------------------------------------------------- + yield from _collect_projections( + spec, + detector, + stage, + *mirrors, + pmac, + pandabrick, + exposure_time, + {DARKS_UID_KEY: darks_uid, FLATS_UID_KEY: flats_uid, **metadata}, + scan_type, + ) + + +def _restoring_mirrors(plan: MsgGenerator, *mirrors: FlipMirror) -> MsgGenerator: + """Run plan, then put the mirrors back where they were, even if it fails.""" + initial_positions = [] + for mirror in mirrors: + initial_positions += [mirror, (yield from bps.rd(mirror.position))] + return (yield from bpp.finalize_wrapper(plan, bps.mv(*initial_positions))) def _wind_back(stage: XThetaStage, start_angle: float) -> MsgGenerator[None]: @@ -169,7 +270,6 @@ def _wind_back(stage: XThetaStage, start_angle: float) -> MsgGenerator[None]: def tomography( - light_source: LightSource, detector: AravisDetector = detector, stage: XThetaStage = stage, tomography_flip_mirror: FlipMirror = tomography_flip_mirror, @@ -183,58 +283,25 @@ def tomography( metadata: dict[str, Any] | None = None, *, scan_type: ScanType, -) -> MsgGenerator[None]: - """Do a tomography scan, taking darks and flats first. - - The darks, flats and projections are each their own run, sharing a - "tomography_group" id, and the projections run records the uids of the other - two. The flip mirrors are put back as they were afterwards. - - :param light_source: illumination the flats and projections are taken under. - """ - yield from load_settings( - device=detector, - design_name="tomography_detector_baseline", - whitelist_pvs=DETECTOR_WHITELIST, - ) - - # mv rather than prepare, because prepare cannot be used outside a run. - # See: https://github.com/DiamondLightSource/blueapi/issues/1211 - yield from bps.mv( - *(detector.driver.acquire_time, exposure_time), - *( - detector.driver.acquire_period, - exposure_time + MAKO_ACQUIRE_PERIOD_PAD, - ), - group="detector_acquire", - ) - +) -> MsgGenerator: + """Do a tomography scan: darks, then flats with the sample out, then projections.""" + # ---- set up detector and stage ------------------------------------------- + yield from _setup_detector(detector, exposure_time) yield from load_settings( device=stage, design_name="tomography_stage_baseline", whitelist_pvs=STAGE_WHITELIST, ) + # ---- define scan --------------------------------------------------------- # One projection per degree over a full rotation spec = spec or Line(stage.theta, 0, 359, 360) # type: ignore + shared_metadata = {GROUP_KEY: str(uuid4()), **(metadata or {})} - # Tells the workflow which calibration images belong with these projections. - shared_metadata: dict[str, Any] = { - GROUP_KEY: str(uuid4()), - LIGHT_SOURCE_KEY: light_source.value, - **(metadata or {}), - } - - initial_mirror_positions = [ - (tomography_flip_mirror, (yield from bps.rd(tomography_flip_mirror.position))), - ( - spectroscopy_flip_mirror, - (yield from bps.rd(spectroscopy_flip_mirror.position)), - ), - ] - - yield from bpp.finalize_wrapper( - _darks_flats_projections( + # ---- collect darks, flats and projections, then restore mirrors ---------- + yield from _restoring_mirrors( + _collect_images( + spec, detector, stage, tomography_flip_mirror, @@ -242,18 +309,18 @@ def tomography( pmac, pandabrick, exposure_time, - spec, num_darks, num_flats, shared_metadata, scan_type, ), - _restore_mirrors(initial_mirror_positions), + tomography_flip_mirror, + spectroscopy_flip_mirror, ) def demo_tomography( - light_source: LightSource, + light_source: LightSource, # deprecated detector: AravisDetector = detector, stage: XThetaStage = stage, tomography_flip_mirror: FlipMirror = tomography_flip_mirror, @@ -270,19 +337,22 @@ def demo_tomography( ) -> MsgGenerator[None]: """Tomography plan intended for use in Visr demonstrations to visitors. Takes darks and flats, then number_of_projections evenly spaced over - angular_range from start_angle, then winds the stage back to start_angle. - """ + angular_range from start_angle, then winds the stage back to start_angle.""" + del light_source + + # ---- define projections -------------------------------------------------- angular_step = angular_range / number_of_projections last_angle = start_angle + angular_range - angular_step spec = Line(stage.theta, start_angle, last_angle, number_of_projections) # type: ignore + # ---- move to start ------------------------------------------------------- # PMAC handles the run-up, but move to the nominal start first so a failure # here is reported before the trajectory is built. yield from bps.mv(stage.theta, start_angle, group="initial_move") + # ---- run tomography, then wind back -------------------------------------- yield from bpp.finalize_wrapper( tomography( - light_source, detector=detector, stage=stage, tomography_flip_mirror=tomography_flip_mirror, diff --git a/tests/unit_tests/test_plans.py b/tests/unit_tests/test_plans.py index 15456d8..234c90b 100644 --- a/tests/unit_tests/test_plans.py +++ b/tests/unit_tests/test_plans.py @@ -40,7 +40,6 @@ from test_rig_bluesky.tomography_plans import ( SAMPLE_IN_X, SAMPLE_OUT_X, - LightSource, tomography, ) @@ -520,7 +519,6 @@ def test_tomography_takes_darks_then_flats_then_projections( run_engine( tomography( - LightSource.LED, detector=tomography_detector, stage=tomography_stage, tomography_flip_mirror=tomography_flip_mirror, @@ -536,7 +534,6 @@ def test_tomography_takes_darks_then_flats_then_projections( assert [start["image_key"] for start in starts] == ["dark", "flat", "projection"] assert [start["num_points"] for start in starts] == [2, 3, 4] assert len({start["tomography_group"] for start in starts}) == 1 - assert all(start["light_source"] == "led" for start in starts) assert projections["darks_uid"] == darks["uid"] assert projections["flats_uid"] == flats["uid"] @@ -552,13 +549,20 @@ def test_tomography_takes_darks_then_flats_then_projections( ) ] assert sets_and_opens == [ + # Darks: sample out, no light ("set", tomography_stage.x, (SAMPLE_OUT_X,)), ("set", spectroscopy_flip_mirror, (False,)), ("set", tomography_flip_mirror, (False,)), ("open_run", None, ()), + # Flats: sample out, light on + ("set", tomography_stage.x, (SAMPLE_OUT_X,)), + ("set", spectroscopy_flip_mirror, (False,)), ("set", tomography_flip_mirror, (True,)), ("open_run", None, ()), + # Projections: sample in, light on ("set", tomography_stage.x, (SAMPLE_IN_X,)), + ("set", spectroscopy_flip_mirror, (False,)), + ("set", tomography_flip_mirror, (True,)), ("open_run", None, ()), # Mirrors put back as they were ("set", tomography_flip_mirror, (True,)), @@ -579,7 +583,6 @@ async def test_tomography_restores_mirrors_on_failure( with pytest.raises(RuntimeError, match="boom"): run_engine( tomography( - LightSource.LED, detector=tomography_detector, stage=tomography_stage, tomography_flip_mirror=tomography_flip_mirror, From 53947a6a31c9684d0e5bc98d9ef0cc6d47bef496 Mon Sep 17 00:00:00 2001 From: Tom Kane Date: Wed, 4 Nov 2026 14:30:47 +0000 Subject: [PATCH 04/12] simplify by removing unnecessary restoring of states --- src/test_rig_bluesky/spectroscopy_plans.py | 2 + src/test_rig_bluesky/tomography_plans.py | 124 ++++++++------------- tests/unit_tests/test_plans.py | 34 +----- 3 files changed, 51 insertions(+), 109 deletions(-) diff --git a/src/test_rig_bluesky/spectroscopy_plans.py b/src/test_rig_bluesky/spectroscopy_plans.py index b4b5fd2..9f429d5 100644 --- a/src/test_rig_bluesky/spectroscopy_plans.py +++ b/src/test_rig_bluesky/spectroscopy_plans.py @@ -194,6 +194,8 @@ def spectroscopy( spec, # type: ignore metadata=metadata, ) + else: + raise ValueError(f"Unsupported scan type: {scan_type}") def demo_spectroscopy( diff --git a/src/test_rig_bluesky/tomography_plans.py b/src/test_rig_bluesky/tomography_plans.py index 2fef75a..4887878 100644 --- a/src/test_rig_bluesky/tomography_plans.py +++ b/src/test_rig_bluesky/tomography_plans.py @@ -104,7 +104,9 @@ class ImageType(StrEnum): PANDA_WHITELIST = ["incenc-8-val_dataset"] -def _setup_detector(detector: AravisDetector, exposure_time: float) -> MsgGenerator: +def _setup_detector( + detector: AravisDetector, exposure_time: float +) -> MsgGenerator[None]: """Put the detector in its scanning state.""" # ---- load baseline settings ---------------------------------------------- yield from load_settings( @@ -131,7 +133,7 @@ def _move_for( stage: XThetaStage, tomography_flip_mirror: FlipMirror, spectroscopy_flip_mirror: FlipMirror, -) -> MsgGenerator: +) -> MsgGenerator[None]: """Move the sample and mirrors into place for this type of image.""" sample_x, light_on = SAMPLE_X_AND_LIGHT[image_type] yield from bps.mv( @@ -158,7 +160,7 @@ def _collect_calibration( ) return ( yield from count( - [detector], num=num_images, md={IMAGE_KEY: image_type.value, **metadata} + [detector], num=num_images, md={**metadata, IMAGE_KEY: image_type.value} ) ) @@ -174,7 +176,7 @@ def _collect_projections( exposure_time: float, metadata: dict[str, Any], scan_type: ScanType, -) -> MsgGenerator: +) -> MsgGenerator[None]: """Collect projections over spec, by fly or step scan.""" # ---- move sample in, light on -------------------------------------------- yield from _move_for( @@ -183,13 +185,13 @@ def _collect_projections( # ---- build metadata ------------------------------------------------------ metadata = { + **metadata, IMAGE_KEY: ImageType.PROJECTION.value, "spec": serialize_spec(spec), - **metadata, } - # ---- scan ---------------------------------------------------------------- if scan_type == ScanType.FLY: + # ---- fly scan -------------------------------------------------------- LOGGER.info("Performing a tomography fly scan.") yield from fly_scan( spec, @@ -204,61 +206,16 @@ def _collect_projections( metadata=metadata, ) elif scan_type == ScanType.STEP: + # ---- step scan ------------------------------------------------------- LOGGER.info("Performing a tomography step scan.") yield from spec_scan( {detector, stage}, spec, # type: ignore metadata=metadata, ) - - -def _collect_images( - spec: Spec[Movable], - detector: AravisDetector, - stage: XThetaStage, - tomography_flip_mirror: FlipMirror, - spectroscopy_flip_mirror: FlipMirror, - pmac: PmacIO, - pandabrick: HDFPanda, - exposure_time: float, - num_darks: int, - num_flats: int, - metadata: dict[str, Any], - scan_type: ScanType, -) -> MsgGenerator: - """Darks, then flats, then projections linked to both.""" - mirrors = (tomography_flip_mirror, spectroscopy_flip_mirror) - - # ---- collect darks ------------------------------------------------------- - darks_uid = yield from _collect_calibration( - ImageType.DARK, num_darks, detector, stage, *mirrors, metadata - ) - - # ---- collect flats ------------------------------------------------------- - flats_uid = yield from _collect_calibration( - ImageType.FLAT, num_flats, detector, stage, *mirrors, metadata - ) - - # ---- collect projections ------------------------------------------------- - yield from _collect_projections( - spec, - detector, - stage, - *mirrors, - pmac, - pandabrick, - exposure_time, - {DARKS_UID_KEY: darks_uid, FLATS_UID_KEY: flats_uid, **metadata}, - scan_type, - ) - - -def _restoring_mirrors(plan: MsgGenerator, *mirrors: FlipMirror) -> MsgGenerator: - """Run plan, then put the mirrors back where they were, even if it fails.""" - initial_positions = [] - for mirror in mirrors: - initial_positions += [mirror, (yield from bps.rd(mirror.position))] - return (yield from bpp.finalize_wrapper(plan, bps.mv(*initial_positions))) + else: + # ---- unsupported scan type ------------------------------------------- + raise ValueError(f"Unsupported scan type: {scan_type}") def _wind_back(stage: XThetaStage, start_angle: float) -> MsgGenerator[None]: @@ -278,12 +235,11 @@ def tomography( pandabrick: HDFPanda = pandabrick, exposure_time: float = 0.1, spec: Spec[Movable] | None = None, - num_darks: int = 20, - num_flats: int = 20, + num_calibration_images: int = 20, metadata: dict[str, Any] | None = None, *, scan_type: ScanType, -) -> MsgGenerator: +) -> MsgGenerator[None]: """Do a tomography scan: darks, then flats with the sample out, then projections.""" # ---- set up detector and stage ------------------------------------------- yield from _setup_detector(detector, exposure_time) @@ -296,26 +252,40 @@ def tomography( # ---- define scan --------------------------------------------------------- # One projection per degree over a full rotation spec = spec or Line(stage.theta, 0, 359, 360) # type: ignore - shared_metadata = {GROUP_KEY: str(uuid4()), **(metadata or {})} + shared_metadata = {**(metadata or {}), GROUP_KEY: str(uuid4())} + mirrors = (tomography_flip_mirror, spectroscopy_flip_mirror) - # ---- collect darks, flats and projections, then restore mirrors ---------- - yield from _restoring_mirrors( - _collect_images( - spec, - detector, - stage, - tomography_flip_mirror, - spectroscopy_flip_mirror, - pmac, - pandabrick, - exposure_time, - num_darks, - num_flats, - shared_metadata, - scan_type, - ), - tomography_flip_mirror, - spectroscopy_flip_mirror, + # ---- collect darks ------------------------------------------------------- + darks_uid = yield from _collect_calibration( + ImageType.DARK, + num_calibration_images, + detector, + stage, + *mirrors, + shared_metadata, + ) + + # ---- collect flats ------------------------------------------------------- + flats_uid = yield from _collect_calibration( + ImageType.FLAT, + num_calibration_images, + detector, + stage, + *mirrors, + shared_metadata, + ) + + # ---- collect projections ------------------------------------------------- + yield from _collect_projections( + spec, + detector, + stage, + *mirrors, + pmac, + pandabrick, + exposure_time, + {**shared_metadata, DARKS_UID_KEY: darks_uid, FLATS_UID_KEY: flats_uid}, + scan_type, ) diff --git a/tests/unit_tests/test_plans.py b/tests/unit_tests/test_plans.py index 234c90b..07a6def 100644 --- a/tests/unit_tests/test_plans.py +++ b/tests/unit_tests/test_plans.py @@ -524,15 +524,14 @@ def test_tomography_takes_darks_then_flats_then_projections( tomography_flip_mirror=tomography_flip_mirror, spectroscopy_flip_mirror=spectroscopy_flip_mirror, spec=Line(tomography_stage.theta, 0, 3, 4), - num_darks=2, - num_flats=3, + num_calibration_images=2, scan_type=ScanType.STEP, ) ) darks, flats, projections = starts assert [start["image_key"] for start in starts] == ["dark", "flat", "projection"] - assert [start["num_points"] for start in starts] == [2, 3, 4] + assert [start["num_points"] for start in starts] == [2, 2, 4] assert len({start["tomography_group"] for start in starts}) == 1 assert projections["darks_uid"] == darks["uid"] assert projections["flats_uid"] == flats["uid"] @@ -564,33 +563,4 @@ def test_tomography_takes_darks_then_flats_then_projections( ("set", spectroscopy_flip_mirror, (False,)), ("set", tomography_flip_mirror, (True,)), ("open_run", None, ()), - # Mirrors put back as they were - ("set", tomography_flip_mirror, (True,)), - ("set", spectroscopy_flip_mirror, (True,)), ] - - -async def test_tomography_restores_mirrors_on_failure( - run_engine: RunEngine, - tomography_detector: AravisDetector, - tomography_stage: XThetaStage, - tomography_flip_mirror: FlipMirror, - spectroscopy_flip_mirror: FlipMirror, -): - with patch( - "test_rig_bluesky.tomography_plans.count", side_effect=RuntimeError("boom") - ): - with pytest.raises(RuntimeError, match="boom"): - run_engine( - tomography( - detector=tomography_detector, - stage=tomography_stage, - tomography_flip_mirror=tomography_flip_mirror, - spectroscopy_flip_mirror=spectroscopy_flip_mirror, - spec=Line(tomography_stage.theta, 0, 3, 4), - scan_type=ScanType.STEP, - ) - ) - - assert await tomography_flip_mirror.position.get_value() is True - assert await spectroscopy_flip_mirror.position.get_value() is True From 339625a113cc84d3b0810f5342b4eae6b1b7d81f Mon Sep 17 00:00:00 2001 From: Tom Kane Date: Mon, 9 Nov 2026 12:54:34 +0000 Subject: [PATCH 05/12] remove light source entirely --- src/test_rig_bluesky/tomography_plans.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/src/test_rig_bluesky/tomography_plans.py b/src/test_rig_bluesky/tomography_plans.py index 4887878..dd2a43f 100644 --- a/src/test_rig_bluesky/tomography_plans.py +++ b/src/test_rig_bluesky/tomography_plans.py @@ -32,14 +32,6 @@ spectroscopy_flip_mirror = inject("spectroscopy_flip_mirror") -class LightSource(StrEnum): - """Deprecated: only kept so the demo_tomography API doesn't change. Ignored.""" - - LED = "led" - SR = "sr" - DARK = "dark" - - class ImageType(StrEnum): """What a run of images is, so the reconstruction workflow can tell them apart.""" @@ -290,7 +282,6 @@ def tomography( def demo_tomography( - light_source: LightSource, # deprecated detector: AravisDetector = detector, stage: XThetaStage = stage, tomography_flip_mirror: FlipMirror = tomography_flip_mirror, @@ -308,7 +299,6 @@ def demo_tomography( """Tomography plan intended for use in Visr demonstrations to visitors. Takes darks and flats, then number_of_projections evenly spaced over angular_range from start_angle, then winds the stage back to start_angle.""" - del light_source # ---- define projections -------------------------------------------------- angular_step = angular_range / number_of_projections From 7f72693b485d2556fbf4ce476592b2a0302d9823 Mon Sep 17 00:00:00 2001 From: Tom Kane Date: Tue, 3 Nov 2026 11:50:47 +0000 Subject: [PATCH 06/12] remove redundant terms from function arguments, update with changes from dodal (XThetaStage for tmography and sample stage rename --- ....yaml => spectroscopy_stage_baseline.yaml} | 0 src/test_rig_bluesky/plans.py | 2 +- src/test_rig_bluesky/spectroscopy_plans.py | 66 ++++++++-------- src/test_rig_bluesky/tomography_plans.py | 50 ++++++------ tests/system_tests/test_plans_system.py | 20 +++-- tests/unit_tests/test_plans.py | 77 ++++++++++--------- 6 files changed, 107 insertions(+), 108 deletions(-) rename src/test_rig_bluesky/baselines/{sample_stage_baseline.yaml => spectroscopy_stage_baseline.yaml} (100%) diff --git a/src/test_rig_bluesky/baselines/sample_stage_baseline.yaml b/src/test_rig_bluesky/baselines/spectroscopy_stage_baseline.yaml similarity index 100% rename from src/test_rig_bluesky/baselines/sample_stage_baseline.yaml rename to src/test_rig_bluesky/baselines/spectroscopy_stage_baseline.yaml diff --git a/src/test_rig_bluesky/plans.py b/src/test_rig_bluesky/plans.py index ddd2474..d05972b 100644 --- a/src/test_rig_bluesky/plans.py +++ b/src/test_rig_bluesky/plans.py @@ -106,7 +106,7 @@ def serialize_spec(spec: Spec[Any]) -> Any: """Render a Spec as something JSON-serializable, for scan metadata. Axes that are devices are written out by their ``name`` (e.g. - ``"sample_stage-x"``), so consumers of the run's start document can tell + ``"spectroscopy_stage-x"``), so consumers of the run's start document can tell which axis is which. """ return spec.serialize() diff --git a/src/test_rig_bluesky/spectroscopy_plans.py b/src/test_rig_bluesky/spectroscopy_plans.py index cd11372..7b459f9 100644 --- a/src/test_rig_bluesky/spectroscopy_plans.py +++ b/src/test_rig_bluesky/spectroscopy_plans.py @@ -35,10 +35,10 @@ LOGGER = logging.getLogger(__name__) -spectroscopy_detector = inject("spectroscopy_detector") -spectroscopy_diodes = inject("spectroscopy_diodes") +detector = inject("spectroscopy_detector") +diodes = inject("spectroscopy_diodes") pandabox = inject("pandabox") -sample_stage = inject("spectroscopy_stage") +stage = inject("spectroscopy_stage") # The encoder entries map to specific axes, so despite the PandA being shared # hardware this list is spectroscopy's. @@ -62,12 +62,12 @@ def spectroscopy( - spectroscopy_detector: AravisDetector = spectroscopy_detector, - sample_stage: XYZStage = sample_stage, + detector: AravisDetector = detector, + stage: XYZStage = stage, pmac: PmacIO = pmac, pandabrick: HDFPanda = pandabrick, pandabox: HDFPanda = pandabox, - spectroscopy_diodes: RGBDiodes = spectroscopy_diodes, + diodes: RGBDiodes = diodes, spec: Spec[Movable] | None = None, exposure_time: float = 0.1, amplifier_range: EnzLocumRange | None = None, @@ -78,19 +78,15 @@ def spectroscopy( """Do a spectroscopy scan, acquiring from the camera and the RGB diodes.""" # The range is only recorded once per scan, as configuration, so it must not # change between points - scan_range = amplifier_range or ( - yield from bps.rd(spectroscopy_diodes.amplifier.range) - ) + scan_range = amplifier_range or (yield from bps.rd(diodes.amplifier.range)) if scan_range == EnzLocumRange.AUTO: raise ValueError("The diode amplifier must be on a fixed range, not Auto") if amplifier_range is not None: - yield from bps.abs_set( - spectroscopy_diodes.amplifier.range, amplifier_range, wait=True - ) + yield from bps.abs_set(diodes.amplifier.range, amplifier_range, wait=True) yield from load_settings( - device=spectroscopy_detector, + device=detector, design_name="spectroscopy_detector_baseline", whitelist_pvs=[ "hdf-nd_array_port", @@ -126,17 +122,17 @@ def spectroscopy( # outside of a run. # See: https://github.com/DiamondLightSource/blueapi/issues/1211 yield from bps.mv( - *(spectroscopy_detector.driver.acquire_time, exposure_time), + *(detector.driver.acquire_time, exposure_time), *( - spectroscopy_detector.driver.acquire_period, + detector.driver.acquire_period, exposure_time + MANTA_ACQUIRE_PERIOD_PAD, ), - group="spectroscopy_detector_aquire", + group="detector_aquire", ) params: list[NDAttributeParam] = [] - for channel in list(spectroscopy_detector.roistat.channels.keys()): # type: ignore - roistatn = spectroscopy_detector.roistat.channels[channel] # type: ignore + for channel in list(detector.roistat.channels.keys()): # type: ignore + roistatn = detector.roistat.channels[channel] # type: ignore assert isinstance(roistatn, NDROIStatNIO) channel_name = yield from bps.rd(roistatn.name_) @@ -151,11 +147,11 @@ def spectroscopy( ) ) - yield from setup_ndattributes(spectroscopy_detector.roistat, params) # type: ignore + yield from setup_ndattributes(detector.roistat, params) # type: ignore yield from load_settings( - device=sample_stage, - design_name="sample_stage_baseline", + device=stage, + design_name="spectroscopy_stage_baseline", whitelist_pvs=[ "x-acceleration_time", "x-velocity", @@ -164,7 +160,7 @@ def spectroscopy( ], ) - spec = spec or Line(sample_stage.x, 0, 5, 5) # type: ignore + spec = spec or Line(stage.x, 0, 5, 5) # type: ignore # write json-serializable form of spec into metadata if metadata is None: @@ -181,7 +177,7 @@ def spectroscopy( ) yield from fly_scan( spec, - spectroscopy_detector, + detector, pmac=pmac, pandabrick=pandabrick, exposure_time=exposure_time, @@ -189,24 +185,24 @@ def spectroscopy( pandabrick_whitelist=PANDA_WHITELIST, metadata=metadata, pandabox=pandabox, - diodes=spectroscopy_diodes, + diodes=diodes, ) elif scan_type == ScanType.STEP: LOGGER.info("Performing a step scan.") yield from spec_scan( - {spectroscopy_detector, spectroscopy_diodes, sample_stage}, + {detector, diodes, stage}, spec, # type: ignore metadata=metadata, ) def demo_spectroscopy( - spectroscopy_detector: AravisDetector = spectroscopy_detector, - sample_stage: XYZStage = sample_stage, + detector: AravisDetector = detector, + stage: XYZStage = stage, pmac: PmacIO = pmac, pandabrick: HDFPanda = pandabrick, pandabox: HDFPanda = pandabox, - spectroscopy_diodes: RGBDiodes = spectroscopy_diodes, + diodes: RGBDiodes = diodes, total_number_of_scan_points: int = 25, grid_size: float = 5.0, grid_origin_x: float = 0.0, @@ -239,23 +235,21 @@ def demo_spectroscopy( # Move to the start point # pmac should handle premove - yield from bps.mv( - *(sample_stage.x, xmin), *(sample_stage.y, ymin), group="initial_move" - ) + yield from bps.mv(*(stage.x, xmin), *(stage.y, ymin), group="initial_move") - grid = Line(sample_stage.y, ymin, ymax, ysteps) * ~Line( # type: ignore - sample_stage.x, # type: ignore + grid = Line(stage.y, ymin, ymax, ysteps) * ~Line( # type: ignore + stage.x, # type: ignore xmin, xmax, xsteps, ) yield from spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=detector, + stage=stage, pmac=pmac, pandabrick=pandabrick, pandabox=pandabox, - spectroscopy_diodes=spectroscopy_diodes, + diodes=diodes, spec=grid, exposure_time=exposure_time, metadata=metadata, diff --git a/src/test_rig_bluesky/tomography_plans.py b/src/test_rig_bluesky/tomography_plans.py index 4c74a96..5fb79ad 100644 --- a/src/test_rig_bluesky/tomography_plans.py +++ b/src/test_rig_bluesky/tomography_plans.py @@ -10,10 +10,10 @@ from bluesky.protocols import Movable from bluesky.utils import MsgGenerator from dodal.common import inject +from dodal.devices.motors import XThetaStage from dodal.plans import spec_scan from ophyd_async.core import DetectorTrigger from ophyd_async.epics.adaravis import AravisDetector -from ophyd_async.epics.motor import Motor from ophyd_async.epics.pmac import PmacIO from ophyd_async.fastcs.panda import HDFPanda from pydantic import Field @@ -24,8 +24,8 @@ LOGGER = logging.getLogger(__name__) # pmac and pandabrick come from plans.py -tomography_detector = inject("tomography_detector") -tomography_stage = inject("tomography_stage") +detector = inject("tomography_detector") +stage = inject("tomography_stage") class LightSource(StrEnum): @@ -57,7 +57,7 @@ class LightSource(StrEnum): MAKO_ACQUIRE_PERIOD_PAD = 25e-3 MAKO_DETECTOR_DEADTIME = 25e-3 -TOMOGRAPHY_DETECTOR_TRIGGER = DetectorTrigger.EXTERNAL_EDGE +DETECTOR_TRIGGER = DetectorTrigger.EXTERNAL_EDGE DETECTOR_WHITELIST = [ "hdf-nd_array_port", @@ -67,7 +67,7 @@ class LightSource(StrEnum): "driver-trigger_source", ] -TOMOGRAPHY_STAGE_WHITELIST = [ +STAGE_WHITELIST = [ "acceleration_time", "velocity", ] @@ -95,14 +95,14 @@ def _setup_detector( detector.driver.acquire_period, exposure_time + MAKO_ACQUIRE_PERIOD_PAD, ), - group="tomography_detector_acquire", + group="detector_acquire", ) # @attach_data_session_metadata_decorator() def darks_flats( light_source: LightSource, - tomography_detector: AravisDetector = tomography_detector, + detector: AravisDetector = detector, num_images: int = 20, exposure_time: float = 0.1, metadata: dict[str, Any] | None = None, @@ -128,15 +128,15 @@ def darks_flats( light_source.value, ) - yield from _setup_detector(tomography_detector, exposure_time) + yield from _setup_detector(detector, exposure_time) - yield from count([tomography_detector], num=num_images, md=run_metadata) + yield from count([detector], num=num_images, md=run_metadata) def tomography( light_source: LightSource, - tomography_detector: AravisDetector = tomography_detector, - tomography_stage: Motor = tomography_stage, + detector: AravisDetector = detector, + stage: XThetaStage = stage, pmac: PmacIO = pmac, pandabrick: HDFPanda = pandabrick, exposure_time: float = 0.1, @@ -146,16 +146,16 @@ def tomography( scan_type: ScanType, ) -> MsgGenerator[None]: """Do a tomography scan.""" - yield from _setup_detector(tomography_detector, exposure_time) + yield from _setup_detector(detector, exposure_time) yield from load_settings( - device=tomography_stage, + device=stage, design_name="tomography_stage_baseline", - whitelist_pvs=TOMOGRAPHY_STAGE_WHITELIST, + whitelist_pvs=STAGE_WHITELIST, ) # One projection per degree over a full rotation - spec = spec or Line(tomography_stage, 0, 359, 360) # type: ignore + spec = spec or Line(stage, 0, 359, 360) # type: ignore if metadata is None: metadata = {} @@ -167,12 +167,12 @@ def tomography( LOGGER.info("Performing a tomography fly scan.") yield from fly_scan( spec, - tomography_detector, + detector, pmac=pmac, pandabrick=pandabrick, exposure_time=exposure_time, detector_deadtime=MAKO_DETECTOR_DEADTIME, - detector_trigger=TOMOGRAPHY_DETECTOR_TRIGGER, + detector_trigger=DETECTOR_TRIGGER, pandabrick_whitelist=PANDA_WHITELIST, plan_name="tomography_fly_scan", metadata=metadata, @@ -180,7 +180,7 @@ def tomography( elif scan_type == ScanType.STEP: LOGGER.info("Performing a tomography step scan.") yield from spec_scan( - {tomography_detector, tomography_stage}, + {detector, stage}, spec, # type: ignore metadata=metadata, ) @@ -188,8 +188,8 @@ def tomography( def demo_tomography( light_source: LightSource, - tomography_detector: AravisDetector = tomography_detector, - tomography_stage: Motor = tomography_stage, + detector: AravisDetector = detector, + stage: XThetaStage = stage, pmac: PmacIO = pmac, pandabrick: HDFPanda = pandabrick, number_of_projections: NumProjections = 360, @@ -206,24 +206,24 @@ def demo_tomography( """ angular_step = angular_range / number_of_projections last_angle = start_angle + angular_range - angular_step - spec = Line(tomography_stage, start_angle, last_angle, number_of_projections) # type: ignore + spec = Line(stage, start_angle, last_angle, number_of_projections) # type: ignore # PMAC handles the run-up, but move to the nominal start first so a failure # here is reported before the trajectory is built. - yield from bps.mv(tomography_stage, start_angle, group="initial_move") + yield from bps.mv(stage, start_angle, group="initial_move") def wind_back() -> MsgGenerator[None]: # unwind to leave the stage where the next scan expects it. # currently this winds back the full rotation but should look # into changing it so that we just advance one step to the # beginning again and reset the angle to 0? - yield from bps.mv(tomography_stage, start_angle, group="wind_back") + yield from bps.mv(stage, start_angle, group="wind_back") yield from bpp.finalize_wrapper( tomography( light_source, - tomography_detector=tomography_detector, - tomography_stage=tomography_stage, + detector=detector, + stage=stage, pmac=pmac, pandabrick=pandabrick, exposure_time=exposure_time, diff --git a/tests/system_tests/test_plans_system.py b/tests/system_tests/test_plans_system.py index f0417dc..dd9711d 100644 --- a/tests/system_tests/test_plans_system.py +++ b/tests/system_tests/test_plans_system.py @@ -117,14 +117,16 @@ def test_generic_scan( def test_spectroscopy_re(): run_engine = RunEngine() spectroscopy_detector = b01_1.spectroscopy_detector.build(connect_immediately=True) - sample_stage = b01_1.spectroscopy_stage.build(connect_immediately=True) + spectroscopy_stage = b01_1.spectroscopy_stage.build(connect_immediately=True) - scan_spec = Line(sample_stage.y, 0, 5, 50) * Line(sample_stage.x, 2, 5, 30) + scan_spec = Line(spectroscopy_stage.y, 0, 5, 50) * Line( + spectroscopy_stage.x, 2, 5, 30 + ) run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, spec=scan_spec, scan_type=ScanType.FLY, ) @@ -168,14 +170,16 @@ def test_spectroscopy_re_run_engine(): run_engine = RunEngine() spectroscopy_detector = b01_1.spectroscopy_detector.build(connect_immediately=True) - sample_stage = b01_1.spectroscopy_stage.build(connect_immediately=True) + spectroscopy_stage = b01_1.spectroscopy_stage.build(connect_immediately=True) - scan_spec = Line(sample_stage.y, 0, 5, 50) * Line(sample_stage.x, 2, 5, 30) + scan_spec = Line(spectroscopy_stage.y, 0, 5, 50) * Line( + spectroscopy_stage.x, 2, 5, 30 + ) run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, spec=scan_spec, scan_type=ScanType.STEP, ) diff --git a/tests/unit_tests/test_plans.py b/tests/unit_tests/test_plans.py index 36260d8..37df9d9 100644 --- a/tests/unit_tests/test_plans.py +++ b/tests/unit_tests/test_plans.py @@ -46,7 +46,7 @@ def spectroscopy_detector(run_engine: RunEngine) -> AravisDetector: @pytest.fixture -def sample_stage(run_engine: RunEngine) -> XYZStage: +def spectroscopy_stage(run_engine: RunEngine) -> XYZStage: stage = b01_1.spectroscopy_stage.build(connect_immediately=True, mock=True) set_mock_value(stage.x.low_limit_travel, -10.0) @@ -143,12 +143,12 @@ async def test_load_settings( def test_snapshot( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, ): docs = defaultdict(list) run_engine.subscribe(lambda name, doc: docs[name].append(doc)) - run_engine(snapshot([spectroscopy_detector, sample_stage])) + run_engine(snapshot([spectroscopy_detector, spectroscopy_stage])) assert_emitted( docs, start=1, descriptor=1, stream_resource=1, stream_datum=1, event=1, stop=1 @@ -164,7 +164,7 @@ def test_snapshot( async def test_spectroscopy( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, pandabrick: HDFPanda, spectroscopy_diodes: RGBDiodes, ): @@ -173,11 +173,12 @@ async def test_spectroscopy( run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, pandabrick=pandabrick, - spectroscopy_diodes=spectroscopy_diodes, - spec=Line(sample_stage.y, 4.2, 6, 3) * Line(sample_stage.x, 0, 5, 10), + diodes=spectroscopy_diodes, + spec=Line(spectroscopy_stage.y, 4.2, 6, 3) + * Line(spectroscopy_stage.x, 0, 5, 10), exposure_time=0.2, scan_type=ScanType.STEP, ) @@ -199,7 +200,7 @@ async def test_spectroscopy( async def test_spectroscopy_defaults( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, pandabrick: HDFPanda, spectroscopy_diodes: RGBDiodes, ): @@ -208,10 +209,10 @@ async def test_spectroscopy_defaults( run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, pandabrick=pandabrick, - spectroscopy_diodes=spectroscopy_diodes, + diodes=spectroscopy_diodes, scan_type=ScanType.STEP, ) ) @@ -232,7 +233,7 @@ async def test_spectroscopy_defaults( def test_spectroscopy_datasets( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, pandabrick: HDFPanda, spectroscopy_diodes: RGBDiodes, ): @@ -241,10 +242,10 @@ def test_spectroscopy_datasets( run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, pandabrick=pandabrick, - spectroscopy_diodes=spectroscopy_diodes, + diodes=spectroscopy_diodes, scan_type=ScanType.STEP, ) ) @@ -264,16 +265,16 @@ def test_spectroscopy_datasets( async def test_spectroscopy_sets_exposure_time_and_acquire_period( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, pandabrick: HDFPanda, spectroscopy_diodes: RGBDiodes, ): run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, pandabrick=pandabrick, - spectroscopy_diodes=spectroscopy_diodes, + diodes=spectroscopy_diodes, exposure_time=1.0, scan_type=ScanType.STEP, ) @@ -292,8 +293,8 @@ def test_demo_spectroscopy(): ) as mock_spec: # Call the generator function and exhaust it generator = demo_spectroscopy( - spectroscopy_detector=fake_detector, - sample_stage=fake_stage, + detector=fake_detector, + stage=fake_stage, total_number_of_scan_points=25, scan_type=ScanType.FLY, ) @@ -304,31 +305,31 @@ def test_demo_spectroscopy(): mock_spec.assert_called_once() called_kwargs = mock_spec.call_args.kwargs - assert called_kwargs["spectroscopy_detector"] is fake_detector - assert called_kwargs["sample_stage"] is fake_stage + assert called_kwargs["detector"] is fake_detector + assert called_kwargs["stage"] is fake_stage assert called_kwargs["spec"] == Line(fake_stage.y, 0.0, 5.0, 5) * ~Line( fake_stage.x, 0.0, 5.0, 5 ) assert called_kwargs["scan_type"] == ScanType.FLY -def test_serialize_spec_names_device_axes(sample_stage: XYZStage): - spec = Line(sample_stage.y, 4.2, 6, 3) * Line(sample_stage.x, 0, 5, 10) +def test_serialize_spec_names_device_axes(spectroscopy_stage: XYZStage): + spec = Line(spectroscopy_stage.y, 4.2, 6, 3) * Line(spectroscopy_stage.x, 0, 5, 10) serialized = serialize_spec(spec) - assert serialized["outer"]["axis"] == sample_stage.y.name - assert serialized["inner"]["axis"] == sample_stage.x.name + assert serialized["outer"]["axis"] == spectroscopy_stage.y.name + assert serialized["inner"]["axis"] == spectroscopy_stage.x.name assert Spec.deserialize(serialized).axes() == [ - sample_stage.y.name, - sample_stage.x.name, + spectroscopy_stage.y.name, + spectroscopy_stage.x.name, ] async def test_spectroscopy_sets_amplifier_range( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, pandabrick: HDFPanda, spectroscopy_diodes: RGBDiodes, ): @@ -337,10 +338,10 @@ async def test_spectroscopy_sets_amplifier_range( run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, pandabrick=pandabrick, - spectroscopy_diodes=spectroscopy_diodes, + diodes=spectroscopy_diodes, amplifier_range=EnzLocumRange.RANGE_10NA, scan_type=ScanType.STEP, ) @@ -356,7 +357,7 @@ async def test_spectroscopy_sets_amplifier_range( def test_spectroscopy_rejects_auto_amplifier_range( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, spectroscopy_diodes: RGBDiodes, ): set_mock_value(spectroscopy_diodes.amplifier.range, EnzLocumRange.AUTO) @@ -364,9 +365,9 @@ def test_spectroscopy_rejects_auto_amplifier_range( with pytest.raises(ValueError, match="Auto"): run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, - spectroscopy_diodes=spectroscopy_diodes, + detector=spectroscopy_detector, + stage=spectroscopy_stage, + diodes=spectroscopy_diodes, scan_type=ScanType.STEP, ) ) From 5211610479109532c526902acd081c9a64c0682c Mon Sep 17 00:00:00 2001 From: Tom Kane Date: Tue, 3 Nov 2026 16:37:28 +0000 Subject: [PATCH 07/12] incorporate darks and flats inside tomography plan --- .../baselines/tomography_stage_baseline.yaml | 16 +- src/test_rig_bluesky/tomography_plans.py | 231 +++++++++++------- tests/unit_tests/test_plans.py | 154 +++++++++++- 3 files changed, 300 insertions(+), 101 deletions(-) diff --git a/src/test_rig_bluesky/baselines/tomography_stage_baseline.yaml b/src/test_rig_bluesky/baselines/tomography_stage_baseline.yaml index 6b9b570..f16bb5f 100644 --- a/src/test_rig_bluesky/baselines/tomography_stage_baseline.yaml +++ b/src/test_rig_bluesky/baselines/tomography_stage_baseline.yaml @@ -1,8 +1,8 @@ -acceleration_time: 0.001 -high_limit_travel: 1000 -low_limit_travel: -1000 -offset: 0.0 -offset_freeze_switch: Variable -set_use_switch: Use -user_setpoint: 5.00004 -velocity: 30 +theta.acceleration_time: 0.001 +theta.high_limit_travel: 1000 +theta.low_limit_travel: -1000 +theta.offset: 0.0 +theta.offset_freeze_switch: Variable +theta.set_use_switch: Use +theta.user_setpoint: 5.00004 +theta.velocity: 30 diff --git a/src/test_rig_bluesky/tomography_plans.py b/src/test_rig_bluesky/tomography_plans.py index 5fb79ad..4d81d57 100644 --- a/src/test_rig_bluesky/tomography_plans.py +++ b/src/test_rig_bluesky/tomography_plans.py @@ -2,7 +2,9 @@ import logging from enum import StrEnum +from itertools import chain from typing import Annotated, Any +from uuid import uuid4 import bluesky.plan_stubs as bps import bluesky.preprocessors as bpp @@ -10,6 +12,7 @@ from bluesky.protocols import Movable from bluesky.utils import MsgGenerator from dodal.common import inject +from dodal.devices.beamlines.b01_1.flip_mirror import FlipMirror from dodal.devices.motors import XThetaStage from dodal.plans import spec_scan from ophyd_async.core import DetectorTrigger @@ -26,6 +29,8 @@ # pmac and pandabrick come from plans.py detector = inject("tomography_detector") stage = inject("tomography_stage") +tomography_flip_mirror = inject("tomography_flip_mirror") +spectroscopy_flip_mirror = inject("spectroscopy_flip_mirror") class LightSource(StrEnum): @@ -39,6 +44,17 @@ class LightSource(StrEnum): # Metadata keys the reconstruction workflow reads to pair projections with the # calibration images taken under the same illumination LIGHT_SOURCE_KEY = "light_source" +# "dark", "flat" or "projection" +IMAGE_KEY = "image_key" +# Shared by the darks, flats and projections runs of one tomography scan +GROUP_KEY = "tomography_group" +# Recorded on the projections run, pointing at its calibration runs +DARKS_UID_KEY = "darks_uid" +FLATS_UID_KEY = "flats_uid" + +# Tomography stage x positions, in mm +SAMPLE_OUT_X = 0.0 # TODO: placeholder, sample out of the field of view for flats +SAMPLE_IN_X = 0.0 # TODO: placeholder, sample back in the field of view MIN_PROJECTIONS = 30 MAX_PROJECTIONS = 1440 @@ -68,85 +84,130 @@ class LightSource(StrEnum): ] STAGE_WHITELIST = [ - "acceleration_time", - "velocity", + "theta-acceleration_time", + "theta-velocity", ] # TODO: confirm which PandA encoder input the rotation axis is wired to. PANDA_WHITELIST = ["incenc-8-val_dataset"] -def _setup_detector( +def _darks_flats_projections( detector: AravisDetector, + stage: XThetaStage, + tomography_flip_mirror: FlipMirror, + spectroscopy_flip_mirror: FlipMirror, + pmac: PmacIO, + pandabrick: HDFPanda, exposure_time: float, + spec: Spec[Movable], + num_darks: int, + num_flats: int, + metadata: dict[str, Any], + scan_type: ScanType, ) -> MsgGenerator[None]: - """Put the detector in its scanning state.""" - yield from load_settings( - device=detector, - design_name="tomography_detector_baseline", - whitelist_pvs=DETECTOR_WHITELIST, + """Take darks, then flats with the sample out of the way, then projections.""" + LOGGER.info("Collecting %d dark images.", num_darks) + yield from bps.mv( + *(stage.x, SAMPLE_OUT_X), + *(spectroscopy_flip_mirror, False), + *(tomography_flip_mirror, False), + group="darks", + ) + darks_uid = yield from count( + [detector], num=num_darks, md={IMAGE_KEY: "dark", **metadata} ) - # mv rather than prepare, because prepare cannot be used outside a run. - # See: https://github.com/DiamondLightSource/blueapi/issues/1211 - yield from bps.mv( - *(detector.driver.acquire_time, exposure_time), - *( - detector.driver.acquire_period, - exposure_time + MAKO_ACQUIRE_PERIOD_PAD, - ), - group="detector_acquire", + LOGGER.info("Collecting %d flat images.", num_flats) + yield from bps.mv(tomography_flip_mirror, True, group="flats") + flats_uid = yield from count( + [detector], num=num_flats, md={IMAGE_KEY: "flat", **metadata} ) + yield from bps.mv(stage.x, SAMPLE_IN_X, group="projections") + projections_metadata = { + IMAGE_KEY: "projection", + DARKS_UID_KEY: darks_uid, + FLATS_UID_KEY: flats_uid, + "spec": serialize_spec(spec), + **metadata, + } -# @attach_data_session_metadata_decorator() -def darks_flats( - light_source: LightSource, - detector: AravisDetector = detector, - num_images: int = 20, - exposure_time: float = 0.1, - metadata: dict[str, Any] | None = None, -) -> MsgGenerator[None]: - """Collect flat or dark field images for the reconstruction to normalise with. + if scan_type == ScanType.FLY: + LOGGER.info("Performing a tomography fly scan.") + yield from fly_scan( + spec, + detector, + pmac=pmac, + pandabrick=pandabrick, + exposure_time=exposure_time, + detector_deadtime=MAKO_DETECTOR_DEADTIME, + detector_trigger=DETECTOR_TRIGGER, + pandabrick_whitelist=PANDA_WHITELIST, + plan_name="tomography_fly_scan", + metadata=projections_metadata, + ) + elif scan_type == ScanType.STEP: + LOGGER.info("Performing a tomography step scan.") + yield from spec_scan( + {detector, stage}, + spec, # type: ignore + metadata=projections_metadata, + ) - Recorded in the run's metadata under "light_source" so the workflow can find - the right images for a given scan. LightSource.NONE means the beam is off, - i.e. dark fields, and one such set serves every illumination. - :param light_source: illumination the images are taken under, or - LightSource.NONE for dark fields. - """ - run_metadata: dict[str, Any] = {LIGHT_SOURCE_KEY: light_source.value} - run_metadata.update(metadata or {}) - - if light_source is LightSource.DARK: - LOGGER.info("Collecting %d dark images.", num_images) - else: - LOGGER.info( - "Collecting %d flat images under %s illumination.", - num_images, - light_source.value, - ) +def _restore_mirrors(positions: list[tuple[FlipMirror, bool]]) -> MsgGenerator[None]: + yield from bps.mv(*chain.from_iterable(positions)) - yield from _setup_detector(detector, exposure_time) - yield from count([detector], num=num_images, md=run_metadata) +def _wind_back(stage: XThetaStage, start_angle: float) -> MsgGenerator[None]: + # unwind to leave the stage where the next scan expects it. + # currently this winds back the full rotation but should look + # into changing it so that we just advance one step to the + # beginning again and reset the angle to 0? + yield from bps.mv(stage.theta, start_angle, group="wind_back") def tomography( light_source: LightSource, detector: AravisDetector = detector, stage: XThetaStage = stage, + tomography_flip_mirror: FlipMirror = tomography_flip_mirror, + spectroscopy_flip_mirror: FlipMirror = spectroscopy_flip_mirror, pmac: PmacIO = pmac, pandabrick: HDFPanda = pandabrick, exposure_time: float = 0.1, spec: Spec[Movable] | None = None, + num_darks: int = 20, + num_flats: int = 20, metadata: dict[str, Any] | None = None, *, scan_type: ScanType, ) -> MsgGenerator[None]: - """Do a tomography scan.""" - yield from _setup_detector(detector, exposure_time) + """Do a tomography scan, taking darks and flats first. + + The darks, flats and projections are each their own run, sharing a + "tomography_group" id, and the projections run records the uids of the other + two. The flip mirrors are put back as they were afterwards. + + :param light_source: illumination the flats and projections are taken under. + """ + yield from load_settings( + device=detector, + design_name="tomography_detector_baseline", + whitelist_pvs=DETECTOR_WHITELIST, + ) + + # mv rather than prepare, because prepare cannot be used outside a run. + # See: https://github.com/DiamondLightSource/blueapi/issues/1211 + yield from bps.mv( + *(detector.driver.acquire_time, exposure_time), + *( + detector.driver.acquire_period, + exposure_time + MAKO_ACQUIRE_PERIOD_PAD, + ), + group="detector_acquire", + ) yield from load_settings( device=stage, @@ -155,41 +216,48 @@ def tomography( ) # One projection per degree over a full rotation - spec = spec or Line(stage, 0, 359, 360) # type: ignore + spec = spec or Line(stage.theta, 0, 359, 360) # type: ignore - if metadata is None: - metadata = {} - metadata["spec"] = serialize_spec(spec) # Tells the workflow which calibration images belong with these projections. - metadata[LIGHT_SOURCE_KEY] = light_source.value + shared_metadata: dict[str, Any] = { + GROUP_KEY: str(uuid4()), + LIGHT_SOURCE_KEY: light_source.value, + **(metadata or {}), + } + + initial_mirror_positions = [ + (tomography_flip_mirror, (yield from bps.rd(tomography_flip_mirror.position))), + ( + spectroscopy_flip_mirror, + (yield from bps.rd(spectroscopy_flip_mirror.position)), + ), + ] - if scan_type == ScanType.FLY: - LOGGER.info("Performing a tomography fly scan.") - yield from fly_scan( - spec, + yield from bpp.finalize_wrapper( + _darks_flats_projections( detector, - pmac=pmac, - pandabrick=pandabrick, - exposure_time=exposure_time, - detector_deadtime=MAKO_DETECTOR_DEADTIME, - detector_trigger=DETECTOR_TRIGGER, - pandabrick_whitelist=PANDA_WHITELIST, - plan_name="tomography_fly_scan", - metadata=metadata, - ) - elif scan_type == ScanType.STEP: - LOGGER.info("Performing a tomography step scan.") - yield from spec_scan( - {detector, stage}, - spec, # type: ignore - metadata=metadata, - ) + stage, + tomography_flip_mirror, + spectroscopy_flip_mirror, + pmac, + pandabrick, + exposure_time, + spec, + num_darks, + num_flats, + shared_metadata, + scan_type, + ), + _restore_mirrors(initial_mirror_positions), + ) def demo_tomography( light_source: LightSource, detector: AravisDetector = detector, stage: XThetaStage = stage, + tomography_flip_mirror: FlipMirror = tomography_flip_mirror, + spectroscopy_flip_mirror: FlipMirror = spectroscopy_flip_mirror, pmac: PmacIO = pmac, pandabrick: HDFPanda = pandabrick, number_of_projections: NumProjections = 360, @@ -201,29 +269,24 @@ def demo_tomography( scan_type: ScanType, ) -> MsgGenerator[None]: """Tomography plan intended for use in Visr demonstrations to visitors. - Takes number_of_projections evenly spaced over angular_range from start_angle, - then winds the stage back to start_angle. + Takes darks and flats, then number_of_projections evenly spaced over + angular_range from start_angle, then winds the stage back to start_angle. """ angular_step = angular_range / number_of_projections last_angle = start_angle + angular_range - angular_step - spec = Line(stage, start_angle, last_angle, number_of_projections) # type: ignore + spec = Line(stage.theta, start_angle, last_angle, number_of_projections) # type: ignore # PMAC handles the run-up, but move to the nominal start first so a failure # here is reported before the trajectory is built. - yield from bps.mv(stage, start_angle, group="initial_move") - - def wind_back() -> MsgGenerator[None]: - # unwind to leave the stage where the next scan expects it. - # currently this winds back the full rotation but should look - # into changing it so that we just advance one step to the - # beginning again and reset the angle to 0? - yield from bps.mv(stage, start_angle, group="wind_back") + yield from bps.mv(stage.theta, start_angle, group="initial_move") yield from bpp.finalize_wrapper( tomography( light_source, detector=detector, stage=stage, + tomography_flip_mirror=tomography_flip_mirror, + spectroscopy_flip_mirror=spectroscopy_flip_mirror, pmac=pmac, pandabrick=pandabrick, exposure_time=exposure_time, @@ -231,5 +294,5 @@ def wind_back() -> MsgGenerator[None]: metadata=metadata, scan_type=scan_type, ), - wind_back, + _wind_back(stage, start_angle), ) diff --git a/tests/unit_tests/test_plans.py b/tests/unit_tests/test_plans.py index 37df9d9..15456d8 100644 --- a/tests/unit_tests/test_plans.py +++ b/tests/unit_tests/test_plans.py @@ -9,8 +9,9 @@ import yaml from bluesky import RunEngine from dodal.devices.beamlines.b01_1.diodes import RGBDiodes +from dodal.devices.beamlines.b01_1.flip_mirror import FlipMirror from dodal.devices.current_amplifiers import EnzLocumRange -from dodal.devices.motors import XYZStage +from dodal.devices.motors import XThetaStage, XYZStage from ophyd_async.core import ( callback_on_mock_put, get_mock_put, @@ -36,6 +37,12 @@ demo_spectroscopy, spectroscopy, ) +from test_rig_bluesky.tomography_plans import ( + SAMPLE_IN_X, + SAMPLE_OUT_X, + LightSource, + tomography, +) @pytest.fixture @@ -92,8 +99,14 @@ async def on_acquire(acquire: bool) -> None: if acquire: asyncio.create_task(mock_acquisition()) + def on_capture(capture: bool) -> None: + # Like the IOC, start counting from zero for each new file + if capture: + set_mock_value(hdf.num_captured, 0) + set_mock_value(hdf.file_path_exists, True) callback_on_mock_put(detector.driver.acquire, on_acquire) + callback_on_mock_put(hdf.capture, on_capture) @patch("test_rig_bluesky.plans.YamlSettingsProvider") @@ -376,13 +389,13 @@ def test_spectroscopy_rejects_auto_amplifier_range( def test_fly_scan_rejects_pandabox_without_diodes( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, pandabox: HDFPanda, ): with pytest.raises(ValueError, match="pandabox was given without diodes"): run_engine( fly_scan( - Line(sample_stage.x, 0, 5, 5), + Line(spectroscopy_stage.x, 0, 5, 5), spectroscopy_detector, pandabox=pandabox, ) @@ -392,13 +405,13 @@ def test_fly_scan_rejects_pandabox_without_diodes( def test_fly_scan_rejects_diodes_without_pandabox( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, spectroscopy_diodes: RGBDiodes, ): with pytest.raises(ValueError, match="diodes were given without pandabox"): run_engine( fly_scan( - Line(sample_stage.x, 0, 5, 5), + Line(spectroscopy_stage.x, 0, 5, 5), spectroscopy_detector, diodes=spectroscopy_diodes, ) @@ -408,7 +421,7 @@ def test_fly_scan_rejects_diodes_without_pandabox( def test_spectroscopy_fly_scan_uses_pandabox_and_diodes( run_engine: RunEngine, spectroscopy_detector: AravisDetector, - sample_stage: XYZStage, + spectroscopy_stage: XYZStage, pandabrick: HDFPanda, pandabox: HDFPanda, spectroscopy_diodes: RGBDiodes, @@ -424,11 +437,11 @@ def test_spectroscopy_fly_scan_uses_pandabox_and_diodes( ): run_engine( spectroscopy( - spectroscopy_detector=spectroscopy_detector, - sample_stage=sample_stage, + detector=spectroscopy_detector, + stage=spectroscopy_stage, pandabrick=pandabrick, pandabox=pandabox, - spectroscopy_diodes=spectroscopy_diodes, + diodes=spectroscopy_diodes, scan_type=ScanType.FLY, ) ) @@ -455,3 +468,126 @@ def test_pandabox_baseline_captures_diodes_under_their_names( assert datasets["fmc_in.val5_dataset"] == spectroscopy_diodes.red.name assert datasets["fmc_in.val6_dataset"] == spectroscopy_diodes.green.name assert datasets["fmc_in.val7_dataset"] == spectroscopy_diodes.blue.name + + +@pytest.fixture +def tomography_detector(run_engine: RunEngine) -> AravisDetector: + det = b01_1.tomography_detector.build(connect_immediately=True, mock=True) + _mock_detector_behavior(det) + return det + + +@pytest.fixture +def tomography_stage(run_engine: RunEngine) -> XThetaStage: + stage = b01_1.tomography_stage.build(connect_immediately=True, mock=True) + + set_mock_value(stage.x.low_limit_travel, -10.0) + set_mock_value(stage.x.high_limit_travel, 10.0) + set_mock_value(stage.x.velocity, 1.0) + set_mock_value(stage.x.user_readback, 5.0) + set_mock_value(stage.theta.low_limit_travel, -400.0) + set_mock_value(stage.theta.high_limit_travel, 400.0) + set_mock_value(stage.theta.velocity, 10.0) + + return stage + + +@pytest.fixture +def tomography_flip_mirror(run_engine: RunEngine) -> FlipMirror: + mirror = b01_1.tomography_flip_mirror.build(connect_immediately=True, mock=True) + set_mock_value(mirror.position, True) + return mirror + + +@pytest.fixture +def spectroscopy_flip_mirror(run_engine: RunEngine) -> FlipMirror: + mirror = b01_1.spectroscopy_flip_mirror.build(connect_immediately=True, mock=True) + set_mock_value(mirror.position, True) + return mirror + + +def test_tomography_takes_darks_then_flats_then_projections( + run_engine: RunEngine, + tomography_detector: AravisDetector, + tomography_stage: XThetaStage, + tomography_flip_mirror: FlipMirror, + spectroscopy_flip_mirror: FlipMirror, +): + messages = [] + run_engine.msg_hook = messages.append + starts = [] + run_engine.subscribe(lambda name, doc: starts.append(doc), "start") + + run_engine( + tomography( + LightSource.LED, + detector=tomography_detector, + stage=tomography_stage, + tomography_flip_mirror=tomography_flip_mirror, + spectroscopy_flip_mirror=spectroscopy_flip_mirror, + spec=Line(tomography_stage.theta, 0, 3, 4), + num_darks=2, + num_flats=3, + scan_type=ScanType.STEP, + ) + ) + + darks, flats, projections = starts + assert [start["image_key"] for start in starts] == ["dark", "flat", "projection"] + assert [start["num_points"] for start in starts] == [2, 3, 4] + assert len({start["tomography_group"] for start in starts}) == 1 + assert all(start["light_source"] == "led" for start in starts) + assert projections["darks_uid"] == darks["uid"] + assert projections["flats_uid"] == flats["uid"] + + # Each run is preceded by the moves that set it up + sets_and_opens = [ + (msg.command, msg.obj, msg.args) + for msg in messages + if msg.command == "open_run" + or ( + msg.command == "set" + and msg.obj + in (tomography_stage.x, tomography_flip_mirror, spectroscopy_flip_mirror) + ) + ] + assert sets_and_opens == [ + ("set", tomography_stage.x, (SAMPLE_OUT_X,)), + ("set", spectroscopy_flip_mirror, (False,)), + ("set", tomography_flip_mirror, (False,)), + ("open_run", None, ()), + ("set", tomography_flip_mirror, (True,)), + ("open_run", None, ()), + ("set", tomography_stage.x, (SAMPLE_IN_X,)), + ("open_run", None, ()), + # Mirrors put back as they were + ("set", tomography_flip_mirror, (True,)), + ("set", spectroscopy_flip_mirror, (True,)), + ] + + +async def test_tomography_restores_mirrors_on_failure( + run_engine: RunEngine, + tomography_detector: AravisDetector, + tomography_stage: XThetaStage, + tomography_flip_mirror: FlipMirror, + spectroscopy_flip_mirror: FlipMirror, +): + with patch( + "test_rig_bluesky.tomography_plans.count", side_effect=RuntimeError("boom") + ): + with pytest.raises(RuntimeError, match="boom"): + run_engine( + tomography( + LightSource.LED, + detector=tomography_detector, + stage=tomography_stage, + tomography_flip_mirror=tomography_flip_mirror, + spectroscopy_flip_mirror=spectroscopy_flip_mirror, + spec=Line(tomography_stage.theta, 0, 3, 4), + scan_type=ScanType.STEP, + ) + ) + + assert await tomography_flip_mirror.position.get_value() is True + assert await spectroscopy_flip_mirror.position.get_value() is True From e454c4d150dc0074bf4b0a54f7555effbdfddb34 Mon Sep 17 00:00:00 2001 From: Tom Kane Date: Wed, 4 Nov 2026 14:30:47 +0000 Subject: [PATCH 08/12] refactor tomography --- src/test_rig_bluesky/tomography_plans.py | 248 +++++++++++++++-------- tests/unit_tests/test_plans.py | 11 +- 2 files changed, 166 insertions(+), 93 deletions(-) diff --git a/src/test_rig_bluesky/tomography_plans.py b/src/test_rig_bluesky/tomography_plans.py index 4d81d57..ec5feac 100644 --- a/src/test_rig_bluesky/tomography_plans.py +++ b/src/test_rig_bluesky/tomography_plans.py @@ -2,7 +2,6 @@ import logging from enum import StrEnum -from itertools import chain from typing import Annotated, Any from uuid import uuid4 @@ -34,17 +33,22 @@ class LightSource(StrEnum): - """The illumination in use.""" + """Deprecated: only kept so the demo_tomography API doesn't change. Ignored.""" LED = "led" SR = "sr" DARK = "dark" -# Metadata keys the reconstruction workflow reads to pair projections with the -# calibration images taken under the same illumination -LIGHT_SOURCE_KEY = "light_source" -# "dark", "flat" or "projection" +class ImageType(StrEnum): + """What a run of images is, so the reconstruction workflow can tell them apart.""" + + DARK = "dark" + FLAT = "flat" + PROJECTION = "projection" + + +# Metadata key holding the run's ImageType IMAGE_KEY = "image_key" # Shared by the darks, flats and projections runs of one tomography scan GROUP_KEY = "tomography_group" @@ -56,6 +60,14 @@ class LightSource(StrEnum): SAMPLE_OUT_X = 0.0 # TODO: placeholder, sample out of the field of view for flats SAMPLE_IN_X = 0.0 # TODO: placeholder, sample back in the field of view +# Where the sample sits and whether the tomography light is on for each image type. +# The spectroscopy mirror is always flipped out of the way. +SAMPLE_X_AND_LIGHT: dict[ImageType, tuple[float, bool]] = { + ImageType.DARK: (SAMPLE_OUT_X, False), + ImageType.FLAT: (SAMPLE_OUT_X, True), + ImageType.PROJECTION: (SAMPLE_IN_X, True), +} + MIN_PROJECTIONS = 30 MAX_PROJECTIONS = 1440 @@ -70,8 +82,8 @@ class LightSource(StrEnum): # TODO: conservative estimates for now -MAKO_ACQUIRE_PERIOD_PAD = 25e-3 -MAKO_DETECTOR_DEADTIME = 25e-3 +ALVIUM_ACQUIRE_PERIOD_PAD = 25e-3 +ALVIUM_DETECTOR_DEADTIME = 25e-3 DETECTOR_TRIGGER = DetectorTrigger.EXTERNAL_EDGE @@ -92,7 +104,67 @@ class LightSource(StrEnum): PANDA_WHITELIST = ["incenc-8-val_dataset"] -def _darks_flats_projections( +def _setup_detector(detector: AravisDetector, exposure_time: float) -> MsgGenerator: + """Put the detector in its scanning state.""" + # ---- load baseline settings ---------------------------------------------- + yield from load_settings( + device=detector, + design_name="tomography_detector_baseline", + whitelist_pvs=DETECTOR_WHITELIST, + ) + + # ---- set exposure -------------------------------------------------------- + # mv rather than prepare, because prepare cannot be used outside a run. + # See: https://github.com/DiamondLightSource/blueapi/issues/1211 + yield from bps.mv( + *(detector.driver.acquire_time, exposure_time), + *( + detector.driver.acquire_period, + exposure_time + ALVIUM_ACQUIRE_PERIOD_PAD, + ), + group="detector_acquire", + ) + + +def _move_for( + image_type: ImageType, + stage: XThetaStage, + tomography_flip_mirror: FlipMirror, + spectroscopy_flip_mirror: FlipMirror, +) -> MsgGenerator: + """Move the sample and mirrors into place for this type of image.""" + sample_x, light_on = SAMPLE_X_AND_LIGHT[image_type] + yield from bps.mv( + *(stage.x, sample_x), + *(spectroscopy_flip_mirror, False), + *(tomography_flip_mirror, light_on), + group=f"move_for_{image_type.value}", + ) + + +def _collect_calibration( + image_type: ImageType, + num_images: int, + detector: AravisDetector, + stage: XThetaStage, + tomography_flip_mirror: FlipMirror, + spectroscopy_flip_mirror: FlipMirror, + metadata: dict[str, Any], +) -> MsgGenerator[str]: + """Collect dark or flat images, returning the run's uid.""" + LOGGER.info("Collecting %d %s images.", num_images, image_type.value) + yield from _move_for( + image_type, stage, tomography_flip_mirror, spectroscopy_flip_mirror + ) + return ( + yield from count( + [detector], num=num_images, md={IMAGE_KEY: image_type.value, **metadata} + ) + ) + + +def _collect_projections( + spec: Spec[Movable], detector: AravisDetector, stage: XThetaStage, tomography_flip_mirror: FlipMirror, @@ -100,39 +172,23 @@ def _darks_flats_projections( pmac: PmacIO, pandabrick: HDFPanda, exposure_time: float, - spec: Spec[Movable], - num_darks: int, - num_flats: int, metadata: dict[str, Any], scan_type: ScanType, -) -> MsgGenerator[None]: - """Take darks, then flats with the sample out of the way, then projections.""" - LOGGER.info("Collecting %d dark images.", num_darks) - yield from bps.mv( - *(stage.x, SAMPLE_OUT_X), - *(spectroscopy_flip_mirror, False), - *(tomography_flip_mirror, False), - group="darks", - ) - darks_uid = yield from count( - [detector], num=num_darks, md={IMAGE_KEY: "dark", **metadata} - ) - - LOGGER.info("Collecting %d flat images.", num_flats) - yield from bps.mv(tomography_flip_mirror, True, group="flats") - flats_uid = yield from count( - [detector], num=num_flats, md={IMAGE_KEY: "flat", **metadata} +) -> MsgGenerator: + """Collect projections over spec, by fly or step scan.""" + # ---- move sample in, light on -------------------------------------------- + yield from _move_for( + ImageType.PROJECTION, stage, tomography_flip_mirror, spectroscopy_flip_mirror ) - yield from bps.mv(stage.x, SAMPLE_IN_X, group="projections") - projections_metadata = { - IMAGE_KEY: "projection", - DARKS_UID_KEY: darks_uid, - FLATS_UID_KEY: flats_uid, + # ---- build metadata ------------------------------------------------------ + metadata = { + IMAGE_KEY: ImageType.PROJECTION.value, "spec": serialize_spec(spec), **metadata, } + # ---- scan ---------------------------------------------------------------- if scan_type == ScanType.FLY: LOGGER.info("Performing a tomography fly scan.") yield from fly_scan( @@ -141,23 +197,68 @@ def _darks_flats_projections( pmac=pmac, pandabrick=pandabrick, exposure_time=exposure_time, - detector_deadtime=MAKO_DETECTOR_DEADTIME, + detector_deadtime=ALVIUM_DETECTOR_DEADTIME, detector_trigger=DETECTOR_TRIGGER, pandabrick_whitelist=PANDA_WHITELIST, plan_name="tomography_fly_scan", - metadata=projections_metadata, + metadata=metadata, ) elif scan_type == ScanType.STEP: LOGGER.info("Performing a tomography step scan.") yield from spec_scan( {detector, stage}, spec, # type: ignore - metadata=projections_metadata, + metadata=metadata, ) -def _restore_mirrors(positions: list[tuple[FlipMirror, bool]]) -> MsgGenerator[None]: - yield from bps.mv(*chain.from_iterable(positions)) +def _collect_images( + spec: Spec[Movable], + detector: AravisDetector, + stage: XThetaStage, + tomography_flip_mirror: FlipMirror, + spectroscopy_flip_mirror: FlipMirror, + pmac: PmacIO, + pandabrick: HDFPanda, + exposure_time: float, + num_darks: int, + num_flats: int, + metadata: dict[str, Any], + scan_type: ScanType, +) -> MsgGenerator: + """Darks, then flats, then projections linked to both.""" + mirrors = (tomography_flip_mirror, spectroscopy_flip_mirror) + + # ---- collect darks ------------------------------------------------------- + darks_uid = yield from _collect_calibration( + ImageType.DARK, num_darks, detector, stage, *mirrors, metadata + ) + + # ---- collect flats ------------------------------------------------------- + flats_uid = yield from _collect_calibration( + ImageType.FLAT, num_flats, detector, stage, *mirrors, metadata + ) + + # ---- collect projections ------------------------------------------------- + yield from _collect_projections( + spec, + detector, + stage, + *mirrors, + pmac, + pandabrick, + exposure_time, + {DARKS_UID_KEY: darks_uid, FLATS_UID_KEY: flats_uid, **metadata}, + scan_type, + ) + + +def _restoring_mirrors(plan: MsgGenerator, *mirrors: FlipMirror) -> MsgGenerator: + """Run plan, then put the mirrors back where they were, even if it fails.""" + initial_positions = [] + for mirror in mirrors: + initial_positions += [mirror, (yield from bps.rd(mirror.position))] + return (yield from bpp.finalize_wrapper(plan, bps.mv(*initial_positions))) def _wind_back(stage: XThetaStage, start_angle: float) -> MsgGenerator[None]: @@ -169,7 +270,6 @@ def _wind_back(stage: XThetaStage, start_angle: float) -> MsgGenerator[None]: def tomography( - light_source: LightSource, detector: AravisDetector = detector, stage: XThetaStage = stage, tomography_flip_mirror: FlipMirror = tomography_flip_mirror, @@ -183,58 +283,25 @@ def tomography( metadata: dict[str, Any] | None = None, *, scan_type: ScanType, -) -> MsgGenerator[None]: - """Do a tomography scan, taking darks and flats first. - - The darks, flats and projections are each their own run, sharing a - "tomography_group" id, and the projections run records the uids of the other - two. The flip mirrors are put back as they were afterwards. - - :param light_source: illumination the flats and projections are taken under. - """ - yield from load_settings( - device=detector, - design_name="tomography_detector_baseline", - whitelist_pvs=DETECTOR_WHITELIST, - ) - - # mv rather than prepare, because prepare cannot be used outside a run. - # See: https://github.com/DiamondLightSource/blueapi/issues/1211 - yield from bps.mv( - *(detector.driver.acquire_time, exposure_time), - *( - detector.driver.acquire_period, - exposure_time + MAKO_ACQUIRE_PERIOD_PAD, - ), - group="detector_acquire", - ) - +) -> MsgGenerator: + """Do a tomography scan: darks, then flats with the sample out, then projections.""" + # ---- set up detector and stage ------------------------------------------- + yield from _setup_detector(detector, exposure_time) yield from load_settings( device=stage, design_name="tomography_stage_baseline", whitelist_pvs=STAGE_WHITELIST, ) + # ---- define scan --------------------------------------------------------- # One projection per degree over a full rotation spec = spec or Line(stage.theta, 0, 359, 360) # type: ignore + shared_metadata = {GROUP_KEY: str(uuid4()), **(metadata or {})} - # Tells the workflow which calibration images belong with these projections. - shared_metadata: dict[str, Any] = { - GROUP_KEY: str(uuid4()), - LIGHT_SOURCE_KEY: light_source.value, - **(metadata or {}), - } - - initial_mirror_positions = [ - (tomography_flip_mirror, (yield from bps.rd(tomography_flip_mirror.position))), - ( - spectroscopy_flip_mirror, - (yield from bps.rd(spectroscopy_flip_mirror.position)), - ), - ] - - yield from bpp.finalize_wrapper( - _darks_flats_projections( + # ---- collect darks, flats and projections, then restore mirrors ---------- + yield from _restoring_mirrors( + _collect_images( + spec, detector, stage, tomography_flip_mirror, @@ -242,18 +309,18 @@ def tomography( pmac, pandabrick, exposure_time, - spec, num_darks, num_flats, shared_metadata, scan_type, ), - _restore_mirrors(initial_mirror_positions), + tomography_flip_mirror, + spectroscopy_flip_mirror, ) def demo_tomography( - light_source: LightSource, + light_source: LightSource, # deprecated detector: AravisDetector = detector, stage: XThetaStage = stage, tomography_flip_mirror: FlipMirror = tomography_flip_mirror, @@ -270,19 +337,22 @@ def demo_tomography( ) -> MsgGenerator[None]: """Tomography plan intended for use in Visr demonstrations to visitors. Takes darks and flats, then number_of_projections evenly spaced over - angular_range from start_angle, then winds the stage back to start_angle. - """ + angular_range from start_angle, then winds the stage back to start_angle.""" + del light_source + + # ---- define projections -------------------------------------------------- angular_step = angular_range / number_of_projections last_angle = start_angle + angular_range - angular_step spec = Line(stage.theta, start_angle, last_angle, number_of_projections) # type: ignore + # ---- move to start ------------------------------------------------------- # PMAC handles the run-up, but move to the nominal start first so a failure # here is reported before the trajectory is built. yield from bps.mv(stage.theta, start_angle, group="initial_move") + # ---- run tomography, then wind back -------------------------------------- yield from bpp.finalize_wrapper( tomography( - light_source, detector=detector, stage=stage, tomography_flip_mirror=tomography_flip_mirror, diff --git a/tests/unit_tests/test_plans.py b/tests/unit_tests/test_plans.py index 15456d8..234c90b 100644 --- a/tests/unit_tests/test_plans.py +++ b/tests/unit_tests/test_plans.py @@ -40,7 +40,6 @@ from test_rig_bluesky.tomography_plans import ( SAMPLE_IN_X, SAMPLE_OUT_X, - LightSource, tomography, ) @@ -520,7 +519,6 @@ def test_tomography_takes_darks_then_flats_then_projections( run_engine( tomography( - LightSource.LED, detector=tomography_detector, stage=tomography_stage, tomography_flip_mirror=tomography_flip_mirror, @@ -536,7 +534,6 @@ def test_tomography_takes_darks_then_flats_then_projections( assert [start["image_key"] for start in starts] == ["dark", "flat", "projection"] assert [start["num_points"] for start in starts] == [2, 3, 4] assert len({start["tomography_group"] for start in starts}) == 1 - assert all(start["light_source"] == "led" for start in starts) assert projections["darks_uid"] == darks["uid"] assert projections["flats_uid"] == flats["uid"] @@ -552,13 +549,20 @@ def test_tomography_takes_darks_then_flats_then_projections( ) ] assert sets_and_opens == [ + # Darks: sample out, no light ("set", tomography_stage.x, (SAMPLE_OUT_X,)), ("set", spectroscopy_flip_mirror, (False,)), ("set", tomography_flip_mirror, (False,)), ("open_run", None, ()), + # Flats: sample out, light on + ("set", tomography_stage.x, (SAMPLE_OUT_X,)), + ("set", spectroscopy_flip_mirror, (False,)), ("set", tomography_flip_mirror, (True,)), ("open_run", None, ()), + # Projections: sample in, light on ("set", tomography_stage.x, (SAMPLE_IN_X,)), + ("set", spectroscopy_flip_mirror, (False,)), + ("set", tomography_flip_mirror, (True,)), ("open_run", None, ()), # Mirrors put back as they were ("set", tomography_flip_mirror, (True,)), @@ -579,7 +583,6 @@ async def test_tomography_restores_mirrors_on_failure( with pytest.raises(RuntimeError, match="boom"): run_engine( tomography( - LightSource.LED, detector=tomography_detector, stage=tomography_stage, tomography_flip_mirror=tomography_flip_mirror, From a171e1e55d0f86ca1ddfb4f923b909e0943c2f95 Mon Sep 17 00:00:00 2001 From: Tom Kane Date: Wed, 4 Nov 2026 14:30:47 +0000 Subject: [PATCH 09/12] simplify by removing unnecessary restoring of states --- src/test_rig_bluesky/spectroscopy_plans.py | 2 + src/test_rig_bluesky/tomography_plans.py | 124 ++++++++------------- tests/unit_tests/test_plans.py | 34 +----- 3 files changed, 51 insertions(+), 109 deletions(-) diff --git a/src/test_rig_bluesky/spectroscopy_plans.py b/src/test_rig_bluesky/spectroscopy_plans.py index 7b459f9..dffe2c1 100644 --- a/src/test_rig_bluesky/spectroscopy_plans.py +++ b/src/test_rig_bluesky/spectroscopy_plans.py @@ -194,6 +194,8 @@ def spectroscopy( spec, # type: ignore metadata=metadata, ) + else: + raise ValueError(f"Unsupported scan type: {scan_type}") def demo_spectroscopy( diff --git a/src/test_rig_bluesky/tomography_plans.py b/src/test_rig_bluesky/tomography_plans.py index ec5feac..e9a06a4 100644 --- a/src/test_rig_bluesky/tomography_plans.py +++ b/src/test_rig_bluesky/tomography_plans.py @@ -104,7 +104,9 @@ class ImageType(StrEnum): PANDA_WHITELIST = ["incenc-8-val_dataset"] -def _setup_detector(detector: AravisDetector, exposure_time: float) -> MsgGenerator: +def _setup_detector( + detector: AravisDetector, exposure_time: float +) -> MsgGenerator[None]: """Put the detector in its scanning state.""" # ---- load baseline settings ---------------------------------------------- yield from load_settings( @@ -131,7 +133,7 @@ def _move_for( stage: XThetaStage, tomography_flip_mirror: FlipMirror, spectroscopy_flip_mirror: FlipMirror, -) -> MsgGenerator: +) -> MsgGenerator[None]: """Move the sample and mirrors into place for this type of image.""" sample_x, light_on = SAMPLE_X_AND_LIGHT[image_type] yield from bps.mv( @@ -158,7 +160,7 @@ def _collect_calibration( ) return ( yield from count( - [detector], num=num_images, md={IMAGE_KEY: image_type.value, **metadata} + [detector], num=num_images, md={**metadata, IMAGE_KEY: image_type.value} ) ) @@ -174,7 +176,7 @@ def _collect_projections( exposure_time: float, metadata: dict[str, Any], scan_type: ScanType, -) -> MsgGenerator: +) -> MsgGenerator[None]: """Collect projections over spec, by fly or step scan.""" # ---- move sample in, light on -------------------------------------------- yield from _move_for( @@ -183,13 +185,13 @@ def _collect_projections( # ---- build metadata ------------------------------------------------------ metadata = { + **metadata, IMAGE_KEY: ImageType.PROJECTION.value, "spec": serialize_spec(spec), - **metadata, } - # ---- scan ---------------------------------------------------------------- if scan_type == ScanType.FLY: + # ---- fly scan -------------------------------------------------------- LOGGER.info("Performing a tomography fly scan.") yield from fly_scan( spec, @@ -204,61 +206,16 @@ def _collect_projections( metadata=metadata, ) elif scan_type == ScanType.STEP: + # ---- step scan ------------------------------------------------------- LOGGER.info("Performing a tomography step scan.") yield from spec_scan( {detector, stage}, spec, # type: ignore metadata=metadata, ) - - -def _collect_images( - spec: Spec[Movable], - detector: AravisDetector, - stage: XThetaStage, - tomography_flip_mirror: FlipMirror, - spectroscopy_flip_mirror: FlipMirror, - pmac: PmacIO, - pandabrick: HDFPanda, - exposure_time: float, - num_darks: int, - num_flats: int, - metadata: dict[str, Any], - scan_type: ScanType, -) -> MsgGenerator: - """Darks, then flats, then projections linked to both.""" - mirrors = (tomography_flip_mirror, spectroscopy_flip_mirror) - - # ---- collect darks ------------------------------------------------------- - darks_uid = yield from _collect_calibration( - ImageType.DARK, num_darks, detector, stage, *mirrors, metadata - ) - - # ---- collect flats ------------------------------------------------------- - flats_uid = yield from _collect_calibration( - ImageType.FLAT, num_flats, detector, stage, *mirrors, metadata - ) - - # ---- collect projections ------------------------------------------------- - yield from _collect_projections( - spec, - detector, - stage, - *mirrors, - pmac, - pandabrick, - exposure_time, - {DARKS_UID_KEY: darks_uid, FLATS_UID_KEY: flats_uid, **metadata}, - scan_type, - ) - - -def _restoring_mirrors(plan: MsgGenerator, *mirrors: FlipMirror) -> MsgGenerator: - """Run plan, then put the mirrors back where they were, even if it fails.""" - initial_positions = [] - for mirror in mirrors: - initial_positions += [mirror, (yield from bps.rd(mirror.position))] - return (yield from bpp.finalize_wrapper(plan, bps.mv(*initial_positions))) + else: + # ---- unsupported scan type ------------------------------------------- + raise ValueError(f"Unsupported scan type: {scan_type}") def _wind_back(stage: XThetaStage, start_angle: float) -> MsgGenerator[None]: @@ -278,12 +235,11 @@ def tomography( pandabrick: HDFPanda = pandabrick, exposure_time: float = 0.1, spec: Spec[Movable] | None = None, - num_darks: int = 20, - num_flats: int = 20, + num_calibration_images: int = 20, metadata: dict[str, Any] | None = None, *, scan_type: ScanType, -) -> MsgGenerator: +) -> MsgGenerator[None]: """Do a tomography scan: darks, then flats with the sample out, then projections.""" # ---- set up detector and stage ------------------------------------------- yield from _setup_detector(detector, exposure_time) @@ -296,26 +252,40 @@ def tomography( # ---- define scan --------------------------------------------------------- # One projection per degree over a full rotation spec = spec or Line(stage.theta, 0, 359, 360) # type: ignore - shared_metadata = {GROUP_KEY: str(uuid4()), **(metadata or {})} + shared_metadata = {**(metadata or {}), GROUP_KEY: str(uuid4())} + mirrors = (tomography_flip_mirror, spectroscopy_flip_mirror) - # ---- collect darks, flats and projections, then restore mirrors ---------- - yield from _restoring_mirrors( - _collect_images( - spec, - detector, - stage, - tomography_flip_mirror, - spectroscopy_flip_mirror, - pmac, - pandabrick, - exposure_time, - num_darks, - num_flats, - shared_metadata, - scan_type, - ), - tomography_flip_mirror, - spectroscopy_flip_mirror, + # ---- collect darks ------------------------------------------------------- + darks_uid = yield from _collect_calibration( + ImageType.DARK, + num_calibration_images, + detector, + stage, + *mirrors, + shared_metadata, + ) + + # ---- collect flats ------------------------------------------------------- + flats_uid = yield from _collect_calibration( + ImageType.FLAT, + num_calibration_images, + detector, + stage, + *mirrors, + shared_metadata, + ) + + # ---- collect projections ------------------------------------------------- + yield from _collect_projections( + spec, + detector, + stage, + *mirrors, + pmac, + pandabrick, + exposure_time, + {**shared_metadata, DARKS_UID_KEY: darks_uid, FLATS_UID_KEY: flats_uid}, + scan_type, ) diff --git a/tests/unit_tests/test_plans.py b/tests/unit_tests/test_plans.py index 234c90b..07a6def 100644 --- a/tests/unit_tests/test_plans.py +++ b/tests/unit_tests/test_plans.py @@ -524,15 +524,14 @@ def test_tomography_takes_darks_then_flats_then_projections( tomography_flip_mirror=tomography_flip_mirror, spectroscopy_flip_mirror=spectroscopy_flip_mirror, spec=Line(tomography_stage.theta, 0, 3, 4), - num_darks=2, - num_flats=3, + num_calibration_images=2, scan_type=ScanType.STEP, ) ) darks, flats, projections = starts assert [start["image_key"] for start in starts] == ["dark", "flat", "projection"] - assert [start["num_points"] for start in starts] == [2, 3, 4] + assert [start["num_points"] for start in starts] == [2, 2, 4] assert len({start["tomography_group"] for start in starts}) == 1 assert projections["darks_uid"] == darks["uid"] assert projections["flats_uid"] == flats["uid"] @@ -564,33 +563,4 @@ def test_tomography_takes_darks_then_flats_then_projections( ("set", spectroscopy_flip_mirror, (False,)), ("set", tomography_flip_mirror, (True,)), ("open_run", None, ()), - # Mirrors put back as they were - ("set", tomography_flip_mirror, (True,)), - ("set", spectroscopy_flip_mirror, (True,)), ] - - -async def test_tomography_restores_mirrors_on_failure( - run_engine: RunEngine, - tomography_detector: AravisDetector, - tomography_stage: XThetaStage, - tomography_flip_mirror: FlipMirror, - spectroscopy_flip_mirror: FlipMirror, -): - with patch( - "test_rig_bluesky.tomography_plans.count", side_effect=RuntimeError("boom") - ): - with pytest.raises(RuntimeError, match="boom"): - run_engine( - tomography( - detector=tomography_detector, - stage=tomography_stage, - tomography_flip_mirror=tomography_flip_mirror, - spectroscopy_flip_mirror=spectroscopy_flip_mirror, - spec=Line(tomography_stage.theta, 0, 3, 4), - scan_type=ScanType.STEP, - ) - ) - - assert await tomography_flip_mirror.position.get_value() is True - assert await spectroscopy_flip_mirror.position.get_value() is True From 2af0e6ac27d0ef35a1ccb5c1afb3d86759a94fe9 Mon Sep 17 00:00:00 2001 From: Tom Kane Date: Mon, 9 Nov 2026 12:54:34 +0000 Subject: [PATCH 10/12] remove light source entirely --- src/test_rig_bluesky/tomography_plans.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/src/test_rig_bluesky/tomography_plans.py b/src/test_rig_bluesky/tomography_plans.py index e9a06a4..6154c77 100644 --- a/src/test_rig_bluesky/tomography_plans.py +++ b/src/test_rig_bluesky/tomography_plans.py @@ -32,14 +32,6 @@ spectroscopy_flip_mirror = inject("spectroscopy_flip_mirror") -class LightSource(StrEnum): - """Deprecated: only kept so the demo_tomography API doesn't change. Ignored.""" - - LED = "led" - SR = "sr" - DARK = "dark" - - class ImageType(StrEnum): """What a run of images is, so the reconstruction workflow can tell them apart.""" @@ -290,7 +282,6 @@ def tomography( def demo_tomography( - light_source: LightSource, # deprecated detector: AravisDetector = detector, stage: XThetaStage = stage, tomography_flip_mirror: FlipMirror = tomography_flip_mirror, @@ -308,7 +299,6 @@ def demo_tomography( """Tomography plan intended for use in Visr demonstrations to visitors. Takes darks and flats, then number_of_projections evenly spaced over angular_range from start_angle, then winds the stage back to start_angle.""" - del light_source # ---- define projections -------------------------------------------------- angular_step = angular_range / number_of_projections From 81b28334fbf162dc4452c307c0415d56ad19c2ae Mon Sep 17 00:00:00 2001 From: Tom Kane Date: Mon, 9 Nov 2026 13:36:38 +0000 Subject: [PATCH 11/12] update dodal for flipper mirror imports --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 239e45b..6a48335 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,7 @@ classifiers = [ "Programming Language :: Python :: 3.12", ] description = "Bluesky plans to be run on Diamond's test rigs e.g. ViSR, P45, etc" -dependencies = ["dls-dodal>=3.0.7", "scanspec>=1.0.0"] +dependencies = ["dls-dodal>=3.0.8", "scanspec>=1.0.0"] dynamic = ["version"] license.file = "LICENSE" readme = "README.md" From fcb148f44eddd89cc873d05507b089a8e30123ab Mon Sep 17 00:00:00 2001 From: Tom Kane Date: Mon, 9 Nov 2026 14:25:16 +0000 Subject: [PATCH 12/12] use msg mutator to avoid typing error --- tests/unit_tests/test_plans.py | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/tests/unit_tests/test_plans.py b/tests/unit_tests/test_plans.py index 07a6def..336856c 100644 --- a/tests/unit_tests/test_plans.py +++ b/tests/unit_tests/test_plans.py @@ -4,10 +4,12 @@ from pathlib import Path from unittest.mock import ANY, AsyncMock, Mock, patch +import bluesky.preprocessors as bpp import dodal.beamlines.b01_1 as b01_1 import pytest import yaml from bluesky import RunEngine +from bluesky.utils import Msg from dodal.devices.beamlines.b01_1.diodes import RGBDiodes from dodal.devices.beamlines.b01_1.flip_mirror import FlipMirror from dodal.devices.current_amplifiers import EnzLocumRange @@ -512,20 +514,27 @@ def test_tomography_takes_darks_then_flats_then_projections( tomography_flip_mirror: FlipMirror, spectroscopy_flip_mirror: FlipMirror, ): - messages = [] - run_engine.msg_hook = messages.append + messages: list[Msg] = [] + + def record(msg: Msg) -> Msg: + messages.append(msg) + return msg + starts = [] run_engine.subscribe(lambda name, doc: starts.append(doc), "start") run_engine( - tomography( - detector=tomography_detector, - stage=tomography_stage, - tomography_flip_mirror=tomography_flip_mirror, - spectroscopy_flip_mirror=spectroscopy_flip_mirror, - spec=Line(tomography_stage.theta, 0, 3, 4), - num_calibration_images=2, - scan_type=ScanType.STEP, + bpp.msg_mutator( + tomography( + detector=tomography_detector, + stage=tomography_stage, + tomography_flip_mirror=tomography_flip_mirror, + spectroscopy_flip_mirror=spectroscopy_flip_mirror, + spec=Line(tomography_stage.theta, 0, 3, 4), + num_calibration_images=2, + scan_type=ScanType.STEP, + ), + record, ) )