ophyd-async 0.9.0a1__py3-none-any.whl → 0.10.0a1__py3-none-any.whl

ophyd_async/plan_stubs/_settings.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable, Mapping
+from typing import Any
+
+import bluesky.plan_stubs as bps
+import numpy as np
+from bluesky.utils import MsgGenerator, plan
+
+from ophyd_async.core import (
+    Device,
+    Settings,
+    SettingsProvider,
+    SignalRW,
+    walk_config_signals,
+    walk_rw_signals,
+)
+from ophyd_async.core._table import Table
+
+from ._utils import T
+from ._wait_for_awaitable import wait_for_awaitable
+
+
+@plan
+def _get_values_of_signals(
+    signals: Mapping[T, SignalRW],
+) -> MsgGenerator[dict[T, Any]]:
+    coros = [sig.get_value() for sig in signals.values()]
+    values = yield from wait_for_awaitable(asyncio.gather(*coros))
+    named_values = dict(zip(signals, values, strict=True))
+    return named_values
+
+
+@plan
+def get_current_settings(
+    device: Device, only_config: bool = False
+) -> MsgGenerator[Settings]:
+    """Get current settings on `Device`.
+
+    If `only_config` is True, get current configuration settings on `Configurable`.
+    """
+    if only_config:
+        signals = yield from wait_for_awaitable(walk_config_signals(device))
+    else:
+        signals = walk_rw_signals(device)
+    named_values = yield from _get_values_of_signals(signals)
+    signal_values = {signals[name]: value for name, value in named_values.items()}
+    return Settings(device, signal_values)
+
+
+@plan
+def store_settings(
+    provider: SettingsProvider, name: str, device: Device, only_config: bool = False
+) -> MsgGenerator[None]:
+    """Walk a Device for SignalRWs and store their values.
+
+    If `only_config` is True, store only configuration settings on `Configurable`.
+
+    :param provider: The provider to store the settings with.
+    :param name: The name to store the settings under.
+    :param device: The Device to walk for SignalRWs.
+    :param only_config: If True, store only configuration settings.
+    """
+    if only_config:
+        signals = yield from wait_for_awaitable(walk_config_signals(device))
+    else:
+        signals = walk_rw_signals(device)
+    named_values = yield from _get_values_of_signals(signals)
+    yield from wait_for_awaitable(provider.store(name, named_values))
+
+
+@plan
+def retrieve_settings(
+    provider: SettingsProvider, name: str, device: Device, only_config: bool = False
+) -> MsgGenerator[Settings]:
+    """Retrieve named Settings for a Device from a provider.
+
+    If `only_config` is True, retrieve only configuration settings on `Configurable`.
+
+    :param provider: The provider to retrieve the settings from.
+    :param name: The name of the settings to retrieve.
+    :param device: The Device to retrieve the settings for.
+    :param only_config: If True, retrieve only configuration settings.
+    """
+    named_values = yield from wait_for_awaitable(provider.retrieve(name))
+    if only_config:
+        signals = yield from wait_for_awaitable(walk_config_signals(device))
+    else:
+        signals = walk_rw_signals(device)
+    unknown_names = set(named_values) - set(signals)
+    if unknown_names:
+        raise NameError(f"Unknown signal names {sorted(unknown_names)}")
+    signal_values = {signals[name]: value for name, value in named_values.items()}
+    return Settings(device, signal_values)
+
+
+@plan
+def apply_settings(settings: Settings) -> MsgGenerator[None]:
+    """Set every SignalRW to the given value in Settings. If value is None ignore it."""
+    signal_values = {
+        signal: value for signal, value in settings.items() if value is not None
+    }
+    if signal_values:
+        for signal, value in signal_values.items():
+            yield from bps.abs_set(signal, value, group="apply_settings")
+        yield from bps.wait("apply_settings")
+
+
+@plan
+def apply_settings_if_different(
+    settings: Settings,
+    apply_plan: Callable[[Settings], MsgGenerator[None]],
+    current_settings: Settings | None = None,
+) -> MsgGenerator[None]:
+    """Set every SignalRW in settings, only if it is different to the current value.
+
+    :param apply_plan:
+        A device specific plan which takes the Settings to apply and applies them to
+        the Device. Used to add device specific ordering to setting the signals.
+    :param current_settings:
+        If given, should be a superset of settings containing the current value of
+        the Settings in the Device. If not given it will be created by reading just
+        the signals given in settings.
+    """
+    if current_settings is None:
+        # If we aren't give the current settings, then get the
+        # values of just the signals we were asked to change.
+        # This allows us to use this plan with Settings for a subset
+        # of signals in the Device without retrieving them all
+        signal_values = yield from _get_values_of_signals(
+            {sig: sig for sig in settings}
+        )
+        current_settings = Settings(settings.device, signal_values)
+
+    def _is_different(current, required) -> bool:
+        if isinstance(current, Table):
+            current = current.model_dump()
+            if isinstance(required, Table):
+                required = required.model_dump()
+            return current.keys() != required.keys() or any(
+                _is_different(current[k], required[k]) for k in current
+            )
+        elif isinstance(current, np.ndarray):
+            return not np.array_equal(current, required)
+        else:
+            return current != required
+
+    settings_to_change, _ = settings.partition(
+        lambda sig: _is_different(current_settings[sig], settings[sig])
+    )
+    yield from apply_plan(settings_to_change)

