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

ophyd_async/panda/_panda_controller.py

@@ -0,0 +1,37 @@
+import asyncio
+from typing import Optional
+
+from ophyd_async.core import (
+    AsyncStatus,
+    DetectorControl,
+    DetectorTrigger,
+    wait_for_value,
+)
+from ophyd_async.panda import PcapBlock
+
+
+class PandaPcapController(DetectorControl):
+    def __init__(self, pcap: PcapBlock) -> None:
+        self.pcap = pcap
+
+    def get_deadtime(self, exposure: float) -> float:
+        return 0.000000008
+
+    async def arm(
+        self,
+        num: int,
+        trigger: DetectorTrigger = DetectorTrigger.constant_gate,
+        exposure: Optional[float] = None,
+    ) -> AsyncStatus:
+        assert trigger in (
+            DetectorTrigger.constant_gate,
+            trigger == DetectorTrigger.variable_gate,
+        ), "Only constant_gate and variable_gate triggering is supported on the PandA"
+        await asyncio.gather(self.pcap.arm.set(True))
+        await wait_for_value(self.pcap.active, True, timeout=1)
+        return AsyncStatus(wait_for_value(self.pcap.active, False, timeout=None))
+
+    async def disarm(self) -> AsyncStatus:
+        await asyncio.gather(self.pcap.arm.set(False))
+        await wait_for_value(self.pcap.active, False, timeout=1)
+        return AsyncStatus(wait_for_value(self.pcap.active, False, timeout=None))

ophyd_async/panda/_table.py

