ophyd-async 0.1.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

ophyd_async/core/detector.py

@@ -0,0 +1,341 @@
+"""Module which defines abstract classes to work with detectors"""
+
+import asyncio
+import time
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import (
+    AsyncGenerator,
+    AsyncIterator,
+    Callable,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    Sequence,
+    TypeVar,
+)
+
+from bluesky.protocols import (
+    Collectable,
+    DataKey,
+    Flyable,
+    Preparable,
+    Reading,
+    Stageable,
+    StreamAsset,
+    Triggerable,
+    WritesStreamAssets,
+)
+
+from ophyd_async.protocols import AsyncConfigurable, AsyncReadable
+
+from .async_status import AsyncStatus, WatchableAsyncStatus
+from .device import Device
+from .utils import DEFAULT_TIMEOUT, WatcherUpdate, merge_gathered_dicts
+
+T = TypeVar("T")
+
+
+class DetectorTrigger(str, Enum):
+    """Type of mechanism for triggering a detector to take frames"""
+
+    #: Detector generates internal trigger for given rate
+    internal = "internal"
+    #: Expect a series of arbitrary length trigger signals
+    edge_trigger = "edge_trigger"
+    #: Expect a series of constant width external gate signals
+    constant_gate = "constant_gate"
+    #: Expect a series of variable width external gate signals
+    variable_gate = "variable_gate"
+
+
+@dataclass(frozen=True)
+class TriggerInfo:
+    """Minimal set of information required to setup triggering on a detector"""
+
+    #: Number of triggers that will be sent
+    num: int
+    #: Sort of triggers that will be sent
+    trigger: DetectorTrigger
+    #: What is the minimum deadtime between triggers
+    deadtime: float
+    #: What is the maximum high time of the triggers
+    livetime: float
+
+
+class DetectorControl(ABC):
+    """
+    Classes implementing this interface should hold the logic for
+    arming and disarming a detector
+    """
+
+    @abstractmethod
+    def get_deadtime(self, exposure: float) -> float:
+        """For a given exposure, how long should the time between exposures be"""
+
+    @abstractmethod
+    async def arm(
+        self,
+        num: int,
+        trigger: DetectorTrigger = DetectorTrigger.internal,
+        exposure: Optional[float] = None,
+    ) -> AsyncStatus:
+        """
+        Arm detector, do all necessary steps to prepare detector for triggers.
+
+        Args:
+            num: Expected number of frames
+            trigger: Type of trigger for which to prepare the detector. Defaults to
+            DetectorTrigger.internal.
+            exposure: Exposure time with which to set up the detector. Defaults to None
+            if not applicable or the detector is expected to use its previously-set
+            exposure time.
+
+        Returns:
+            AsyncStatus: Status representing the arm operation. This function returning
+            represents the start of the arm. The returned status completing means
+            the detector is now armed.
+        """
+
+    @abstractmethod
+    async def disarm(self):
+        """Disarm the detector, return detector to an idle state"""
+
+
+class DetectorWriter(ABC):
+    """Logic for making a detector write data to somewhere persistent
+    (e.g. an HDF5 file)"""
+
+    @abstractmethod
+    async def open(self, multiplier: int = 1) -> Dict[str, DataKey]:
+        """Open writer and wait for it to be ready for data.
+
+        Args:
+            multiplier: Each StreamDatum index corresponds to this many
+                written exposures
+
+        Returns:
+            Output for ``describe()``
+        """
+
+    @abstractmethod
+    def observe_indices_written(
+        self, timeout=DEFAULT_TIMEOUT
+    ) -> AsyncGenerator[int, None]:
+        """Yield the index of each frame (or equivalent data point) as it is written"""
+
+    @abstractmethod
+    async def get_indices_written(self) -> int:
+        """Get the number of indices written"""
+
+    @abstractmethod
+    def collect_stream_docs(self, indices_written: int) -> AsyncIterator[StreamAsset]:
+        """Create Stream docs up to given number written"""
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Close writer, blocks until I/O is complete"""
+
+
+class StandardDetector(
+    Device,
+    Stageable,
+    AsyncConfigurable,
+    AsyncReadable,
+    Triggerable,
+    Preparable,
+    Flyable,
+    Collectable,
+    WritesStreamAssets,
+    Generic[T],
+):
+    """
+    Useful detector base class for step and fly scanning detectors.
+    Aggregates controller and writer logic together.
+    """
+
+    def __init__(
+        self,
+        controller: DetectorControl,
+        writer: DetectorWriter,
+        config_sigs: Sequence[AsyncReadable] = (),
+        name: str = "",
+        writer_timeout: float = DEFAULT_TIMEOUT,
+    ) -> None:
+        """
+        Constructor
+
+        Args:
+            controller: Logic for arming and disarming the detector
+            writer: Logic for making the detector write persistent data
+            config_sigs: Signals to read when describe and read
+            configuration are called. Defaults to ().
+            name: Device name. Defaults to "".
+            writer_timeout: Timeout for frame writing to start, if the
+            timeout is reached, ophyd-async assumes the detector
+            has a problem and raises an error.
+            Defaults to DEFAULT_TIMEOUT.
+        """
+        self._controller = controller
+        self._writer = writer
+        self._describe: Dict[str, DataKey] = {}
+        self._config_sigs = list(config_sigs)
+        self._frame_writing_timeout = writer_timeout
+        # For prepare
+        self._arm_status: Optional[AsyncStatus] = None
+        self._trigger_info: Optional[TriggerInfo] = None
+        # For kickoff
+        self._watchers: List[Callable] = []
+        self._fly_status: Optional[WatchableAsyncStatus] = None
+        self._fly_start: float
+
+        self._intial_frame: int
+        self._last_frame: int
+        super().__init__(name)
+
+    @property
+    def controller(self) -> DetectorControl:
+        return self._controller
+
+    @property
+    def writer(self) -> DetectorWriter:
+        return self._writer
+
+    @AsyncStatus.wrap
+    async def stage(self) -> None:
+        # Disarm the detector, stop filewriting, and open file for writing.
+        await self._check_config_sigs()
+        await asyncio.gather(self.writer.close(), self.controller.disarm())
+        self._describe = await self.writer.open()
+
+    async def _check_config_sigs(self):
+        """Checks configuration signals are named and connected."""
+        for signal in self._config_sigs:
+            if signal.name == "":
+                raise Exception(
+                    "config signal must be named before it is passed to the detector"
+                )
+            try:
+                await signal.get_value()
+            except NotImplementedError:
+                raise Exception(
+                    f"config signal {signal._name} must be connected before it is "
+                    + "passed to the detector"
+                )
+
+    @AsyncStatus.wrap
+    async def unstage(self) -> None:
+        # Stop data writing.
+        await self.writer.close()
+
+    async def read_configuration(self) -> Dict[str, Reading]:
+        return await merge_gathered_dicts(sig.read() for sig in self._config_sigs)
+
+    async def describe_configuration(self) -> Dict[str, DataKey]:
+        return await merge_gathered_dicts(sig.describe() for sig in self._config_sigs)
+
+    async def read(self) -> Dict[str, Reading]:
+        # All data is in StreamResources, not Events, so nothing to output here
+        return {}
+
+    async def describe(self) -> Dict[str, DataKey]:
+        return self._describe
+
+    @AsyncStatus.wrap
+    async def trigger(self) -> None:
+        # Arm the detector and wait for it to finish.
+        indices_written = await self.writer.get_indices_written()
+        written_status = await self.controller.arm(
+            num=1,
+            trigger=DetectorTrigger.internal,
+        )
+        await written_status
+        end_observation = indices_written + 1
+
+        async for index in self.writer.observe_indices_written(
+            self._frame_writing_timeout
+        ):
+            if index >= end_observation:
+                break
+
+    def prepare(
+        self,
+        value: T,
+    ) -> AsyncStatus:
+        # Just arm detector for the time being
+        return AsyncStatus(self._prepare(value))
+
+    async def _prepare(self, value: T) -> None:
+        """
+        Arm detector.
+
+        Prepare the detector with trigger information. This is determined at and passed
+        in from the plan level.
+
+        This currently only prepares detectors for flyscans and stepscans just use the
+        trigger information determined in trigger.
+
+        To do: Unify prepare to be use for both fly and step scans.
+
+        Args:
+            value: TriggerInfo describing how to trigger the detector
+        """
+        assert type(value) is TriggerInfo
+        self._trigger_info = value
+        self._initial_frame = await self.writer.get_indices_written()
+        self._last_frame = self._initial_frame + self._trigger_info.num
+
+        required = self.controller.get_deadtime(self._trigger_info.livetime)
+        assert required <= self._trigger_info.deadtime, (
+            f"Detector {self.controller} needs at least {required}s deadtime, "
+            f"but trigger logic provides only {self._trigger_info.deadtime}s"
+        )
+        self._arm_status = await self.controller.arm(
+            num=self._trigger_info.num,
+            trigger=self._trigger_info.trigger,
+            exposure=self._trigger_info.livetime,
+        )
+        self._fly_start = time.monotonic()
+
+    @AsyncStatus.wrap
+    async def kickoff(self):
+        if not self._arm_status:
+            raise Exception("Detector not armed!")
+
+    @WatchableAsyncStatus.wrap
+    async def complete(self):
+        assert self._arm_status, "Prepare not run"
+        assert self._trigger_info
+        async for index in self.writer.observe_indices_written(
+            self._frame_writing_timeout
+        ):
+            yield WatcherUpdate(
+                name=self.name,
+                current=index,
+                initial=self._initial_frame,
+                target=self._trigger_info.num,
+                unit="",
+                precision=0,
+                time_elapsed=time.monotonic() - self._fly_start,
+            )
+            if index >= self._trigger_info.num:
+                break
+
+    async def describe_collect(self) -> Dict[str, DataKey]:
+        return self._describe
+
+    async def collect_asset_docs(
+        self, index: Optional[int] = None
+    ) -> AsyncIterator[StreamAsset]:
+        # Collect stream datum documents for all indices written.
+        # The index is optional, and provided for fly scans, however this needs to be
+        # retrieved for step scans.
+        if index is None:
+            index = await self.writer.get_indices_written()
+        async for doc in self.writer.collect_stream_docs(index):
+            yield doc
+
+    async def get_index(self) -> int:
+        return await self.writer.get_indices_written()

ophyd_async/core/device.py

@@ -0,0 +1,226 @@
+"""Base device"""
+
+import asyncio
+import sys
+from functools import cached_property
+from logging import LoggerAdapter, getLogger
+from typing import (
+    Any,
+    Coroutine,
+    Dict,
+    Generator,
+    Iterator,
+    Optional,
+    Set,
+    Tuple,
+    TypeVar,
+)
+
+from bluesky.protocols import HasName
+from bluesky.run_engine import call_in_bluesky_event_loop
+
+from .utils import DEFAULT_TIMEOUT, NotConnected, wait_for_connection
+
+
+class Device(HasName):
+    """Common base class for all Ophyd Async Devices.
+
+    By default, names and connects all Device children.
+    """
+
+    _name: str = ""
+    #: The parent Device if it exists
+    parent: Optional["Device"] = None
+    # None if connect hasn't started, a Task if it has
+    _connect_task: Optional[asyncio.Task] = None
+    _connect_mock_arg: bool = False
+
+    def __init__(self, name: str = "") -> None:
+        self.set_name(name)
+
+    @property
+    def name(self) -> str:
+        """Return the name of the Device"""
+        return self._name
+
+    @cached_property
+    def log(self):
+        return LoggerAdapter(
+            getLogger("ophyd_async.devices"), {"ophyd_async_device_name": self.name}
+        )
+
+    def children(self) -> Iterator[Tuple[str, "Device"]]:
+        for attr_name, attr in self.__dict__.items():
+            if attr_name != "parent" and isinstance(attr, Device):
+                yield attr_name, attr
+
+    def set_name(self, name: str):
+        """Set ``self.name=name`` and each ``self.child.name=name+"-child"``.
+
+        Parameters
+        ----------
+        name:
+            New name to set
+        """
+
+        # Ensure self.log is recreated after a name change
+        if hasattr(self, "log"):
+            del self.log
+
+        self._name = name
+        for attr_name, child in self.children():
+            child_name = f"{name}-{attr_name.rstrip('_')}" if name else ""
+            child.set_name(child_name)
+            child.parent = self
+
+    async def connect(
+        self,
+        mock: bool = False,
+        timeout: float = DEFAULT_TIMEOUT,
+        force_reconnect: bool = False,
+    ):
+        """Connect self and all child Devices.
+
+        Contains a timeout that gets propagated to child.connect methods.
+
+        Parameters
+        ----------
+        mock:
+            If True then use ``MockSignalBackend`` for all Signals
+        timeout:
+            Time to wait before failing with a TimeoutError.
+        """
+        # If previous connect with same args has started and not errored, can use it
+        can_use_previous_connect = (
+            self._connect_task
+            and not (self._connect_task.done() and self._connect_task.exception())
+            and self._connect_mock_arg == mock
+        )
+        if force_reconnect or not can_use_previous_connect:
+            # Kick off a connection
+            coros = {
+                name: child_device.connect(
+                    mock, timeout=timeout, force_reconnect=force_reconnect
+                )
+                for name, child_device in self.children()
+            }
+            self._connect_task = asyncio.create_task(wait_for_connection(**coros))
+            self._connect_mock_arg = mock
+
+        assert self._connect_task, "Connect task not created, this shouldn't happen"
+        # Wait for it to complete
+        await self._connect_task
+
+
+VT = TypeVar("VT", bound=Device)
+
+
+class DeviceVector(Dict[int, VT], Device):
+    """
+    Defines device components with indices.
+
+    In the below example, foos becomes a dictionary on the parent device
+    at runtime, so parent.foos[2] returns a FooDevice. For example usage see
+    :class:`~ophyd_async.epics.demo.DynamicSensorGroup`
+    """
+
+    def children(self) -> Generator[Tuple[str, Device], None, None]:
+        for attr_name, attr in self.items():
+            if isinstance(attr, Device):
+                yield str(attr_name), attr
+
+
+class DeviceCollector:
+    """Collector of top level Device instances to be used as a context manager
+
+    Parameters
+    ----------
+    set_name:
+        If True, call ``device.set_name(variable_name)`` on all collected
+        Devices
+    connect:
+        If True, call ``device.connect(mock)`` in parallel on all
+        collected Devices
+    mock:
+        If True, connect Signals in simulation mode
+    timeout:
+        How long to wait for connect before logging an exception
+
+    Notes
+    -----
+    Example usage::
+
+        [async] with DeviceCollector():
+            t1x = motor.Motor("BLxxI-MO-TABLE-01:X")
+            t1y = motor.Motor("pva://BLxxI-MO-TABLE-01:Y")
+            # Names and connects devices here
+        assert t1x.comm.velocity.source
+        assert t1x.name == "t1x"
+
+    """
+
+    def __init__(
+        self,
+        set_name=True,
+        connect=True,
+        mock=False,
+        timeout: float = 10.0,
+    ):
+        self._set_name = set_name
+        self._connect = connect
+        self._mock = mock
+        self._timeout = timeout
+        self._names_on_enter: Set[str] = set()
+        self._objects_on_exit: Dict[str, Any] = {}
+
+    def _caller_locals(self):
+        """Walk up until we find a stack frame that doesn't have us as self"""
+        try:
+            raise ValueError
+        except ValueError:
+            _, _, tb = sys.exc_info()
+            assert tb, "Can't get traceback, this shouldn't happen"
+            caller_frame = tb.tb_frame
+            while caller_frame.f_locals.get("self", None) is self:
+                caller_frame = caller_frame.f_back
+            return caller_frame.f_locals
+
+    def __enter__(self) -> "DeviceCollector":
+        # Stash the names that were defined before we were called
+        self._names_on_enter = set(self._caller_locals())
+        return self
+
+    async def __aenter__(self) -> "DeviceCollector":
+        return self.__enter__()
+
+    async def _on_exit(self) -> None:
+        # Name and kick off connect for devices
+        connect_coroutines: Dict[str, Coroutine] = {}
+        for name, obj in self._objects_on_exit.items():
+            if name not in self._names_on_enter and isinstance(obj, Device):
+                if self._set_name and not obj.name:
+                    obj.set_name(name)
+                if self._connect:
+                    connect_coroutines[name] = obj.connect(
+                        self._mock, timeout=self._timeout
+                    )
+
+        # Connect to all the devices
+        if connect_coroutines:
+            await wait_for_connection(**connect_coroutines)
+
+    async def __aexit__(self, type, value, traceback):
+        self._objects_on_exit = self._caller_locals()
+        await self._on_exit()
+
+    def __exit__(self, type_, value, traceback):
+        self._objects_on_exit = self._caller_locals()
+        try:
+            fut = call_in_bluesky_event_loop(self._on_exit())
+        except RuntimeError:
+            raise NotConnected(
+                "Could not connect devices. Is the bluesky event loop running? See "
+                "https://blueskyproject.io/ophyd-async/main/"
+                "user/explanations/event-loop-choice.html for more info."
+            )
+        return fut