ophyd_async/plan_stubs/_utils.py

@@ -0,0 +1,3 @@
+from typing import TypeVar
+
+T = TypeVar("T")

ophyd_async/plan_stubs/_wait_for_awaitable.py

@@ -0,0 +1,13 @@
+from collections.abc import Awaitable
+
+import bluesky.plan_stubs as bps
+from bluesky.utils import MsgGenerator, plan
+
+from ._utils import T
+
+
+@plan
+def wait_for_awaitable(coro: Awaitable[T]) -> MsgGenerator[T]:
+    """Wait for a single awaitable to complete, and return the result."""
+    (task,) = yield from bps.wait_for([lambda: coro])
+    return task.result()

ophyd_async/sim/__init__.py

@@ -0,0 +1,29 @@
+"""Some simulated devices to be used in tutorials and testing."""
+
+from ._blob_detector import SimBlobDetector
+from ._mirror_horizontal import HorizontalMirror, HorizontalMirrorDerived
+from ._mirror_vertical import (
+    TwoJackDerived,
+    TwoJackRaw,
+    TwoJackTransform,
+    VerticalMirror,
+)
+from ._motor import FlySimMotorInfo, SimMotor
+from ._pattern_generator import PatternGenerator
+from ._point_detector import SimPointDetector
+from ._stage import SimStage
+
+__all__ = [
+    "SimMotor",
+    "FlySimMotorInfo",
+    "SimStage",
+    "PatternGenerator",
+    "SimPointDetector",
+    "SimBlobDetector",
+    "VerticalMirror",
+    "HorizontalMirror",
+    "HorizontalMirrorDerived",
+    "TwoJackTransform",
+    "TwoJackDerived",
+    "TwoJackRaw",
+]

ophyd_async/sim/__main__.py