@@ -0,0 +1,158 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Sequence, Type, TypedDict, TypeVar
+
+import numpy as np
+import numpy.typing as npt
+
+
+class SeqTrigger(str, Enum):
+    IMMEDIATE = "Immediate"
+    BITA_0 = "BITA=0"
+    BITA_1 = "BITA=1"
+    BITB_0 = "BITB=0"
+    BITB_1 = "BITB=1"
+    BITC_0 = "BITC=0"
+    BITC_1 = "BITC=1"
+    POSA_GT = "POSA>=POSITION"
+    POSA_LT = "POSA<=POSITION"
+    POSB_GT = "POSB>=POSITION"
+    POSB_LT = "POSB<=POSITION"
+    POSC_GT = "POSC>=POSITION"
+    POSC_LT = "POSC<=POSITION"
+
+
+@dataclass
+class SeqTableRow:
+    repeats: int = 1
+    trigger: SeqTrigger = SeqTrigger.IMMEDIATE
+    position: int = 0
+    time1: int = 0
+    outa1: bool = False
+    outb1: bool = False
+    outc1: bool = False
+    outd1: bool = False
+    oute1: bool = False
+    outf1: bool = False
+    time2: int = 0
+    outa2: bool = False
+    outb2: bool = False
+    outc2: bool = False
+    outd2: bool = False
+    oute2: bool = False
+    outf2: bool = False
+
+
+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_]
+
+
+def seq_table_from_rows(*rows: SeqTableRow):
+    """
+    Constructs a sequence table from a series of rows.
+    """
+    return seq_table_from_arrays(
+        repeats=np.array([row.repeats for row in rows], dtype=np.uint16),
+        trigger=[row.trigger for row in rows],
+        position=np.array([row.position for row in rows], dtype=np.int32),
+        time1=np.array([row.time1 for row in rows], dtype=np.uint32),
+        outa1=np.array([row.outa1 for row in rows], dtype=np.bool_),
+        outb1=np.array([row.outb1 for row in rows], dtype=np.bool_),
+        outc1=np.array([row.outc1 for row in rows], dtype=np.bool_),
+        outd1=np.array([row.outd1 for row in rows], dtype=np.bool_),
+        oute1=np.array([row.oute1 for row in rows], dtype=np.bool_),
+        outf1=np.array([row.outf1 for row in rows], dtype=np.bool_),
+        time2=np.array([row.time2 for row in rows], dtype=np.uint32),
+        outa2=np.array([row.outa2 for row in rows], dtype=np.bool_),
+        outb2=np.array([row.outb2 for row in rows], dtype=np.bool_),
+        outc2=np.array([row.outc2 for row in rows], dtype=np.bool_),
+        outd2=np.array([row.outd2 for row in rows], dtype=np.bool_),
+        oute2=np.array([row.oute2 for row in rows], dtype=np.bool_),
+        outf2=np.array([row.outf2 for row in rows], dtype=np.bool_),
+    )
+
+
+T = TypeVar("T", bound=np.generic)
+
+
+def seq_table_from_arrays(
+    *,
+    repeats: Optional[npt.NDArray[np.uint16]] = None,
+    trigger: Optional[Sequence[SeqTrigger]] = None,
+    position: Optional[npt.NDArray[np.int32]] = None,
+    time1: Optional[npt.NDArray[np.uint32]] = None,
+    outa1: Optional[npt.NDArray[np.bool_]] = None,
+    outb1: Optional[npt.NDArray[np.bool_]] = None,
+    outc1: Optional[npt.NDArray[np.bool_]] = None,
+    outd1: Optional[npt.NDArray[np.bool_]] = None,
+    oute1: Optional[npt.NDArray[np.bool_]] = None,
+    outf1: Optional[npt.NDArray[np.bool_]] = None,
+    time2: npt.NDArray[np.uint32],
+    outa2: Optional[npt.NDArray[np.bool_]] = None,
+    outb2: Optional[npt.NDArray[np.bool_]] = None,
+    outc2: Optional[npt.NDArray[np.bool_]] = None,
+    outd2: Optional[npt.NDArray[np.bool_]] = None,
+    oute2: Optional[npt.NDArray[np.bool_]] = None,
+    outf2: Optional[npt.NDArray[np.bool_]] = None,
+) -> SeqTable:
+    """
+    Constructs a sequence table from a series of columns as arrays.
+    time2 is the only required argument and must not be None.
+    All other provided arguments must be of equal length to time2.
+    If any other argument is not given, or else given as None or empty,
+    an array of length len(time2) filled with the following is defaulted:
+    repeats: 1
+    trigger: SeqTrigger.IMMEDIATE
+    all others: 0/False as appropriate
+    """
+    assert time2 is not None, "time2 must be provided"
+    length = len(time2)
+    assert 0 < length < 4096, f"Length {length} not in range"
+
+    def or_default(
+        value: Optional[npt.NDArray[T]], dtype: Type[T], default_value: int = 0
+    ) -> npt.NDArray[T]:
+        if value is None or len(value) == 0:
+            return np.full(length, default_value, dtype=dtype)
+        return value
+
+    table = SeqTable(
+        repeats=or_default(repeats, np.uint16, 1),
+        trigger=trigger or [SeqTrigger.IMMEDIATE] * length,
+        position=or_default(position, np.int32),
+        time1=or_default(time1, np.uint32),
+        outa1=or_default(outa1, np.bool_),
+        outb1=or_default(outb1, np.bool_),
+        outc1=or_default(outc1, np.bool_),
+        outd1=or_default(outd1, np.bool_),
+        oute1=or_default(oute1, np.bool_),
+        outf1=or_default(outf1, np.bool_),
+        time2=time2,
+        outa2=or_default(outa2, np.bool_),
+        outb2=or_default(outb2, np.bool_),
+        outc2=or_default(outc2, np.bool_),
+        outd2=or_default(outd2, np.bool_),
+        oute2=or_default(oute2, np.bool_),
+        outf2=or_default(outf2, np.bool_),
+    )
+    for k, v in table.items():
+        size = len(v)  # type: ignore
+        if size != length:
+            raise ValueError(f"{k}: has length {size} not {length}")
+    return table

ophyd_async/panda/_trigger.py

