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

ophyd_async/core/_detector.py

@@ -1,14 +1,20 @@
-"""Module which defines abstract classes to work with detectors"""
+"""Module which defines abstract classes to work with detectors."""
 
 import asyncio
 import time
 from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterator, Sequence
+from enum import Enum
 from functools import cached_property
+from typing import (
+    Generic,
+    TypeVar,
+)
 
 from bluesky.protocols import (
     Collectable,
     Flyable,
+    Hints,
     Preparable,
     Reading,
     Stageable,
@@ -23,48 +29,62 @@ from ._device import Device, DeviceConnector
 from ._protocol import AsyncConfigurable, AsyncReadable
 from ._signal import SignalR
 from ._status import AsyncStatus, WatchableAsyncStatus
-from ._utils import DEFAULT_TIMEOUT, StrictEnum, WatcherUpdate, merge_gathered_dicts
+from ._utils import DEFAULT_TIMEOUT, WatcherUpdate, merge_gathered_dicts
+
+
+class DetectorTrigger(Enum):
+    """Type of mechanism for triggering a detector to take frames."""
+
+    INTERNAL = "INTERNAL"
+    """Detector generates internal trigger for given rate"""
 
+    EDGE_TRIGGER = "EDGE_TRIGGER"
+    """Expect a series of arbitrary length trigger signals"""
 
-class DetectorTrigger(StrictEnum):
-    """Type of mechanism for triggering a detector to take frames"""
+    CONSTANT_GATE = "CONSTANT_GATE"
+    """Expect a series of constant width external gate signals"""
 
-    #: 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"
+    VARIABLE_GATE = "VARIABLE_GATE"
+    """Expect a series of variable width external gate signals"""
 
 
 class TriggerInfo(BaseModel):
-    """Minimal set of information required to setup triggering on a detector"""
-
-    #: Number of triggers that will be sent, (0 means infinite) Can be:
-    #  - A single integer or
-    #  - A list of integers for multiple triggers
-    # Example for tomography: TriggerInfo(number=[2,3,100,3])
-    #:     This would trigger:
-    #:     - 2 times for dark field images
-    #:     - 3 times for initial flat field images
-    #:     - 100 times for projections
-    #:     - 3 times for final flat field images
-    number_of_triggers: NonNegativeInt | list[NonNegativeInt]
-    #: Sort of triggers that will be sent
+    """Minimal set of information required to setup triggering on a detector."""
+
+    number_of_triggers: NonNegativeInt | list[NonNegativeInt] = Field(default=1)
+    """Number of triggers that will be sent, (0 means infinite).
+
+    Can be:
+    - A single integer or
+    - A list of integers for multiple triggers
+
+    Example for tomography: ``TriggerInfo(number=[2,3,100,3])``.
+    This would trigger:
+
+    - 2 times for dark field images
+    - 3 times for initial flat field images
+    - 100 times for projections
+    - 3 times for final flat field images
+    """
+
     trigger: DetectorTrigger = Field(default=DetectorTrigger.INTERNAL)
-    #: What is the minimum deadtime between triggers
-    deadtime: float | None = Field(default=None, ge=0)
-    #: What is the maximum high time of the triggers
+    """Sort of triggers that will be sent"""
+
+    deadtime: float = Field(default=0.0, ge=0)
+    """What is the minimum deadtime between triggers"""
+
     livetime: float | None = Field(default=None, ge=0)
-    #: What is the maximum timeout on waiting for a frame
+    """What is the maximum high time of the triggers"""
+
     frame_timeout: float | None = Field(default=None, gt=0)
-    #: How many triggers make up a single StreamDatum index, to allow multiple frames
-    #: from a faster detector to be zipped with a single frame from a slow detector
-    #: e.g. if num=10 and multiplier=5 then the detector will take 10 frames,
-    #: but publish 2 indices, and describe() will show a shape of (5, h, w)
+    """What is the maximum timeout on waiting for a frame"""
+
     multiplier: int = 1
+    """How many triggers make up a single StreamDatum index, to allow multiple frames
+    from a faster detector to be zipped with a single frame from a slow detector
+    e.g. if num=10 and multiplier=5 then the detector will take 10 frames,
+    but publish 2 indices, and describe() will show a shape of (5, h, w)
+    """
 
     @computed_field
     @cached_property
@@ -77,85 +97,81 @@ class TriggerInfo(BaseModel):
 
 
 class DetectorController(ABC):
-    """
-    Classes implementing this interface should hold the logic for
-    arming and disarming a detector
-    """
+    """Detector logic for arming and disarming the detector."""
 
     @abstractmethod
     def get_deadtime(self, exposure: float | None) -> float:
-        """For a given exposure, how long should the time between exposures be"""
+        """For a given exposure, how long should the time between exposures be."""
 
     @abstractmethod
-    async def prepare(self, trigger_info: TriggerInfo):
-        """
-        Do all necessary steps to prepare the detector for triggers.
-
-        Args:
-            trigger_info: This is a Pydantic model which contains
-                number Expected number of frames.
-                trigger Type of trigger for which to prepare the detector. Defaults
-                to DetectorTrigger.internal.
-                livetime Livetime / 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.
-                deadtime Defaults to None. This is the minimum deadtime between
-                triggers.
-                multiplier The number of triggers grouped into a single StreamDatum
-                index.
+    async def prepare(self, trigger_info: TriggerInfo) -> None:
+        """Do all necessary steps to prepare the detector for triggers.
+
+        :param trigger_info: The sort of triggers to expect.
         """
 
     @abstractmethod
     async def arm(self) -> None:
-        """
-        Arm the detector
-        """
+        """Arm the detector."""
 
     @abstractmethod
     async def wait_for_idle(self):
-        """
-        This will wait on the internal _arm_status and wait for it to get disarmed/idle
-        """
+        """Wait on the internal _arm_status and wait for it to get disarmed/idle."""
 
     @abstractmethod
     async def disarm(self):
-        """Disarm the detector, return detector to an idle state"""
+        """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)"""
+    """Logic for making detector write data to somewhere persistent (e.g. 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()``
+        :param multiplier:
+            Each StreamDatum index corresponds to this many written exposures
+        :return: 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"""
+    @property
+    def hints(self) -> Hints:
+        """The hints to be used for the detector."""
+        return {}
 
     @abstractmethod
     async def get_indices_written(self) -> int:
-        """Get the number of indices written"""
+        """Get the number of indices written."""
+
+    # Note: this method is really async, but if we make it async here then we
+    # need to give it a body with a redundant yield statement, which is a bit
+    # awkward. So we just leave it as a regular method and let the user
+    # implement it as async.
+    @abstractmethod
+    def observe_indices_written(self, timeout: float) -> AsyncGenerator[int, None]:
+        """Yield the index of each frame (or equivalent data point) as it is written."""
 
     @abstractmethod
     def collect_stream_docs(self, indices_written: int) -> AsyncIterator[StreamAsset]:
-        """Create Stream docs up to given number written"""
+        """Create Stream docs up to given number written."""
 
     @abstractmethod
     async def close(self) -> None:
-        """Close writer, blocks until I/O is complete"""
+        """Close writer, blocks until I/O is complete."""
+
+
+# Add type var for controller so we can define
+# StandardDetector[KinetixController, ADWriter] for example
+DetectorControllerT = TypeVar("DetectorControllerT", bound=DetectorController)
+DetectorWriterT = TypeVar("DetectorWriterT", bound=DetectorWriter)
+
+
+def _ensure_trigger_info_exists(trigger_info: TriggerInfo | None) -> TriggerInfo:
+    # make absolute sure we realy have a valid TriggerInfo ... mostly for pylance
+    if trigger_info is None:
+        raise RuntimeError("Trigger info must be set before calling this method.")
+    return trigger_info
 
 
 class StandardDetector(
@@ -168,30 +184,26 @@ class StandardDetector(
     Flyable,
     Collectable,
     WritesStreamAssets,
+    Generic[DetectorControllerT, DetectorWriterT],
 ):
-    """
-    Useful detector base class for step and fly scanning detectors.
+    """Detector base class for step and fly scanning detectors.
+
     Aggregates controller and writer logic together.
+
+    :param controller: Logic for arming and disarming the detector
+    :param writer: Logic for making the detector write persistent data
+    :param config_sigs: Signals to read when describe and read configuration are called
+    :param name: Device name
     """
 
     def __init__(
         self,
-        controller: DetectorController,
-        writer: DetectorWriter,
+        controller: DetectorControllerT,
+        writer: DetectorWriterT,
         config_sigs: Sequence[SignalR] = (),
         name: str = "",
         connector: DeviceConnector | None = None,
     ) -> 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 "".
-        """
         self._controller = controller
         self._writer = writer
         self._describe: dict[str, DataKey] = {}
@@ -211,23 +223,15 @@ class StandardDetector(
         self._initial_frame: int = 0
         super().__init__(name, connector=connector)
 
-    @property
-    def controller(self) -> DetectorController:
-        return self._controller
-
-    @property
-    def writer(self) -> DetectorWriter:
-        return self._writer
-
     @AsyncStatus.wrap
     async def stage(self) -> None:
-        # Disarm the detector, stop file writing.
+        """Make sure the detector is idle and ready to be used."""
         await self._check_config_sigs()
-        await asyncio.gather(self.writer.close(), self.controller.disarm())
+        await asyncio.gather(self._writer.close(), self._controller.disarm())
         self._trigger_info = None
 
     async def _check_config_sigs(self):
-        """Checks configuration signals are named and connected."""
+        """Check configuration signals are named and connected."""
         for signal in self._config_sigs:
             if signal.name == "":
                 raise Exception(
@@ -243,8 +247,8 @@ class StandardDetector(
 
     @AsyncStatus.wrap
     async def unstage(self) -> None:
-        # Stop data writing.
-        await asyncio.gather(self.writer.close(), self.controller.disarm())
+        """Disarm the detector and stop file writing."""
+        await asyncio.gather(self._writer.close(), self._controller.disarm())
 
     async def read_configuration(self) -> dict[str, Reading]:
         return await merge_gathered_dicts(sig.read() for sig in self._config_sigs)
@@ -253,6 +257,7 @@ class StandardDetector(
         return await merge_gathered_dicts(sig.describe() for sig in self._config_sigs)
 
     async def read(self) -> dict[str, Reading]:
+        """There is no data to be placed in events, so this is empty."""
         # All data is in StreamResources, not Events, so nothing to output here
         return {}
 
@@ -266,71 +271,67 @@ class StandardDetector(
                 TriggerInfo(
                     number_of_triggers=1,
                     trigger=DetectorTrigger.INTERNAL,
-                    deadtime=None,
-                    livetime=None,
-                    frame_timeout=None,
                 )
             )
-        assert self._trigger_info
-        assert self._trigger_info.trigger is DetectorTrigger.INTERNAL
+
+        trigger_info = _ensure_trigger_info_exists(self._trigger_info)
+        if trigger_info.trigger is not DetectorTrigger.INTERNAL:
+            msg = "The trigger method can only be called with INTERNAL triggering"
+            raise ValueError(msg)
+
         # Arm the detector and wait for it to finish.
-        indices_written = await self.writer.get_indices_written()
-        await self.controller.arm()
-        await self.controller.wait_for_idle()
+        indices_written = await self._writer.get_indices_written()
+        await self._controller.arm()
+        await self._controller.wait_for_idle()
         end_observation = indices_written + 1
 
-        async for index in self.writer.observe_indices_written(
-            DEFAULT_TIMEOUT
-            + (self._trigger_info.livetime or 0)
-            + (self._trigger_info.deadtime or 0)
+        async for index in self._writer.observe_indices_written(
+            DEFAULT_TIMEOUT + (trigger_info.livetime or 0) + trigger_info.deadtime
         ):
             if index >= end_observation:
                 break
 
     @AsyncStatus.wrap
     async def prepare(self, value: TriggerInfo) -> None:
-        """
-        Arm detector.
+        """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
+        :param value: TriggerInfo describing how to trigger the detector
         """
-        if value.trigger != DetectorTrigger.INTERNAL:
-            assert (
-                value.deadtime
-            ), "Deadtime must be supplied when in externally triggered mode"
-        if value.deadtime:
-            required = self.controller.get_deadtime(value.livetime)
-            assert required <= value.deadtime, (
-                f"Detector {self.controller} needs at least {required}s deadtime, "
-                f"but trigger logic provides only {value.deadtime}s"
+        if value.trigger != DetectorTrigger.INTERNAL and not value.deadtime:
+            msg = "Deadtime must be supplied when in externally triggered mode"
+            raise ValueError(msg)
+        required_deadtime = self._controller.get_deadtime(value.livetime)
+        if value.deadtime and required_deadtime > value.deadtime:
+            msg = (
+                f"Detector {self._controller} needs at least {required_deadtime}s "
+                f"deadtime, but trigger logic provides only {value.deadtime}s"
             )
-        self._trigger_info = value
+            raise ValueError(msg)
+        elif not value.deadtime:
+            value.deadtime = self._controller.get_deadtime(value.livetime)
         self._number_of_triggers_iter = iter(
-            self._trigger_info.number_of_triggers
-            if isinstance(self._trigger_info.number_of_triggers, list)
-            else [self._trigger_info.number_of_triggers]
+            value.number_of_triggers
+            if isinstance(value.number_of_triggers, list)
+            else [value.number_of_triggers]
         )
-        self._initial_frame = await self.writer.get_indices_written()
         self._describe, _ = await asyncio.gather(
-            self.writer.open(value.multiplier), self.controller.prepare(value)
+            self._writer.open(value.multiplier), self._controller.prepare(value)
         )
+        self._initial_frame = await self._writer.get_indices_written()
         if value.trigger != DetectorTrigger.INTERNAL:
-            await self.controller.arm()
-            self._fly_start = time.monotonic()
+            await self._controller.arm()
+        self._trigger_info = value
 
     @AsyncStatus.wrap
     async def kickoff(self):
         if self._trigger_info is None or self._number_of_triggers_iter is None:
             raise RuntimeError("Prepare must be called before kickoff!")
+        if self._trigger_info.trigger == DetectorTrigger.INTERNAL:
+            await self._controller.arm()
+        self._fly_start = time.monotonic()
         try:
             self._frames_to_complete = next(self._number_of_triggers_iter)
             self._completable_frames += self._frames_to_complete
@@ -342,13 +343,13 @@ class StandardDetector(
 
     @WatchableAsyncStatus.wrap
     async def complete(self):
-        assert self._trigger_info
-        indices_written = self.writer.observe_indices_written(
-            self._trigger_info.frame_timeout
+        trigger_info = _ensure_trigger_info_exists(self._trigger_info)
+        indices_written = self._writer.observe_indices_written(
+            trigger_info.frame_timeout
             or (
                 DEFAULT_TIMEOUT
-                + (self._trigger_info.livetime or 0)
-                + (self._trigger_info.deadtime or 0)
+                + (trigger_info.livetime or 0)
+                + (trigger_info.deadtime or 0)
             )
         )
         try:
@@ -368,11 +369,11 @@ class StandardDetector(
                     break
         finally:
             await indices_written.aclose()
-            if self._completable_frames >= self._trigger_info.total_number_of_triggers:
+            if self._completable_frames >= trigger_info.total_number_of_triggers:
                 self._completable_frames = 0
                 self._frames_to_complete = 0
                 self._number_of_triggers_iter = None
-                await self.controller.wait_for_idle()
+                await self._controller.wait_for_idle()
 
     async def describe_collect(self) -> dict[str, DataKey]:
         return self._describe
@@ -384,9 +385,13 @@ class StandardDetector(
         # 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):
+            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()
+        return await self._writer.get_indices_written()
+
+    @property
+    def hints(self) -> Hints:
+        return self._writer.hints