mx-bluesky 0.3.1__py3-none-any.whl → 1.2.0__py3-none-any.whl

mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py

@@ -0,0 +1,537 @@
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Callable
+from functools import partial
+from pathlib import Path
+from time import time
+from typing import Protocol
+
+import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
+import numpy as np
+from blueapi.core import BlueskyContext, MsgGenerator
+from dodal.devices.aperturescatterguard import (
+    ApertureScatterguard,
+    ApertureValue,
+)
+from dodal.devices.attenuator import Attenuator
+from dodal.devices.backlight import Backlight
+from dodal.devices.dcm import DCM
+from dodal.devices.eiger import EigerDetector
+from dodal.devices.fast_grid_scan import (
+    FastGridScanCommon,
+    PandAFastGridScan,
+    ZebraFastGridScan,
+)
+from dodal.devices.fast_grid_scan import (
+    set_fast_grid_scan_params as set_flyscan_params,
+)
+from dodal.devices.flux import Flux
+from dodal.devices.robot import BartRobot
+from dodal.devices.s4_slit_gaps import S4SlitGaps
+from dodal.devices.smargon import Smargon, StubPosition
+from dodal.devices.synchrotron import Synchrotron
+from dodal.devices.undulator import Undulator
+from dodal.devices.xbpm_feedback import XBPMFeedback
+from dodal.devices.zebra import Zebra
+from dodal.devices.zebra_controlled_shutter import ZebraShutter
+from dodal.devices.zocalo.zocalo_results import (
+    ZOCALO_READING_PLAN_NAME,
+    ZOCALO_STAGE_GROUP,
+    ZocaloResults,
+    get_processing_result,
+)
+from dodal.plans.check_topup import check_topup_and_wait_if_necessary
+from ophyd_async.fastcs.panda import HDFPanda
+from scanspec.core import AxesPoints, Axis
+
+from mx_bluesky.hyperion.device_setup_plans.manipulate_sample import move_x_y_z
+from mx_bluesky.hyperion.device_setup_plans.read_hardware_for_setup import (
+    read_hardware_during_collection,
+    read_hardware_for_zocalo,
+    read_hardware_pre_collection,
+)
+from mx_bluesky.hyperion.device_setup_plans.setup_panda import (
+    disarm_panda_for_gridscan,
+    set_panda_directory,
+    setup_panda_for_flyscan,
+)
+from mx_bluesky.hyperion.device_setup_plans.setup_zebra import (
+    setup_zebra_for_gridscan,
+    setup_zebra_for_panda_flyscan,
+    tidy_up_zebra_after_gridscan,
+)
+from mx_bluesky.hyperion.device_setup_plans.xbpm_feedback import (
+    transmission_and_xbpm_feedback_for_collection_decorator,
+)
+from mx_bluesky.hyperion.exceptions import WarningException
+from mx_bluesky.hyperion.log import LOGGER
+from mx_bluesky.hyperion.parameters.constants import CONST
+from mx_bluesky.hyperion.parameters.gridscan import ThreeDGridScan
+from mx_bluesky.hyperion.tracing import TRACER
+from mx_bluesky.hyperion.utils.context import device_composite_from_context
+
+
+class SmargonSpeedException(Exception):
+    pass
+
+
+class CrystalNotFoundException(WarningException):
+    """Raised if grid detection completed normally but no crystal was found."""
+
+    pass
+
+
+@dataclasses.dataclass
+class FlyScanXRayCentreComposite:
+    """All devices which are directly or indirectly required by this plan"""
+
+    aperture_scatterguard: ApertureScatterguard
+    attenuator: Attenuator
+    backlight: Backlight
+    dcm: DCM
+    eiger: EigerDetector
+    zebra_fast_grid_scan: ZebraFastGridScan
+    flux: Flux
+    s4_slit_gaps: S4SlitGaps
+    smargon: Smargon
+    undulator: Undulator
+    synchrotron: Synchrotron
+    xbpm_feedback: XBPMFeedback
+    zebra: Zebra
+    zocalo: ZocaloResults
+    panda: HDFPanda
+    panda_fast_grid_scan: PandAFastGridScan
+    robot: BartRobot
+    sample_shutter: ZebraShutter
+
+    @property
+    def sample_motors(self) -> Smargon:
+        """Convenience alias with a more user-friendly name"""
+        return self.smargon
+
+
+def create_devices(context: BlueskyContext) -> FlyScanXRayCentreComposite:
+    """Creates the devices required for the plan and connect to them"""
+    return device_composite_from_context(context, FlyScanXRayCentreComposite)
+
+
+def flyscan_xray_centre(
+    composite: FlyScanXRayCentreComposite,
+    parameters: ThreeDGridScan,
+) -> MsgGenerator:
+    """Create the plan to run the grid scan based on provided parameters.
+
+    The ispyb handler should be added to the whole gridscan as we want to capture errors
+    at any point in it.
+
+    Args:
+        parameters (ThreeDGridScan): The parameters to run the scan.
+
+    Returns:
+        Generator: The plan for the gridscan
+    """
+    parameters.features.update_self_from_server()
+    composite.eiger.set_detector_parameters(parameters.detector_params)
+    composite.zocalo.zocalo_environment = parameters.zocalo_environment
+    composite.zocalo.use_cpu_and_gpu = parameters.use_cpu_and_gpu_zocalo
+
+    feature_controlled = _get_feature_controlled(composite, parameters)
+
+    @bpp.set_run_key_decorator(CONST.PLAN.GRIDSCAN_OUTER)
+    @bpp.run_decorator(  # attach experiment metadata to the start document
+        md={
+            "subplan_name": CONST.PLAN.GRIDSCAN_OUTER,
+            CONST.TRIGGER.ZOCALO: CONST.PLAN.DO_FGS,
+            "zocalo_environment": parameters.zocalo_environment,
+            "hyperion_parameters": parameters.model_dump_json(),
+            "activate_callbacks": [
+                "GridscanNexusFileCallback",
+            ],
+        }
+    )
+    @bpp.finalize_decorator(lambda: feature_controlled.tidy_plan(composite))
+    @transmission_and_xbpm_feedback_for_collection_decorator(
+        composite.xbpm_feedback,
+        composite.attenuator,
+        parameters.transmission_frac,
+    )
+    def run_gridscan_and_move_and_tidy(
+        fgs_composite: FlyScanXRayCentreComposite,
+        params: ThreeDGridScan,
+        feature_controlled: _FeatureControlled,
+    ):
+        yield from run_gridscan_and_move(fgs_composite, params, feature_controlled)
+
+    return run_gridscan_and_move_and_tidy(composite, parameters, feature_controlled)
+
+
+@bpp.set_run_key_decorator(CONST.PLAN.GRIDSCAN_AND_MOVE)
+@bpp.run_decorator(md={"subplan_name": CONST.PLAN.GRIDSCAN_AND_MOVE})
+def run_gridscan_and_move(
+    fgs_composite: FlyScanXRayCentreComposite,
+    parameters: ThreeDGridScan,
+    feature_controlled: _FeatureControlled,
+) -> MsgGenerator:
+    """A multi-run plan which runs a gridscan, gets the results from zocalo
+    and moves to the centre of mass determined by zocalo"""
+
+    # We get the initial motor positions so we can return to them on zocalo failure
+    initial_xyz = np.array(
+        [
+            (yield from bps.rd(fgs_composite.sample_motors.x)),
+            (yield from bps.rd(fgs_composite.sample_motors.y)),
+            (yield from bps.rd(fgs_composite.sample_motors.z)),
+        ]
+    )
+
+    yield from feature_controlled.setup_trigger(fgs_composite, parameters, initial_xyz)
+
+    LOGGER.info("Starting grid scan")
+    yield from bps.stage(
+        fgs_composite.zocalo, group=ZOCALO_STAGE_GROUP
+    )  # connect to zocalo and make sure the queue is clear
+    yield from run_gridscan(fgs_composite, parameters, feature_controlled)
+
+    LOGGER.info("Grid scan finished, getting results.")
+
+    try:
+        with TRACER.start_span("wait_for_zocalo"):
+            yield from bps.trigger_and_read(
+                [fgs_composite.zocalo], name=ZOCALO_READING_PLAN_NAME
+            )
+            LOGGER.info("Zocalo triggered and read, interpreting results.")
+            xray_centre, bbox_size = yield from get_processing_result(
+                fgs_composite.zocalo
+            )
+            LOGGER.info(f"Got xray centre: {xray_centre}, bbox size: {bbox_size}")
+            if xray_centre is not None:
+                xray_centre = parameters.FGS_params.grid_position_to_motor_position(
+                    xray_centre
+                )
+            else:
+                LOGGER.warning("No X-ray centre received")
+                raise CrystalNotFoundException()
+            if bbox_size is not None:
+                with TRACER.start_span("change_aperture"):
+                    yield from set_aperture_for_bbox_size(
+                        fgs_composite.aperture_scatterguard, bbox_size
+                    )
+            else:
+                LOGGER.warning("No bounding box size received")
+
+        # once we have the results, go to the appropriate position
+        LOGGER.info("Moving to centre of mass.")
+        with TRACER.start_span("move_to_result"):
+            x, y, z = xray_centre
+            yield from move_x_y_z(fgs_composite.sample_motors, x, y, z, wait=True)
+
+        if parameters.FGS_params.set_stub_offsets:
+            LOGGER.info("Recentring smargon co-ordinate system to this point.")
+            yield from bps.mv(
+                fgs_composite.sample_motors.stub_offsets, StubPosition.CURRENT_AS_CENTER
+            )
+    finally:
+        # Turn off dev/shm streaming to avoid filling disk, see https://github.com/DiamondLightSource/hyperion/issues/1395
+        LOGGER.info("Turning off Eiger dev/shm streaming")
+        yield from bps.abs_set(fgs_composite.eiger.odin.fan.dev_shm_enable, 0)
+
+        # Wait on everything before returning to GDA (particularly apertures), can be removed
+        # when we do not return to GDA here
+        yield from bps.wait()
+
+
+@bpp.set_run_key_decorator(CONST.PLAN.GRIDSCAN_MAIN)
+@bpp.run_decorator(md={"subplan_name": CONST.PLAN.GRIDSCAN_MAIN})
+def run_gridscan(
+    fgs_composite: FlyScanXRayCentreComposite,
+    parameters: ThreeDGridScan,
+    feature_controlled: _FeatureControlled,
+    md={  # noqa
+        "plan_name": CONST.PLAN.GRIDSCAN_MAIN,
+    },
+):
+    sample_motors = fgs_composite.sample_motors
+
+    # Currently gridscan only works for omega 0, see #
+    with TRACER.start_span("moving_omega_to_0"):
+        yield from bps.abs_set(sample_motors.omega, 0)
+
+    # We only subscribe to the communicator callback for run_gridscan, so this is where
+    # we should generate an event reading the values which need to be included in the
+    # ispyb deposition
+    with TRACER.start_span("ispyb_hardware_readings"):
+        yield from read_hardware_pre_collection(
+            fgs_composite.undulator,
+            fgs_composite.synchrotron,
+            fgs_composite.s4_slit_gaps,
+            fgs_composite.robot,
+            fgs_composite.smargon,
+        )
+
+    read_during_collection = partial(
+        read_hardware_during_collection,
+        fgs_composite.aperture_scatterguard,
+        fgs_composite.attenuator,
+        fgs_composite.flux,
+        fgs_composite.dcm,
+        fgs_composite.eiger,
+    )
+
+    LOGGER.info("Setting fgs params")
+    yield from feature_controlled.set_flyscan_params()
+
+    LOGGER.info("Waiting for gridscan validity check")
+    yield from wait_for_gridscan_valid(feature_controlled.fgs_motors)
+
+    LOGGER.info("Waiting for arming to finish")
+    yield from bps.wait(CONST.WAIT.GRID_READY_FOR_DC)
+    yield from bps.stage(fgs_composite.eiger)
+
+    yield from kickoff_and_complete_gridscan(
+        feature_controlled.fgs_motors,
+        fgs_composite.eiger,
+        fgs_composite.synchrotron,
+        [parameters.scan_points_first_grid, parameters.scan_points_second_grid],
+        parameters.scan_indices,
+        do_during_run=read_during_collection,
+    )
+    yield from bps.abs_set(feature_controlled.fgs_motors.z_steps, 0, wait=False)
+
+
+def kickoff_and_complete_gridscan(
+    gridscan: FastGridScanCommon,
+    eiger: EigerDetector,
+    synchrotron: Synchrotron,
+    scan_points: list[AxesPoints[Axis]],
+    scan_start_indices: list[int],
+    do_during_run: Callable[[], MsgGenerator] | None = None,
+):
+    @TRACER.start_as_current_span(CONST.PLAN.DO_FGS)
+    @bpp.set_run_key_decorator(CONST.PLAN.DO_FGS)
+    @bpp.run_decorator(
+        md={
+            "subplan_name": CONST.PLAN.DO_FGS,
+            "scan_points": scan_points,
+            "scan_start_indices": scan_start_indices,
+        }
+    )
+    @bpp.contingency_decorator(
+        except_plan=lambda e: (yield from bps.stop(eiger)),
+        else_plan=lambda: (yield from bps.unstage(eiger)),
+    )
+    def do_fgs():
+        # Check topup gate
+        expected_images = yield from bps.rd(gridscan.expected_images)
+        exposure_sec_per_image = yield from bps.rd(eiger.cam.acquire_time)
+        LOGGER.info("waiting for topup if necessary...")
+        yield from check_topup_and_wait_if_necessary(
+            synchrotron,
+            expected_images * exposure_sec_per_image,
+            30.0,
+        )
+        yield from read_hardware_for_zocalo(eiger)
+        LOGGER.info("Wait for all moves with no assigned group")
+        yield from bps.wait()
+        LOGGER.info("kicking off FGS")
+        yield from bps.kickoff(gridscan, wait=True)
+        gridscan_start_time = time()
+        LOGGER.info("Waiting for Zocalo device queue to have been cleared...")
+        yield from bps.wait(
+            ZOCALO_STAGE_GROUP
+        )  # Make sure ZocaloResults queue is clear and ready to accept our new data
+        if do_during_run:
+            LOGGER.info(f"Running {do_during_run} during FGS")
+            yield from do_during_run()
+        LOGGER.info("completing FGS")
+        yield from bps.complete(gridscan, wait=True)
+
+        # Remove this logging statement once metrics have been added
+        LOGGER.info(
+            f"Gridscan motion program took {round(time()-gridscan_start_time,2)} to complete"
+        )
+
+    yield from do_fgs()
+
+
+def wait_for_gridscan_valid(fgs_motors: FastGridScanCommon, timeout=0.5):
+    LOGGER.info("Waiting for valid fgs_params")
+    SLEEP_PER_CHECK = 0.1
+    times_to_check = int(timeout / SLEEP_PER_CHECK)
+    for _ in range(times_to_check):
+        scan_invalid = yield from bps.rd(fgs_motors.scan_invalid)
+        pos_counter = yield from bps.rd(fgs_motors.position_counter)
+        LOGGER.debug(
+            f"Scan invalid: {scan_invalid} and position counter: {pos_counter}"
+        )
+        if not scan_invalid and pos_counter == 0:
+            LOGGER.info("Gridscan scan valid and position counter reset")
+            return
+        yield from bps.sleep(SLEEP_PER_CHECK)
+    raise WarningException("Scan invalid - pin too long/short/bent and out of range")
+
+
+def set_aperture_for_bbox_size(
+    aperture_device: ApertureScatterguard,
+    bbox_size: list[int] | np.ndarray,
+):
+    # bbox_size is [x,y,z], for i03 we only care about x
+    new_selected_aperture = (
+        ApertureValue.MEDIUM if bbox_size[0] < 2 else ApertureValue.LARGE
+    )
+    LOGGER.info(
+        f"Setting aperture to {new_selected_aperture} based on bounding box size {bbox_size}."
+    )
+
+    @bpp.set_run_key_decorator("change_aperture")
+    @bpp.run_decorator(
+        md={
+            "subplan_name": "change_aperture",
+            "aperture_size": new_selected_aperture.value,
+        }
+    )
+    def set_aperture():
+        yield from bps.abs_set(aperture_device, new_selected_aperture)
+
+    yield from set_aperture()
+
+
+@dataclasses.dataclass
+class _FeatureControlled:
+    class _ZebraSetup(Protocol):
+        def __call__(
+            self, zebra: Zebra, group="setup_zebra_for_gridscan", wait=True
+        ) -> MsgGenerator: ...
+
+    class _ExtraSetup(Protocol):
+        def __call__(
+            self,
+            fgs_composite: FlyScanXRayCentreComposite,
+            parameters: ThreeDGridScan,
+            initial_xyz: np.ndarray,
+        ) -> MsgGenerator: ...
+
+    setup_trigger: _ExtraSetup
+    tidy_plan: Callable[[FlyScanXRayCentreComposite], MsgGenerator]
+    set_flyscan_params: Callable[[], MsgGenerator]
+    fgs_motors: FastGridScanCommon
+
+
+def _get_feature_controlled(
+    fgs_composite: FlyScanXRayCentreComposite,
+    parameters: ThreeDGridScan,
+):
+    if parameters.features.use_panda_for_gridscan:
+        return _FeatureControlled(
+            setup_trigger=_panda_triggering_setup,
+            tidy_plan=_panda_tidy,
+            set_flyscan_params=partial(
+                set_flyscan_params,
+                fgs_composite.panda_fast_grid_scan,
+                parameters.panda_FGS_params,
+            ),
+            fgs_motors=fgs_composite.panda_fast_grid_scan,
+        )
+    else:
+        return _FeatureControlled(
+            setup_trigger=_zebra_triggering_setup,
+            tidy_plan=partial(_generic_tidy, group="flyscan_zebra_tidy", wait=True),
+            set_flyscan_params=partial(
+                set_flyscan_params,
+                fgs_composite.zebra_fast_grid_scan,
+                parameters.FGS_params,
+            ),
+            fgs_motors=fgs_composite.zebra_fast_grid_scan,
+        )
+
+
+def _generic_tidy(
+    fgs_composite: FlyScanXRayCentreComposite, group, wait=True
+) -> MsgGenerator:
+    LOGGER.info("Tidying up Zebra")
+    yield from tidy_up_zebra_after_gridscan(
+        fgs_composite.zebra, fgs_composite.sample_shutter, group=group, wait=wait
+    )
+    LOGGER.info("Tidying up Zocalo")
+    # make sure we don't consume any other results
+    yield from bps.unstage(fgs_composite.zocalo, group=group, wait=wait)
+
+
+def _panda_tidy(fgs_composite: FlyScanXRayCentreComposite):
+    group = "panda_flyscan_tidy"
+    LOGGER.info("Disabling panda blocks")
+    yield from disarm_panda_for_gridscan(fgs_composite.panda, group)
+    yield from _generic_tidy(fgs_composite, group, False)
+    yield from bps.wait(group, timeout=10)
+    yield from bps.unstage(fgs_composite.panda)
+
+
+def _zebra_triggering_setup(
+    fgs_composite: FlyScanXRayCentreComposite,
+    parameters: ThreeDGridScan,
+    initial_xyz: np.ndarray,
+):
+    yield from setup_zebra_for_gridscan(
+        fgs_composite.zebra, fgs_composite.sample_shutter, wait=True
+    )
+
+
+def _panda_triggering_setup(
+    fgs_composite: FlyScanXRayCentreComposite,
+    parameters: ThreeDGridScan,
+    initial_xyz: np.ndarray,
+):
+    LOGGER.info("Setting up Panda for flyscan")
+
+    run_up_distance_mm = yield from bps.rd(
+        fgs_composite.panda_fast_grid_scan.run_up_distance_mm
+    )
+
+    # Set the time between x steps pv
+    DEADTIME_S = 1e-6  # according to https://www.dectris.com/en/detectors/x-ray-detectors/eiger2/eiger2-for-synchrotrons/eiger2-x/
+
+    time_between_x_steps_ms = (DEADTIME_S + parameters.exposure_time_s) * 1e3
+
+    smargon_speed_limit_mm_per_s = yield from bps.rd(
+        fgs_composite.smargon.x.max_velocity
+    )
+
+    sample_velocity_mm_per_s = (
+        parameters.panda_FGS_params.x_step_size * 1e3 / time_between_x_steps_ms
+    )
+    if sample_velocity_mm_per_s > smargon_speed_limit_mm_per_s:
+        raise SmargonSpeedException(
+            f"Smargon speed was calculated from x step size\
+            {parameters.panda_FGS_params.x_step_size} and\
+            time_between_x_steps_ms {time_between_x_steps_ms} as\
+            {sample_velocity_mm_per_s}. The smargon's speed limit is\
+            {smargon_speed_limit_mm_per_s} mm/s."
+        )
+    else:
+        LOGGER.info(
+            f"Panda grid scan: Smargon speed set to {smargon_speed_limit_mm_per_s} mm/s"
+            f" and using a run-up distance of {run_up_distance_mm}"
+        )
+
+    yield from bps.mv(
+        fgs_composite.panda_fast_grid_scan.time_between_x_steps_ms,
+        time_between_x_steps_ms,
+    )
+
+    directory_provider_root = Path(parameters.storage_directory)
+    yield from set_panda_directory(directory_provider_root)
+
+    yield from setup_panda_for_flyscan(
+        fgs_composite.panda,
+        parameters.panda_FGS_params,
+        initial_xyz[0],
+        parameters.exposure_time_s,
+        time_between_x_steps_ms,
+        sample_velocity_mm_per_s,
+    )
+
+    LOGGER.info("Setting up Zebra for panda flyscan")
+    yield from setup_zebra_for_panda_flyscan(
+        fgs_composite.zebra, fgs_composite.sample_shutter, wait=True
+    )