@@ -0,0 +1,39 @@
+import asyncio
+from dataclasses import dataclass
+
+from ophyd_async.core import TriggerLogic, wait_for_value
+from ophyd_async.panda import SeqBlock, SeqTable, TimeUnits
+
+
+@dataclass
+class SeqTableInfo:
+    sequence_table: SeqTable
+    repeats: int
+    prescale_as_us: float = 1  # microseconds
+
+
+class StaticSeqTableTriggerLogic(TriggerLogic[SeqTableInfo]):
+    def __init__(self, seq: SeqBlock) -> None:
+        self.seq = seq
+
+    async def prepare(self, value: SeqTableInfo):
+        await asyncio.gather(
+            self.seq.prescale_units.set(TimeUnits.us),
+            self.seq.enable.set("ZERO"),
+        )
+        await asyncio.gather(
+            self.seq.prescale.set(value.prescale_as_us),
+            self.seq.repeats.set(value.repeats),
+            self.seq.table.set(value.sequence_table),
+        )
+
+    async def kickoff(self) -> None:
+        await self.seq.enable.set("ONE")
+        await wait_for_value(self.seq.active, True, timeout=1)
+
+    async def complete(self) -> None:
+        await wait_for_value(self.seq.active, False, timeout=None)
+
+    async def stop(self):
+        await self.seq.enable.set("ZERO")
+        await wait_for_value(self.seq.active, False, timeout=1)

ophyd_async/panda/_utils.py

@@ -0,0 +1,15 @@
+from typing import Any, Dict, Sequence
+
+
+def phase_sorter(panda_signal_values: Dict[str, Any]) -> Sequence[Dict[str, Any]]:
+    # Panda has two load phases. If the signal name ends in the string "UNITS",
+    # it needs to be loaded first so put in first phase
+    phase_1, phase_2 = {}, {}
+
+    for key, value in panda_signal_values.items():
+        if key.endswith("units"):
+            phase_1[key] = value
+        else:
+            phase_2[key] = value
+
+    return [phase_1, phase_2]

ophyd_async/panda/writers/__init__.py

@@ -0,0 +1,3 @@
+from ._hdf_writer import PandaHDFWriter
+
+__all__ = ["PandaHDFWriter"]

ophyd_async/panda/writers/_hdf_writer.py

