Merge pull request #31 from DiamondLightSource/add_initial_analyserscan

Add initial analyserscan

Author: oliwenmandiamond

.gitignore

@@ -73,3 +73,6 @@ lockfiles/
 
 # ruff cache
 .ruff_cache/
+
+# temp files
+tmp/

pyproject.toml

@@ -17,8 +17,8 @@ dependencies = [
     # it will be auto-pinned to the latest release version by the pre-release workflow
     #
     "bluesky",
-    "dls-dodal",
-    "ophyd-async >= 0.10.0a2",
+    "dls-dodal@git+https://github.com/DiamondLightSource/dodal.git@main",
+    "ophyd-async[sim] >= 0.10.0a2",
     "scanspec",
 ]
 dynamic = ["version"]

src/sm_bluesky/__init__.py

@@ -6,7 +6,7 @@
     Version number as calculated by https://github.com/pypa/setuptools_scm
 """
 
-from . import beamlines, common
+from . import beamlines, common, electron_analyser
 from ._version import __version__
 
-__all__ = ["__version__", "beamlines", "common"]
+__all__ = ["__version__", "beamlines", "common", "electron_analyser"]

tests/common/electron_analyser/__init__.py renamed to src/sm_bluesky/electron_analyser/__init__.py

@@ -0,0 +1,3 @@
+from .analyser_per_step import analyser_nd_step, analyser_shot
+
+__all__ = ["analyser_shot", "analyser_nd_step"]

src/sm_bluesky/common/electron_analyser/energy_source.py renamed to src/sm_bluesky/electron_analyser/energy_source.py

@@ -0,0 +1,81 @@
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from bluesky.plan_stubs import (
+    abs_set,
+    move_per_step,
+    trigger_and_read,
+)
+from bluesky.protocols import Movable, Readable
+from bluesky.utils import (
+    MsgGenerator,
+    plan,
+)
+from dodal.devices.electron_analyser import (
+    ElectronAnalyserRegionDetector,
+    GenericElectronAnalyserRegionDetector,
+)
+from dodal.log import LOGGER
+
+
+@plan
+def analyser_shot(detectors: Sequence[Readable], *args) -> MsgGenerator:
+    yield from analyser_nd_step(detectors, {}, {}, *args)
+
+
+@plan
+def analyser_nd_step(
+    detectors: Sequence[Readable],
+    step: Mapping[Movable, Any],
+    pos_cache: dict[Movable, Any],
+    *args,
+) -> MsgGenerator:
+    """
+    Inner loop of an N-dimensional step scan
+
+    Modified default function for ``per_step`` param in ND plans. Performs an extra for
+    loop for each ElectronAnalyserRegionDetector present so they can be collected one by
+    one.
+
+    Parameters
+    ----------
+    detectors : iterable
+        devices to read
+    step : dict
+        mapping motors to positions in this step
+    pos_cache : dict
+        mapping motors to their last-set positions
+    """
+
+    analyser_detectors: list[GenericElectronAnalyserRegionDetector] = []
+    other_detectors = []
+    for det in detectors:
+        if isinstance(det, ElectronAnalyserRegionDetector):
+            analyser_detectors.append(det)
+        else:
+            other_detectors.append(det)
+
+    # Step provides the map of motors to single position to move to. Move motors to
+    # required positions.
+    yield from move_per_step(step, pos_cache)
+
+    # This is to satisfy type checking. Motors are Moveable and Readable, so make
+    # them Readable so positions can be measured.
+    motors: list[Readable] = [s for s in step.keys() if isinstance(s, Readable)]
+
+    # To get energy sources and open paired shutters, they need to be given in this
+    # plan. They could possibly come from step but we then have to extract them out.
+    # It would also mean forcefully adding in the devices at the wrapper level.
+    # It would easier if they are part of the detector and the plan just calls the
+    # common methods so it is more dynamic and configuration only for device.
+    for analyser_det in analyser_detectors:
+        dets = [analyser_det] + list(other_detectors) + list(motors)
+
+        # This is a work around until we can find a way to get the energy sources to use
+        # the prepare method.
+        yield from abs_set(analyser_det.driver, analyser_det.region)
+        LOGGER.info(f"Scanning region {analyser_det.region.name}.")
+        yield from trigger_and_read(
+            dets,
+            name=analyser_det.region.name,
+        )

src/sm_bluesky/electron_analyser/plan_stubs/__init__.py

@@ -0,0 +1,3 @@
+from .analyser_scans import analysercount, analyserscan, grid_analyserscan
+
+__all__ = ["analysercount", "analyserscan", "grid_analyserscan"]

src/sm_bluesky/electron_analyser/plan_stubs/analyser_per_step.py

@@ -0,0 +1,110 @@
+from collections.abc import Iterable, Sequence
+from typing import Any
+
+from bluesky.plans import PerShot, PerStep, count, grid_scan, scan
+from bluesky.protocols import (
+    Movable,
+    Readable,
+)
+from bluesky.utils import (
+    CustomPlanMetadata,
+    MsgGenerator,
+    ScalarOrIterableFloat,
+    plan,
+)
+from dodal.devices.electron_analyser import (
+    ElectronAnalyserDetector,
+)
+
+from sm_bluesky.electron_analyser.plan_stubs import (
+    analyser_nd_step,
+    analyser_shot,
+)
+
+
+def process_detectors_for_analyserscan(
+    detectors: Sequence[Readable],
+    sequence_file: str,
+) -> Sequence[Readable]:
+    """
+    Check for instance of ElectronAnalyserDetector in the detector list. Provide it with
+    sequence file to read and create list of ElectronAnalyserRegionDetector's. Replace
+    ElectronAnalyserDetector in list of detectors with the
+    list[ElectronAnalyserRegionDetector] and flatten.
+
+    Args:
+        detectors:
+            Devices to measure with for a scan.
+        sequence_file:
+            The file to read to create list[ElectronAnalyserRegionDetector].
+
+    Returns:
+        list of detectors, with any instances of ElectronAnalyserDetector replaced by
+        ElectronAnalyserRegionDetector by the number of regions in the sequence file.
+
+    """
+    analyser_detector = None
+    region_detectors = []
+    for det in detectors:
+        if isinstance(det, ElectronAnalyserDetector):
+            analyser_detector = det
+            region_detectors = det.create_region_detector_list(sequence_file)
+            break
+
+    expansions = (
+        region_detectors if e == analyser_detector else [e] for e in detectors
+    )
+    return [v for vals in expansions for v in vals]
+
+
+def analysercount(
+    detectors: Sequence[Readable],
+    sequence_file: str,
+    num: int = 1,
+    delay: ScalarOrIterableFloat = 0.0,
+    *,
+    per_shot: PerShot | None = None,
+    md: CustomPlanMetadata | None = None,
+) -> MsgGenerator:
+    yield from count(
+        process_detectors_for_analyserscan(detectors, sequence_file),
+        num,
+        delay,
+        per_shot=analyser_shot,
+        md=md,
+    )
+
+
+@plan
+def analyserscan(
+    detectors: Sequence[Readable],
+    sequence_file: str,
+    *args: Movable | Any,
+    num: int | None = None,
+    per_step: PerStep | None = None,
+    md: CustomPlanMetadata | None = None,
+) -> MsgGenerator:
+    yield from scan(
+        process_detectors_for_analyserscan(detectors, sequence_file),
+        *args,
+        num,
+        per_step=analyser_nd_step,
+        md=md,
+    )
+
+
+@plan
+def grid_analyserscan(
+    detectors: Sequence[Readable],
+    sequence_file: str,
+    *args,
+    snake_axes: Iterable | bool | None = None,
+    md: CustomPlanMetadata | None = None,
+) -> MsgGenerator:
+    yield from grid_scan(
+        process_detectors_for_analyserscan(detectors, sequence_file),
+        *args,
+        snake_axes=snake_axes,
+        per_step=analyser_nd_step,
+        md=md,
+    )