@@ -0,0 +1,43 @@
+"""Used for tutorial `Using Devices`."""
+
+# Import bluesky and ophyd
+from tempfile import mkdtemp
+
+import bluesky.plan_stubs as bps  # noqa: F401
+import bluesky.plans as bp  # noqa: F401
+import bluesky.preprocessors as bpp  # noqa: F401
+from bluesky.callbacks.best_effort import BestEffortCallback
+from bluesky.run_engine import RunEngine, autoawait_in_bluesky_event_loop
+
+from ophyd_async import sim
+from ophyd_async.core import StaticPathProvider, UUIDFilenameProvider, init_devices
+
+# Create a run engine and make ipython use it for `await` commands
+RE = RunEngine(call_returns_result=True)
+autoawait_in_bluesky_event_loop()
+
+# Add a callback for plotting
+bec = BestEffortCallback()
+RE.subscribe(bec)
+
+# Make a pattern generator that uses the motor positions
+# to make a test pattern. This simulates the real life process
+# of X-ray scattering off a sample
+pattern_generator = sim.PatternGenerator()
+
+# Make a path provider that makes UUID filenames within a static
+# temporary directory
+path_provider = StaticPathProvider(UUIDFilenameProvider(), mkdtemp())
+
+# All Devices created within this block will be
+# connected and named at the end of the with block
+with init_devices():
+    # Create a sample stage with X and Y motors that report their positions
+    # to the pattern generator
+    stage = sim.SimStage(pattern_generator)
+    # Make a detector device that gives the point value of the pattern generator
+    # when triggered
+    pdet = sim.SimPointDetector(pattern_generator)
+    # Make a detector device that gives a gaussian blob with intensity based
+    # on the point value of the pattern generator when triggered
+    bdet = sim.SimBlobDetector(path_provider, pattern_generator)

ophyd_async/sim/_blob_detector.py

@@ -0,0 +1,33 @@
+from collections.abc import Sequence
+
+from ophyd_async.core import PathProvider, SignalR, StandardDetector
+
+from ._blob_detector_controller import BlobDetectorController
+from ._blob_detector_writer import BlobDetectorWriter
+from ._pattern_generator import PatternGenerator
+
+
+class SimBlobDetector(StandardDetector):
+    """Simulates a detector and writes Blobs to file."""
+
+    def __init__(
+        self,
+        path_provider: PathProvider,
+        pattern_generator: PatternGenerator | None = None,
+        config_sigs: Sequence[SignalR] = (),
+        name: str = "",
+    ) -> None:
+        self.pattern_generator = pattern_generator or PatternGenerator()
+
+        super().__init__(
+            controller=BlobDetectorController(
+                pattern_generator=self.pattern_generator,
+            ),
+            writer=BlobDetectorWriter(
+                pattern_generator=self.pattern_generator,
+                path_provider=path_provider,
+                name_provider=lambda: self.name,
+            ),
+            config_sigs=config_sigs,
+            name=name,
+        )

ophyd_async/sim/_blob_detector_controller.py

@@ -0,0 +1,48 @@
+import asyncio
+from contextlib import suppress
+
+from ophyd_async.core import DetectorController, DetectorTrigger, TriggerInfo
+
+from ._pattern_generator import PatternGenerator
+
+
+class BlobDetectorController(DetectorController):
+    def __init__(self, pattern_generator: PatternGenerator):
+        self.pattern_generator = pattern_generator
+        self.trigger_info: TriggerInfo | None = None
+        self.task: asyncio.Task | None = None
+
+    def get_deadtime(self, exposure):
+        return 0.001
+
+    async def prepare(self, trigger_info: TriggerInfo):
+        # This is a simulation, so only support intenal triggering
+        if trigger_info.trigger != DetectorTrigger.INTERNAL:
+            raise RuntimeError(f"{trigger_info.trigger} not supported by {self}")
+        # Just hold onto the trigger info until we need it
+        self.trigger_info = trigger_info
+
+    async def arm(self):
+        if self.trigger_info is None:
+            raise RuntimeError(f"prepare() not called on {self}")
+        livetime = self.trigger_info.livetime or 0.1
+        # Start a background process off writing the images to file
+        coro = self.pattern_generator.write_images_to_file(
+            exposure=livetime,
+            period=livetime + self.trigger_info.deadtime,
+            number_of_frames=self.trigger_info.total_number_of_triggers,
+        )
+        self.task = asyncio.create_task(coro)
+
+    async def wait_for_idle(self):
+        # Wait for the background task to complete
+        if self.task:
+            await self.task
+
+    async def disarm(self):
+        # Stop the background task and wait for it to finish
+        if self.task:
+            self.task.cancel()
+            with suppress(asyncio.CancelledError):
+                await self.task
+            self.task = None