@@ -0,0 +1,220 @@
+import asyncio
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+from typing import Any, AsyncGenerator, AsyncIterator, Dict, List, Optional
+
+from bluesky.protocols import DataKey, StreamAsset
+from p4p.client.thread import Context
+
+from ophyd_async.core import (
+    DEFAULT_TIMEOUT,
+    DetectorWriter,
+    Device,
+    DirectoryProvider,
+    NameProvider,
+    SignalR,
+    wait_for_value,
+)
+from ophyd_async.core.signal import observe_value
+from ophyd_async.panda import CommonPandaBlocks
+
+from ._panda_hdf_file import _HDFDataset, _HDFFile
+
+
+class Capture(str, Enum):
+    # Capture signals for the HDF Panda
+    No = "No"
+    Value = "Value"
+    Diff = "Diff"
+    Sum = "Sum"
+    Mean = "Mean"
+    Min = "Min"
+    Max = "Max"
+    MinMax = "Min Max"
+    MinMaxMean = "Min Max Mean"
+
+
+def get_capture_signals(
+    block: Device, path_prefix: Optional[str] = ""
+) -> Dict[str, SignalR]:
+    """Get dict mapping a capture signal's name to the signal itself"""
+    if not path_prefix:
+        path_prefix = ""
+    signals: Dict[str, SignalR[Any]] = {}
+    for attr_name, attr in block.children():
+        # Capture signals end in _capture, but num_capture is a red herring
+        if attr_name == "num_capture":
+            continue
+        dot_path = f"{path_prefix}{attr_name}"
+        if isinstance(attr, SignalR) and attr_name.endswith("_capture"):
+            signals[dot_path] = attr
+        attr_signals = get_capture_signals(attr, path_prefix=dot_path + ".")
+        signals.update(attr_signals)
+    return signals
+
+
+@dataclass
+class CaptureSignalWrapper:
+    signal: SignalR
+    capture_type: Capture
+
+
+# This should return a dictionary which contains a dict, containing the Capture
+# signal object, and the value of that signal
+async def get_signals_marked_for_capture(
+    capture_signals: Dict[str, SignalR],
+) -> Dict[str, CaptureSignalWrapper]:
+    # Read signals to see if they should be captured
+    do_read = [signal.get_value() for signal in capture_signals.values()]
+
+    signal_values = await asyncio.gather(*do_read)
+
+    assert len(signal_values) == len(
+        capture_signals
+    ), "Length of read signals are different to length of signals"
+
+    signals_to_capture: Dict[str, CaptureSignalWrapper] = {}
+    for signal_path, signal_object, signal_value in zip(
+        capture_signals.keys(), capture_signals.values(), signal_values
+    ):
+        signal_path = signal_path.replace("_capture", "")
+        if (signal_value in iter(Capture)) and (signal_value != Capture.No):
+            signals_to_capture[signal_path] = CaptureSignalWrapper(
+                signal_object,
+                signal_value,
+            )
+
+    return signals_to_capture
+
+
+class PandaHDFWriter(DetectorWriter):
+    _ctxt: Optional[Context] = None
+
+    def __init__(
+        self,
+        prefix: str,
+        directory_provider: DirectoryProvider,
+        name_provider: NameProvider,
+        panda_device: CommonPandaBlocks,
+    ) -> None:
+        self.panda_device = panda_device
+        self._prefix = prefix
+        self._directory_provider = directory_provider
+        self._name_provider = name_provider
+        self._datasets: List[_HDFDataset] = []
+        self._file: Optional[_HDFFile] = None
+        self._multiplier = 1
+
+    # Triggered on PCAP arm
+    async def open(self, multiplier: int = 1) -> Dict[str, DataKey]:
+        """Retrieve and get descriptor of all PandA signals marked for capture"""
+
+        # Get capture PVs by looking at panda. Gives mapping of dotted attribute path
+        # to Signal object
+        self.capture_signals = get_capture_signals(self.panda_device)
+
+        # Ensure flushes are immediate
+        await self.panda_device.data.flush_period.set(0)
+
+        to_capture = await get_signals_marked_for_capture(self.capture_signals)
+        self._file = None
+        info = self._directory_provider()
+        # 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",
+            ),
+            self.panda_device.data.num_capture.set(0),
+        )
+
+        # Wait for it to start, stashing the status that tells us when it finishes
+        await self.panda_device.data.capture.set(True)
+        name = self._name_provider()
+        if multiplier > 1:
+            raise ValueError(
+                "All PandA datasets should be scalar, multiplier should be 1"
+            )
+        self._datasets = []
+        for attribute_path, capture_signal in to_capture.items():
+            split_path = attribute_path.split(".")
+            signal_name = split_path[-1]
+            # Get block names from numbered blocks, eg INENC[1]
+            block_name = (
+                f"{split_path[-3]}{split_path[-2]}"
+                if split_path[-2].isnumeric()
+                else split_path[-2]
+            )
+
+            for suffix in capture_signal.capture_type.split(" "):
+                self._datasets.append(
+                    _HDFDataset(
+                        name,
+                        block_name,
+                        f"{name}-{block_name}-{signal_name}-{suffix}",
+                        f"{block_name}-{signal_name}".upper() + f"-{suffix}",
+                        [1],
+                        multiplier=1,
+                    )
+                )
+
+        describe = {
+            ds.name: DataKey(
+                source=self.panda_device.data.hdf_directory.source,
+                shape=ds.shape,
+                dtype="array" if ds.shape != [1] else "number",
+                external="STREAM:",
+            )
+            for ds in self._datasets
+        }
+        return describe
+
+    # Next few functions are exactly the same as AD writer. Could move as default
+    # StandardDetector behavior
+    async def wait_for_index(
+        self, index: int, timeout: Optional[float] = DEFAULT_TIMEOUT
+    ):
+        def matcher(value: int) -> bool:
+            return value >= index
+
+        matcher.__name__ = f"index_at_least_{index}"
+        await wait_for_value(
+            self.panda_device.data.num_captured, matcher, timeout=timeout
+        )
+
+    async def get_indices_written(self) -> int:
+        return await self.panda_device.data.num_captured.get_value()
+
+    async def observe_indices_written(
+        self, timeout=DEFAULT_TIMEOUT
+    ) -> AsyncGenerator[int, None]:
+        """Wait until a specific index is ready to be collected"""
+        async for num_captured in observe_value(
+            self.panda_device.data.num_captured, timeout
+        ):
+            yield num_captured // self._multiplier
+
+    async def collect_stream_docs(
+        self, indices_written: int
+    ) -> AsyncIterator[StreamAsset]:
+        # TODO: fail if we get dropped frames
+        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._datasets,
+                )
+                for doc in self._file.stream_resources():
+                    yield "stream_resource", doc
+            for doc in self._file.stream_data(indices_written):
+                yield "stream_datum", doc
+
+    # Could put this function as default for StandardDetector
+    async def close(self):
+        await self.panda_device.data.capture.set(
+            False, wait=True, timeout=DEFAULT_TIMEOUT
+        )

