ophyd-async 0.3.4a1__py3-none-any.whl → 0.4.0__py3-none-any.whl

ophyd_async/epics/motion/motor.py

@@ -1,6 +1,8 @@
 import asyncio
+from typing import Optional
 
-from bluesky.protocols import Movable, Stoppable
+from bluesky.protocols import Flyable, Movable, Preparable, Stoppable
+from pydantic import BaseModel, Field
 
 from ophyd_async.core import (
     ConfigSignal,
@@ -8,6 +10,7 @@ from ophyd_async.core import (
     StandardReadable,
     WatchableAsyncStatus,
 )
+from ophyd_async.core.async_status import AsyncStatus
 from ophyd_async.core.signal import observe_value
 from ophyd_async.core.utils import (
     DEFAULT_TIMEOUT,
@@ -19,7 +22,39 @@ from ophyd_async.core.utils import (
 from ..signal.signal import epics_signal_r, epics_signal_rw, epics_signal_x
 
 
-class Motor(StandardReadable, Movable, Stoppable):
+class MotorLimitsException(Exception):
+    pass
+
+
+class InvalidFlyMotorException(Exception):
+    pass
+
+
+DEFAULT_MOTOR_FLY_TIMEOUT = 60
+DEFAULT_WATCHER_UPDATE_FREQUENCY = 0.2
+
+
+class FlyMotorInfo(BaseModel):
+    """Minimal set of information required to fly a motor:"""
+
+    #: Absolute position of the motor once it finishes accelerating to desired
+    #: velocity, in motor EGUs
+    start_position: float = Field(frozen=True)
+
+    #: Absolute position of the motor once it begins decelerating from desired
+    #: velocity, in EGUs
+    end_position: float = Field(frozen=True)
+
+    #: Time taken for the motor to get from start_position to end_position, excluding
+    #: run-up and run-down, in seconds.
+    time_for_move: float = Field(frozen=True, gt=0)
+
+    #: Maximum time for the complete motor move, including run up and run down.
+    #: Defaults to `time_for_move` + run up and run down times + 10s.
+    timeout: CalculatableTimeout = Field(frozen=True, default=CalculateTimeout)
+
+
+class Motor(StandardReadable, Movable, Stoppable, Flyable, Preparable):
     """Device that moves a motor record"""
 
     def __init__(self, prefix: str, name="") -> None:
@@ -43,6 +78,16 @@ class Motor(StandardReadable, Movable, Stoppable):
         self.motor_stop = epics_signal_x(prefix + ".STOP")
         # Whether set() should complete successfully or not
         self._set_success = True
+
+        # end_position of a fly move, with run_up_distance added on.
+        self._fly_completed_position: Optional[float] = None
+
+        # Set on kickoff(), complete when motor reaches self._fly_completed_position
+        self._fly_status: Optional[WatchableAsyncStatus] = None
+
+        # Set during prepare
+        self._fly_timeout: Optional[CalculatableTimeout] = CalculateTimeout
+
         super().__init__(name=name)
 
     def set_name(self, name: str):
@@ -50,6 +95,44 @@ class Motor(StandardReadable, Movable, Stoppable):
         # Readback should be named the same as its parent in read()
         self.user_readback.set_name(name)
 
+    @AsyncStatus.wrap
+    async def prepare(self, value: FlyMotorInfo):
+        """Calculate required velocity and run-up distance, then if motor limits aren't
+        breached, move to start position minus run-up distance"""
+
+        self._fly_timeout = value.timeout
+
+        # Velocity, at which motor travels from start_position to end_position, in motor
+        # egu/s.
+        fly_velocity = await self._prepare_velocity(
+            value.start_position,
+            value.end_position,
+            value.time_for_move,
+        )
+
+        # start_position with run_up_distance added on.
+        fly_prepared_position = await self._prepare_motor_path(
+            abs(fly_velocity), value.start_position, value.end_position
+        )
+
+        await self.set(fly_prepared_position)
+
+    @AsyncStatus.wrap
+    async def kickoff(self):
+        """Begin moving motor from prepared position to final position."""
+        assert (
+            self._fly_completed_position
+        ), "Motor must be prepared before attempting to kickoff"
+
+        self._fly_status = self.set(
+            self._fly_completed_position, timeout=self._fly_timeout
+        )
+
+    def complete(self) -> WatchableAsyncStatus:
+        """Mark as complete once motor reaches completed position."""
+        assert self._fly_status, "kickoff not called"
+        return self._fly_status
+
     @WatchableAsyncStatus.wrap
     async def set(
         self, new_position: float, timeout: CalculatableTimeout = CalculateTimeout
@@ -95,3 +178,50 @@ class Motor(StandardReadable, Movable, Stoppable):
         # Put with completion will never complete as we are waiting for completion on
         # the move above, so need to pass wait=False
         await self.motor_stop.trigger(wait=False)
+
+    async def _prepare_velocity(
+        self, start_position: float, end_position: float, time_for_move: float
+    ) -> float:
+        fly_velocity = (start_position - end_position) / time_for_move
+        max_speed, egu = await asyncio.gather(
+            self.max_velocity.get_value(), self.motor_egu.get_value()
+        )
+        if abs(fly_velocity) > max_speed:
+            raise MotorLimitsException(
+                f"Motor speed of {abs(fly_velocity)} {egu}/s was requested for a motor "
+                f" with max speed of {max_speed} {egu}/s"
+            )
+        await self.velocity.set(abs(fly_velocity))
+        return fly_velocity
+
+    async def _prepare_motor_path(
+        self, fly_velocity: float, start_position: float, end_position: float
+    ) -> float:
+        # Distance required for motor to accelerate from stationary to fly_velocity, and
+        # distance required for motor to decelerate from fly_velocity to stationary
+        run_up_distance = (await self.acceleration_time.get_value()) * fly_velocity
+
+        self._fly_completed_position = end_position + run_up_distance
+
+        # Prepared position not used after prepare, so no need to store in self
+        fly_prepared_position = start_position - run_up_distance
+
+        motor_lower_limit, motor_upper_limit, egu = await asyncio.gather(
+            self.low_limit_travel.get_value(),
+            self.high_limit_travel.get_value(),
+            self.motor_egu.get_value(),
+        )
+
+        if (
+            not motor_upper_limit >= fly_prepared_position >= motor_lower_limit
+            or not motor_upper_limit
+            >= self._fly_completed_position
+            >= motor_lower_limit
+        ):
+            raise MotorLimitsException(
+                f"Motor trajectory for requested fly is from "
+                f"{fly_prepared_position}{egu} to "
+                f"{self._fly_completed_position}{egu} but motor limits are "
+                f"{motor_lower_limit}{egu} <= x <= {motor_upper_limit}{egu} "
+            )
+        return fly_prepared_position

ophyd_async/panda/__init__.py

@@ -1,7 +1,10 @@
 from ._common_blocks import (
     CommonPandaBlocks,
     DataBlock,
+    EnableDisableOptions,
     PcapBlock,
+    PcompBlock,
+    PcompDirectionOptions,
     PulseBlock,
     SeqBlock,
     TimeUnits,
@@ -15,17 +18,27 @@ from ._table import (
     seq_table_from_arrays,
     seq_table_from_rows,
 )
-from ._trigger import StaticSeqTableTriggerLogic
+from ._trigger import (
+    PcompInfo,
+    SeqTableInfo,
+    StaticPcompTriggerLogic,
+    StaticSeqTableTriggerLogic,
+)
 from ._utils import phase_sorter
 
 __all__ = [
     "CommonPandaBlocks",
     "HDFPanda",
+    "PcompBlock",
+    "PcompInfo",
+    "PcompDirectionOptions",
+    "EnableDisableOptions",
     "PcapBlock",
     "PulseBlock",
     "seq_table_from_arrays",
     "seq_table_from_rows",
     "SeqBlock",
+    "SeqTableInfo",
     "SeqTable",
     "SeqTableRow",
     "SeqTrigger",
@@ -35,4 +48,5 @@ __all__ = [
     "DataBlock",
     "CommonPandABlocks",
     "StaticSeqTableTriggerLogic",
+    "StaticPcompTriggerLogic",
 ]

ophyd_async/panda/_common_blocks.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from enum import Enum
 
 from ophyd_async.core import Device, DeviceVector, SignalR, SignalRW
+from ophyd_async.core.signal_backend import SubsetEnum
 from ophyd_async.panda._table import DatasetTable, SeqTable
 
 
@@ -22,6 +23,25 @@ class PulseBlock(Device):
     width: SignalRW[float]
 
 
+class PcompDirectionOptions(str, Enum):
+    positive = "Positive"
+    negative = "Negative"
+    either = "Either"
+
+
+EnableDisableOptions = SubsetEnum["ZERO", "ONE"]
+
+
+class PcompBlock(Device):
+    active: SignalR[bool]
+    dir: SignalRW[PcompDirectionOptions]
+    enable: SignalRW[EnableDisableOptions]
+    pulses: SignalRW[int]
+    start: SignalRW[int]
+    step: SignalRW[int]
+    width: SignalRW[int]
+
+
 class TimeUnits(str, Enum):
     min = "min"
     s = "s"
@@ -35,7 +55,7 @@ class SeqBlock(Device):
     repeats: SignalRW[int]
     prescale: SignalRW[float]
     prescale_units: SignalRW[TimeUnits]
-    enable: SignalRW[str]
+    enable: SignalRW[EnableDisableOptions]
 
 
 class PcapBlock(Device):
@@ -46,5 +66,6 @@ class PcapBlock(Device):
 class CommonPandaBlocks(Device):
     pulse: DeviceVector[PulseBlock]
     seq: DeviceVector[SeqBlock]
+    pcomp: DeviceVector[PcompBlock]
     pcap: PcapBlock
     data: DataBlock

ophyd_async/panda/_hdf_panda.py

@@ -4,7 +4,7 @@ from typing import Sequence
 
 from ophyd_async.core import (
     DEFAULT_TIMEOUT,
-    DirectoryProvider,
+    PathProvider,
     SignalR,
     StandardDetector,
 )
@@ -19,7 +19,7 @@ class HDFPanda(CommonPandaBlocks, StandardDetector):
     def __init__(
         self,
         prefix: str,
-        directory_provider: DirectoryProvider,
+        path_provider: PathProvider,
         config_sigs: Sequence[SignalR] = (),
         name: str = "",
     ):
@@ -28,7 +28,9 @@ class HDFPanda(CommonPandaBlocks, StandardDetector):
         create_children_from_annotations(self)
         controller = PandaPcapController(pcap=self.pcap)
         writer = PandaHDFWriter(
-            directory_provider=directory_provider,
+            prefix=prefix,
+            path_provider=path_provider,
+            name_provider=lambda: name,
             panda_device=self,
         )
         super().__init__(

ophyd_async/panda/_table.py

@@ -1,9 +1,11 @@
 from dataclasses import dataclass
 from enum import Enum
-from typing import Optional, Sequence, Type, TypedDict, TypeVar
+from typing import Optional, Sequence, Type, TypeVar
 
 import numpy as np
 import numpy.typing as npt
+import pydantic_numpy.typing as pnd
+from typing_extensions import NotRequired, TypedDict
 
 
 class PandaHdf5DatasetType(str, Enum):
@@ -54,23 +56,23 @@ class SeqTableRow:
 
 
 class SeqTable(TypedDict):
-    repeats: npt.NDArray[np.uint16]
-    trigger: Sequence[SeqTrigger]
-    position: npt.NDArray[np.int32]
-    time1: npt.NDArray[np.uint32]
-    outa1: npt.NDArray[np.bool_]
-    outb1: npt.NDArray[np.bool_]
-    outc1: npt.NDArray[np.bool_]
-    outd1: npt.NDArray[np.bool_]
-    oute1: npt.NDArray[np.bool_]
-    outf1: npt.NDArray[np.bool_]
-    time2: npt.NDArray[np.uint32]
-    outa2: npt.NDArray[np.bool_]
-    outb2: npt.NDArray[np.bool_]
-    outc2: npt.NDArray[np.bool_]
-    outd2: npt.NDArray[np.bool_]
-    oute2: npt.NDArray[np.bool_]
-    outf2: npt.NDArray[np.bool_]
+    repeats: NotRequired[pnd.Np1DArrayUint16]
+    trigger: NotRequired[Sequence[SeqTrigger]]
+    position: NotRequired[pnd.Np1DArrayInt32]
+    time1: NotRequired[pnd.Np1DArrayUint32]
+    outa1: NotRequired[pnd.Np1DArrayBool]
+    outb1: NotRequired[pnd.Np1DArrayBool]
+    outc1: NotRequired[pnd.Np1DArrayBool]
+    outd1: NotRequired[pnd.Np1DArrayBool]
+    oute1: NotRequired[pnd.Np1DArrayBool]
+    outf1: NotRequired[pnd.Np1DArrayBool]
+    time2: NotRequired[pnd.Np1DArrayUint32]
+    outa2: NotRequired[pnd.Np1DArrayBool]
+    outb2: NotRequired[pnd.Np1DArrayBool]
+    outc2: NotRequired[pnd.Np1DArrayBool]
+    outd2: NotRequired[pnd.Np1DArrayBool]
+    oute2: NotRequired[pnd.Np1DArrayBool]
+    outf2: NotRequired[pnd.Np1DArrayBool]
 
 
 def seq_table_from_rows(*rows: SeqTableRow):

ophyd_async/panda/_trigger.py

@@ -1,15 +1,22 @@
 import asyncio
-from dataclasses import dataclass
+from typing import Optional
+
+from pydantic import BaseModel, Field
 
 from ophyd_async.core import TriggerLogic, wait_for_value
-from ophyd_async.panda import SeqBlock, SeqTable, TimeUnits
+from ophyd_async.panda import (
+    PcompBlock,
+    PcompDirectionOptions,
+    SeqBlock,
+    SeqTable,
+    TimeUnits,
+)
 
 
-@dataclass
-class SeqTableInfo:
-    sequence_table: SeqTable
-    repeats: int
-    prescale_as_us: float = 1  # microseconds
+class SeqTableInfo(BaseModel):
+    sequence_table: SeqTable = Field(strict=True)
+    repeats: int = Field(ge=0)
+    prescale_as_us: float = Field(default=1, ge=0)  # microseconds
 
 
 class StaticSeqTableTriggerLogic(TriggerLogic[SeqTableInfo]):
@@ -37,3 +44,51 @@ class StaticSeqTableTriggerLogic(TriggerLogic[SeqTableInfo]):
     async def stop(self):
         await self.seq.enable.set("ZERO")
         await wait_for_value(self.seq.active, False, timeout=1)
+
+
+class PcompInfo(BaseModel):
+    start_postion: int = Field(description="start position in counts")
+    pulse_width: int = Field(description="width of a single pulse in counts", gt=0)
+    rising_edge_step: int = Field(
+        description="step between rising edges of pulses in counts", gt=0
+    )  #
+    number_of_pulses: int = Field(
+        description=(
+            "Number of pulses to send before the PCOMP block is disarmed. "
+            "0 means infinite."
+        ),
+        ge=0,
+    )
+    direction: PcompDirectionOptions = Field(
+        description=(
+            "Specifies which direction the motor counts should be "
+            "moving. Pulses won't be sent unless the values are moving in "
+            "this direction."
+        )
+    )
+
+
+class StaticPcompTriggerLogic(TriggerLogic[PcompInfo]):
+    def __init__(self, pcomp: PcompBlock) -> None:
+        self.pcomp = pcomp
+
+    async def prepare(self, value: PcompInfo):
+        await self.pcomp.enable.set("ZERO")
+        await asyncio.gather(
+            self.pcomp.start.set(value.start_postion),
+            self.pcomp.width.set(value.pulse_width),
+            self.pcomp.step.set(value.rising_edge_step),
+            self.pcomp.pulses.set(value.number_of_pulses),
+            self.pcomp.dir.set(value.direction),
+        )
+
+    async def kickoff(self) -> None:
+        await self.pcomp.enable.set("ONE")
+        await wait_for_value(self.pcomp.active, True, timeout=1)
+
+    async def complete(self, timeout: Optional[float] = None) -> None:
+        await wait_for_value(self.pcomp.active, False, timeout=timeout)
+
+    async def stop(self):
+        await self.pcomp.enable.set("ZERO")
+        await wait_for_value(self.pcomp.active, False, timeout=1)

ophyd_async/panda/writers/_hdf_writer.py

@@ -8,13 +8,14 @@ from p4p.client.thread import Context
 from ophyd_async.core import (
     DEFAULT_TIMEOUT,
     DetectorWriter,
-    DirectoryProvider,
+    NameProvider,
+    PathProvider,
     wait_for_value,
 )
 from ophyd_async.core.signal import observe_value
+from ophyd_async.epics.areadetector.writers.general_hdffile import _HDFDataset, _HDFFile
 
 from .._common_blocks import CommonPandaBlocks
-from ._panda_hdf_file import _HDFDataset, _HDFFile
 
 
 class PandaHDFWriter(DetectorWriter):
@@ -22,11 +23,15 @@ class PandaHDFWriter(DetectorWriter):
 
     def __init__(
         self,
-        directory_provider: DirectoryProvider,
+        prefix: str,
+        path_provider: PathProvider,
+        name_provider: NameProvider,
         panda_device: CommonPandaBlocks,
     ) -> None:
         self.panda_device = panda_device
-        self._directory_provider = directory_provider
+        self._prefix = prefix
+        self._path_provider = path_provider
+        self._name_provider = name_provider
         self._datasets: List[_HDFDataset] = []
         self._file: Optional[_HDFFile] = None
         self._multiplier = 1
@@ -39,16 +44,18 @@ class PandaHDFWriter(DetectorWriter):
         await self.panda_device.data.flush_period.set(0)
 
         self._file = None
-        info = self._directory_provider()
+        info = self._path_provider(device_name=self.panda_device.name)
         # Set the initial values
         await asyncio.gather(
             self.panda_device.data.hdf_directory.set(
                 str(info.root / info.resource_dir)
             ),
             self.panda_device.data.hdf_file_name.set(
-                f"{info.prefix}{self.panda_device.name}{info.suffix}.h5",
+                f"{info.filename}.h5",
             ),
             self.panda_device.data.num_capture.set(0),
+            # TODO: Set create_dir_depth once available
+            # https://github.com/bluesky/ophyd-async/issues/317
         )
 
         # Wait for it to start, stashing the status that tells us when it finishes
@@ -71,6 +78,7 @@ class PandaHDFWriter(DetectorWriter):
                 source=self.panda_device.data.hdf_directory.source,
                 shape=ds.shape,
                 dtype="array" if ds.shape != [1] else "number",
+                dtype_numpy="<f8",  # PandA data should always be written as Float64
                 external="STREAM:",
             )
             for ds in self._datasets
@@ -121,8 +129,9 @@ class PandaHDFWriter(DetectorWriter):
         if indices_written:
             if not self._file:
                 self._file = _HDFFile(
-                    self._directory_provider(),
-                    Path(await self.panda_device.data.hdf_file_name.get_value()),
+                    self._path_provider(),
+                    Path(await self.panda_device.data.hdf_directory.get_value())
+                    / Path(await self.panda_device.data.hdf_file_name.get_value()),
                     self._datasets,
                 )
                 for doc in self._file.stream_resources():

ophyd_async/plan_stubs/ensure_connected.py

@@ -10,13 +10,18 @@ def ensure_connected(
     timeout: float = DEFAULT_TIMEOUT,
     force_reconnect=False,
 ):
-    yield from bps.wait_for(
+    (connect_task,) = yield from bps.wait_for(
         [
             lambda: wait_for_connection(
                 **{
-                    device.name: device.connect(mock, timeout, force_reconnect)
+                    device.name: device.connect(
+                        mock=mock, timeout=timeout, force_reconnect=force_reconnect
+                    )
                     for device in devices
                 }
             )
         ]
     )
+
+    if connect_task and connect_task.exception() is not None:
+        raise connect_task.exception()

ophyd_async/plan_stubs/fly.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 import bluesky.plan_stubs as bps
 from bluesky.utils import short_uid
@@ -6,8 +6,33 @@ from bluesky.utils import short_uid
 from ophyd_async.core.detector import DetectorTrigger, StandardDetector, TriggerInfo
 from ophyd_async.core.flyer import HardwareTriggeredFlyable
 from ophyd_async.core.utils import in_micros
-from ophyd_async.panda._table import SeqTable, SeqTableRow, seq_table_from_rows
-from ophyd_async.panda._trigger import SeqTableInfo
+from ophyd_async.panda import (
+    PcompDirectionOptions,
+    PcompInfo,
+    SeqTable,
+    SeqTableInfo,
+    SeqTableRow,
+    seq_table_from_rows,
+)
+
+
+def prepare_static_pcomp_flyer_and_detectors(
+    flyer: HardwareTriggeredFlyable[PcompInfo],
+    detectors: List[StandardDetector],
+    pcomp_info: PcompInfo,
+    trigger_info: TriggerInfo,
+):
+    """Prepare a hardware triggered flyable and one or more detectors.
+
+    Prepare a hardware triggered flyable and one or more detectors with the
+    same trigger.
+
+    """
+
+    for det in detectors:
+        yield from bps.prepare(det, trigger_info, wait=False, group="prep")
+    yield from bps.prepare(flyer, pcomp_info, wait=False, group="prep")
+    yield from bps.wait(group="prep")
 
 
 def prepare_static_seq_table_flyer_and_detectors_with_same_trigger(
@@ -18,7 +43,7 @@ def prepare_static_seq_table_flyer_and_detectors_with_same_trigger(
     shutter_time: float,
     repeats: int = 1,
     period: float = 0.0,
-    frame_timeout: float | None = None,
+    frame_timeout: Optional[float] = None,
 ):
     """Prepare a hardware triggered flyable and one or more detectors.
 
@@ -36,7 +61,7 @@ def prepare_static_seq_table_flyer_and_detectors_with_same_trigger(
     deadtime = max(det.controller.get_deadtime(exposure) for det in detectors)
 
     trigger_info = TriggerInfo(
-        num=number_of_frames * repeats,
+        number=number_of_frames * repeats,
         trigger=DetectorTrigger.constant_gate,
         deadtime=deadtime,
         livetime=exposure,
@@ -65,7 +90,7 @@ def prepare_static_seq_table_flyer_and_detectors_with_same_trigger(
         SeqTableRow(time2=in_micros(shutter_time)),
     )
 
-    table_info = SeqTableInfo(table, repeats)
+    table_info = SeqTableInfo(sequence_table=table, repeats=repeats)
 
     for det in detectors:
         yield from bps.prepare(det, trigger_info, wait=False, group="prep")
@@ -75,7 +100,7 @@ def prepare_static_seq_table_flyer_and_detectors_with_same_trigger(
 
 def fly_and_collect(
     stream_name: str,
-    flyer: HardwareTriggeredFlyable[SeqTableInfo],
+    flyer: HardwareTriggeredFlyable[SeqTableInfo] | HardwareTriggeredFlyable[PcompInfo],
     detectors: List[StandardDetector],
 ):
     """Kickoff, complete and collect with a flyer and multiple detectors.
@@ -113,6 +138,32 @@ def fly_and_collect(
     yield from bps.wait(group=group)
 
 
+def fly_and_collect_with_static_pcomp(
+    stream_name: str,
+    flyer: HardwareTriggeredFlyable[PcompInfo],
+    detectors: List[StandardDetector],
+    number_of_pulses: int,
+    pulse_width: int,
+    rising_edge_step: int,
+    direction: PcompDirectionOptions,
+    trigger_info: TriggerInfo,
+):
+    # Set up scan and prepare trigger
+    pcomp_info = PcompInfo(
+        start_postion=0,
+        pulse_width=pulse_width,
+        rising_edge_step=rising_edge_step,
+        number_of_pulses=number_of_pulses,
+        direction=direction,
+    )
+    yield from prepare_static_pcomp_flyer_and_detectors(
+        flyer, detectors, pcomp_info, trigger_info
+    )
+
+    # Run the fly scan
+    yield from fly_and_collect(stream_name, flyer, detectors)
+
+
 def time_resolved_fly_and_collect_with_static_seq_table(
     stream_name: str,
     flyer: HardwareTriggeredFlyable[SeqTableInfo],