ophyd_async/sim/_blob_detector_writer.py

@@ -0,0 +1,105 @@
+from collections.abc import AsyncGenerator, AsyncIterator
+from pathlib import Path
+
+import numpy as np
+from bluesky.protocols import Hints, StreamAsset
+from event_model import DataKey
+
+from ophyd_async.core import (
+    DetectorWriter,
+    HDFDatasetDescription,
+    HDFDocumentComposer,
+    NameProvider,
+    PathProvider,
+)
+
+from ._pattern_generator import DATA_PATH, SUM_PATH, PatternGenerator
+
+WIDTH = 320
+HEIGHT = 240
+
+
+class BlobDetectorWriter(DetectorWriter):
+    def __init__(
+        self,
+        pattern_generator: PatternGenerator,
+        path_provider: PathProvider,
+        name_provider: NameProvider,
+    ) -> None:
+        self.pattern_generator = pattern_generator
+        self.path_provider = path_provider
+        self.name_provider = name_provider
+        self.path: Path | None = None
+        self.composer: HDFDocumentComposer | None = None
+        self.datasets: list[HDFDatasetDescription] = []
+
+    async def open(self, multiplier: int = 1) -> dict[str, DataKey]:
+        name = self.name_provider()
+        path_info = self.path_provider(name)
+        self.path = path_info.directory_path / f"{path_info.filename}.h5"
+        self.pattern_generator.open_file(self.path, WIDTH, HEIGHT)
+        # We know it will write data and sum, so emit those
+        self.datasets = [
+            HDFDatasetDescription(
+                data_key=name,
+                dataset=DATA_PATH,
+                shape=(HEIGHT, WIDTH),
+                dtype_numpy=np.dtype(np.uint8).str,
+                chunk_shape=(HEIGHT, WIDTH),
+                multiplier=multiplier,
+            ),
+            HDFDatasetDescription(
+                data_key=f"{name}-sum",
+                dataset=SUM_PATH,
+                shape=(),
+                dtype_numpy=np.dtype(np.int64).str,
+                multiplier=multiplier,
+                chunk_shape=(1024,),
+            ),
+        ]
+        self.composer = None
+        outer_shape = (multiplier,) if multiplier > 1 else ()
+        describe = {
+            ds.data_key: DataKey(
+                source="sim://pattern-generator-hdf-file",
+                shape=list(outer_shape) + list(ds.shape),
+                dtype="array" if ds.shape else "number",
+                external="STREAM:",
+            )
+            for ds in self.datasets
+        }
+        return describe
+
+    @property
+    def hints(self) -> Hints:
+        """The hints to be used for the detector."""
+        return {"fields": [self.name_provider()]}
+
+    async def get_indices_written(self) -> int:
+        return self.pattern_generator.get_last_index()
+
+    async def observe_indices_written(
+        self, timeout: float
+    ) -> AsyncGenerator[int, None]:
+        while True:
+            yield self.pattern_generator.get_last_index()
+            await self.pattern_generator.wait_for_next_index(timeout)
+
+    async def collect_stream_docs(
+        self, indices_written: int
+    ) -> AsyncIterator[StreamAsset]:
+        # When we have written something to the file
+        if indices_written:
+            # Only emit stream resource the first time we see frames in
+            # the file
+            if not self.composer:
+                if not self.path:
+                    raise RuntimeError(f"open() not called on {self}")
+                self.composer = HDFDocumentComposer(self.path, self.datasets)
+                for doc in self.composer.stream_resources():
+                    yield "stream_resource", doc
+            for doc in self.composer.stream_data(indices_written):
+                yield "stream_datum", doc
+
+    async def close(self) -> None:
+        self.pattern_generator.close_file()

ophyd_async/sim/_mirror_horizontal.py