ophyd_async/panda/writers/_panda_hdf_file.py

@@ -0,0 +1,58 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterator, List
+
+from event_model import StreamDatum, StreamResource, compose_stream_resource
+
+from ophyd_async.core import DirectoryInfo
+
+
+@dataclass
+class _HDFDataset:
+    device_name: str
+    block: str
+    name: str
+    path: str
+    shape: List[int]
+    multiplier: int
+
+
+class _HDFFile:
+    def __init__(
+        self,
+        directory_info: DirectoryInfo,
+        full_file_name: Path,
+        datasets: List[_HDFDataset],
+    ) -> None:
+        self._last_emitted = 0
+        self._bundles = [
+            compose_stream_resource(
+                spec="AD_HDF5_SWMR_SLICE",
+                root=str(directory_info.root),
+                data_key=ds.name,
+                resource_path=(f"{str(directory_info.root)}/{full_file_name}"),
+                resource_kwargs={
+                    "name": ds.name,
+                    "block": ds.block,
+                    "path": ds.path,
+                    "multiplier": ds.multiplier,
+                    "timestamps": "/entry/instrument/NDAttributes/NDArrayTimeStamp",
+                },
+            )
+            for ds in datasets
+        ]
+
+    def stream_resources(self) -> Iterator[StreamResource]:
+        for bundle in self._bundles:
+            yield bundle.stream_resource_doc
+
+    def stream_data(self, indices_written: int) -> Iterator[StreamDatum]:
+        # Indices are relative to resource
+        if indices_written > self._last_emitted:
+            indices = {
+                "start": self._last_emitted,
+                "stop": indices_written,
+            }
+            self._last_emitted = indices_written
+            for bundle in self._bundles:
+                yield bundle.compose_stream_datum(indices)

ophyd_async/plan_stubs/__init__.py

@@ -0,0 +1,13 @@
+from .ensure_connected import ensure_connected
+from .fly import (
+    fly_and_collect,
+    prepare_static_seq_table_flyer_and_detectors_with_same_trigger,
+    time_resolved_fly_and_collect_with_static_seq_table,
+)
+
+__all__ = [
+    "fly_and_collect",
+    "prepare_static_seq_table_flyer_and_detectors_with_same_trigger",
+    "time_resolved_fly_and_collect_with_static_seq_table",
+    "ensure_connected",
+]

ophyd_async/plan_stubs/ensure_connected.py

@@ -0,0 +1,22 @@
+import bluesky.plan_stubs as bps
+
+from ophyd_async.core.device import Device
+from ophyd_async.core.utils import DEFAULT_TIMEOUT, wait_for_connection
+
+
+def ensure_connected(
+    *devices: Device,
+    mock: bool = False,
+    timeout: float = DEFAULT_TIMEOUT,
+    force_reconnect=False,
+):
+    yield from bps.wait_for(
+        [
+            lambda: wait_for_connection(
+                **{
+                    device.name: device.connect(mock, timeout, force_reconnect)
+                    for device in devices
+                }
+            )
+        ]
+    )