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" 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/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/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..dffe2c1 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,26 @@ 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, ) + else: + raise ValueError(f"Unsupported scan type: {scan_type}") 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 +237,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..6154c77 100644 --- a/src/test_rig_bluesky/tomography_plans.py +++ b/src/test_rig_bluesky/tomography_plans.py @@ -3,6 +3,7 @@ import logging from enum import StrEnum from typing import Annotated, Any +from uuid import uuid4 import bluesky.plan_stubs as bps import bluesky.preprocessors as bpp @@ -10,10 +11,11 @@ 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 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,21 +26,39 @@ 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") +tomography_flip_mirror = inject("tomography_flip_mirror") +spectroscopy_flip_mirror = inject("spectroscopy_flip_mirror") -class LightSource(StrEnum): - """The illumination in use.""" +class ImageType(StrEnum): + """What a run of images is, so the reconstruction workflow can tell them apart.""" - 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" + 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" +# 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 + +# 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 @@ -54,10 +74,10 @@ 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 -TOMOGRAPHY_DETECTOR_TRIGGER = DetectorTrigger.EXTERNAL_EDGE +DETECTOR_TRIGGER = DetectorTrigger.EXTERNAL_EDGE DETECTOR_WHITELIST = [ "hdf-nd_array_port", @@ -67,9 +87,9 @@ class LightSource(StrEnum): "driver-trigger_source", ] -TOMOGRAPHY_STAGE_WHITELIST = [ - "acceleration_time", - "velocity", +STAGE_WHITELIST = [ + "theta-acceleration_time", + "theta-velocity", ] # TODO: confirm which PandA encoder input the rotation axis is wired to. @@ -77,119 +97,195 @@ class LightSource(StrEnum): def _setup_detector( - detector: AravisDetector, - exposure_time: float, + detector: AravisDetector, exposure_time: float ) -> MsgGenerator[None]: """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 + MAKO_ACQUIRE_PERIOD_PAD, + exposure_time + ALVIUM_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, - num_images: int = 20, - exposure_time: float = 0.1, - metadata: dict[str, Any] | None = None, +def _move_for( + image_type: ImageType, + stage: XThetaStage, + tomography_flip_mirror: FlipMirror, + spectroscopy_flip_mirror: FlipMirror, ) -> MsgGenerator[None]: - """Collect flat or dark field images for the reconstruction to normalise with. - - 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. + """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}", + ) - :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 _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={**metadata, IMAGE_KEY: image_type.value} ) - - yield from _setup_detector(tomography_detector, exposure_time) - - yield from count([tomography_detector], num=num_images, md=run_metadata) + ) -def tomography( - light_source: LightSource, - tomography_detector: AravisDetector = tomography_detector, - tomography_stage: Motor = tomography_stage, - pmac: PmacIO = pmac, - pandabrick: HDFPanda = pandabrick, - exposure_time: float = 0.1, - spec: Spec[Movable] | None = None, - metadata: dict[str, Any] | None = None, - *, +def _collect_projections( + spec: Spec[Movable], + detector: AravisDetector, + stage: XThetaStage, + tomography_flip_mirror: FlipMirror, + spectroscopy_flip_mirror: FlipMirror, + pmac: PmacIO, + pandabrick: HDFPanda, + exposure_time: float, + metadata: dict[str, Any], scan_type: ScanType, ) -> MsgGenerator[None]: - """Do a tomography scan.""" - yield from _setup_detector(tomography_detector, exposure_time) - - yield from load_settings( - device=tomography_stage, - design_name="tomography_stage_baseline", - whitelist_pvs=TOMOGRAPHY_STAGE_WHITELIST, + """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 ) - # One projection per degree over a full rotation - spec = spec or Line(tomography_stage, 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 + # ---- build metadata ------------------------------------------------------ + metadata = { + **metadata, + IMAGE_KEY: ImageType.PROJECTION.value, + "spec": serialize_spec(spec), + } if scan_type == ScanType.FLY: + # ---- fly scan -------------------------------------------------------- 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_deadtime=ALVIUM_DETECTOR_DEADTIME, + detector_trigger=DETECTOR_TRIGGER, pandabrick_whitelist=PANDA_WHITELIST, plan_name="tomography_fly_scan", metadata=metadata, ) elif scan_type == ScanType.STEP: + # ---- step scan ------------------------------------------------------- LOGGER.info("Performing a tomography step scan.") yield from spec_scan( - {tomography_detector, tomography_stage}, + {detector, stage}, spec, # type: ignore metadata=metadata, ) + else: + # ---- unsupported scan type ------------------------------------------- + raise ValueError(f"Unsupported scan type: {scan_type}") + + +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( + 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_calibration_images: int = 20, + metadata: dict[str, Any] | None = None, + *, + scan_type: ScanType, +) -> 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) + 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 = {**(metadata or {}), GROUP_KEY: str(uuid4())} + mirrors = (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, + ) def demo_tomography( - light_source: LightSource, - tomography_detector: AravisDetector = tomography_detector, - tomography_stage: Motor = tomography_stage, + 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 +297,26 @@ 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.""" + + # ---- define projections -------------------------------------------------- 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.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(tomography_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.theta, start_angle, group="initial_move") + # ---- run tomography, then wind back -------------------------------------- yield from bpp.finalize_wrapper( tomography( - light_source, - tomography_detector=tomography_detector, - tomography_stage=tomography_stage, + 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 +324,5 @@ def wind_back() -> MsgGenerator[None]: metadata=metadata, scan_type=scan_type, ), - wind_back, + _wind_back(stage, start_angle), ) 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..336856c 100644 --- a/tests/unit_tests/test_plans.py +++ b/tests/unit_tests/test_plans.py @@ -4,13 +4,16 @@ 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 -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 +39,11 @@ demo_spectroscopy, spectroscopy, ) +from test_rig_bluesky.tomography_plans import ( + SAMPLE_IN_X, + SAMPLE_OUT_X, + tomography, +) @pytest.fixture @@ -46,7 +54,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) @@ -92,8 +100,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") @@ -143,12 +157,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 +178,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 +187,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 +214,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 +223,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 +247,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 +256,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 +279,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 +307,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 +319,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 +352,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 +371,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 +379,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, ) ) @@ -375,13 +390,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, ) @@ -391,13 +406,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, ) @@ -407,7 +422,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, @@ -423,11 +438,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, ) ) @@ -454,3 +469,107 @@ 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: 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( + 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, + ) + ) + + darks, flats, projections = starts + assert [start["image_key"] for start in starts] == ["dark", "flat", "projection"] + 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"] + + # 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 == [ + # 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, ()), + ]