@@ -0,0 +1,46 @@
+import asyncio
+from typing import TypedDict
+
+from bluesky.protocols import Movable
+
+from ophyd_async.core import AsyncStatus, DerivedSignalFactory, Device, soft_signal_rw
+
+from ._mirror_vertical import TwoJackDerived, TwoJackTransform
+from ._motor import SimMotor
+
+
+class HorizontalMirrorDerived(TypedDict):
+    x: float
+    roll: float
+
+
+class HorizontalMirror(Device, Movable):
+    def __init__(self, name=""):
+        # Raw signals
+        self.x1 = SimMotor()
+        self.x2 = SimMotor()
+        # Parameter
+        self.x1_x2_distance = soft_signal_rw(float, initial_value=1)
+        # Derived signals
+        self._factory = DerivedSignalFactory(
+            TwoJackTransform,
+            self._set_mirror,
+            jack1=self.x1,
+            jack2=self.x2,
+            distance=self.x1_x2_distance,
+        )
+        self.x = self._factory.derived_signal_rw(float, "height")
+        self.roll = self._factory.derived_signal_rw(float, "angle")
+        super().__init__(name=name)
+
+    async def _set_mirror(self, derived: TwoJackDerived) -> None:
+        transform = await self._factory.transform()
+        raw = transform.derived_to_raw(**derived)
+        await asyncio.gather(
+            self.x1.set(raw["jack1"]),
+            self.x2.set(raw["jack2"]),
+        )
+
+    @AsyncStatus.wrap
+    async def set(self, value: HorizontalMirrorDerived) -> None:
+        await self._set_mirror(TwoJackDerived(height=value["x"], angle=value["roll"]))

ophyd_async/sim/_mirror_vertical.py

@@ -0,0 +1,74 @@
+import asyncio
+import math
+from typing import TypedDict
+
+from bluesky.protocols import Movable
+
+from ophyd_async.core import (
+    AsyncStatus,
+    DerivedSignalFactory,
+    Device,
+    Transform,
+    soft_signal_rw,
+)
+
+from ._motor import SimMotor
+
+
+class TwoJackRaw(TypedDict):
+    jack1: float
+    jack2: float
+
+
+class TwoJackDerived(TypedDict):
+    height: float
+    angle: float
+
+
+class TwoJackTransform(Transform):
+    distance: float
+
+    def raw_to_derived(self, *, jack1: float, jack2: float) -> TwoJackDerived:
+        diff = jack2 - jack1
+        return TwoJackDerived(
+            height=jack1 + diff / 2,
+            # need the cast as returns numpy float rather than float64, but this
+            # is ok at runtime
+            angle=math.atan(diff / self.distance),
+        )
+
+    def derived_to_raw(self, *, height: float, angle: float) -> TwoJackRaw:
+        diff = math.tan(angle) * self.distance
+        return TwoJackRaw(
+            jack1=height - diff / 2,
+            jack2=height + diff / 2,
+        )
+
+
+class VerticalMirror(Device, Movable[TwoJackDerived]):
+    def __init__(self, name=""):
+        # Raw signals
+        self.y1 = SimMotor()
+        self.y2 = SimMotor()
+        # Parameter
+        self.y1_y2_distance = soft_signal_rw(float, initial_value=1)
+        # Derived signals
+        self._factory = DerivedSignalFactory(
+            TwoJackTransform,
+            self.set,
+            jack1=self.y1,
+            jack2=self.y2,
+            distance=self.y1_y2_distance,
+        )
+        self.height = self._factory.derived_signal_rw(float, "height")
+        self.angle = self._factory.derived_signal_rw(float, "angle")
+        super().__init__(name=name)
+
+    @AsyncStatus.wrap
+    async def set(self, derived: TwoJackDerived) -> None:  # type: ignore until bluesky 1.13.2
+        transform = await self._factory.transform()
+        raw = transform.derived_to_raw(**derived)
+        await asyncio.gather(
+            self.y1.set(raw["jack1"]),
+            self.y2.set(raw["jack2"]),
+        )