ophyd-async 0.1.0__py3-none-any.whl → 0.3a1__py3-none-any.whl

ophyd_async/epics/areadetector/drivers/ad_base.py

@@ -0,0 +1,111 @@
+import asyncio
+from enum import Enum
+from typing import FrozenSet, Sequence, Set
+
+from ophyd_async.core import (
+    DEFAULT_TIMEOUT,
+    AsyncStatus,
+    ShapeProvider,
+    set_and_wait_for_value,
+)
+
+from ...signal.signal import epics_signal_rw
+from ..utils import ImageMode, ad_r, ad_rw
+from ..writers.nd_plugin import NDArrayBase
+
+
+class DetectorState(str, Enum):
+    """
+    Default set of states of an AreaDetector driver.
+    See definition in ADApp/ADSrc/ADDriver.h in https://github.com/areaDetector/ADCore
+    """
+
+    Idle = "Idle"
+    Acquire = "Acquire"
+    Readout = "Readout"
+    Correct = "Correct"
+    Saving = "Saving"
+    Aborting = "Aborting"
+    Error = "Error"
+    Waiting = "Waiting"
+    Initializing = "Initializing"
+    Disconnected = "Disconnected"
+    Aborted = "Aborted"
+
+
+#: Default set of states that we should consider "good" i.e. the acquisition
+#  is complete and went well
+DEFAULT_GOOD_STATES: FrozenSet[DetectorState] = frozenset(
+    [DetectorState.Idle, DetectorState.Aborted]
+)
+
+
+class ADBase(NDArrayBase):
+    def __init__(self, prefix: str, name: str = "") -> None:
+        # Define some signals
+        self.acquire = ad_rw(bool, prefix + "Acquire")
+        self.acquire_time = ad_rw(float, prefix + "AcquireTime")
+        self.num_images = ad_rw(int, prefix + "NumImages")
+        self.image_mode = ad_rw(ImageMode, prefix + "ImageMode")
+        self.array_counter = ad_rw(int, prefix + "ArrayCounter")
+        self.array_size_x = ad_r(int, prefix + "ArraySizeX")
+        self.array_size_y = ad_r(int, prefix + "ArraySizeY")
+        self.detector_state = ad_r(DetectorState, prefix + "DetectorState")
+        # There is no _RBV for this one
+        self.wait_for_plugins = epics_signal_rw(bool, prefix + "WaitForPlugins")
+        super().__init__(prefix, name=name)
+
+
+async def start_acquiring_driver_and_ensure_status(
+    driver: ADBase,
+    good_states: Set[DetectorState] = set(DEFAULT_GOOD_STATES),
+    timeout: float = DEFAULT_TIMEOUT,
+) -> AsyncStatus:
+    """
+    Start acquiring driver, raising ValueError if the detector is in a bad state.
+
+    This sets driver.acquire to True, and waits for it to be True up to a timeout.
+    Then, it checks that the DetectorState PV is in DEFAULT_GOOD_STATES, and otherwise
+    raises a ValueError.
+
+    Parameters
+    ----------
+    driver:
+        The driver to start acquiring. Must subclass ADBase.
+    good_states:
+        set of states defined in DetectorState enum which are considered good states.
+    timeout:
+        How long to wait for driver.acquire to readback True (i.e. acquiring).
+
+    Returns
+    -------
+    AsyncStatus:
+        An AsyncStatus that can be awaited to set driver.acquire to True and perform
+        subsequent raising (if applicable) due to detector state.
+    """
+
+    status = await set_and_wait_for_value(driver.acquire, True, timeout=timeout)
+
+    async def complete_acquisition() -> None:
+        """NOTE: possible race condition here between the callback from
+        set_and_wait_for_value and the detector state updating."""
+        await status
+        state = await driver.detector_state.get_value()
+        if state not in good_states:
+            raise ValueError(
+                f"Final detector state {state} not in valid end states: {good_states}"
+            )
+
+    return AsyncStatus(complete_acquisition())
+
+
+class ADBaseShapeProvider(ShapeProvider):
+    def __init__(self, driver: ADBase) -> None:
+        self._driver = driver
+
+    async def __call__(self) -> Sequence[int]:
+        shape = await asyncio.gather(
+            self._driver.array_size_y.get_value(),
+            self._driver.array_size_x.get_value(),
+        )
+        return shape

ophyd_async/epics/areadetector/drivers/pilatus_driver.py

@@ -0,0 +1,18 @@
+from enum import Enum
+
+from ..utils import ad_rw
+from .ad_base import ADBase
+
+
+class TriggerMode(str, Enum):
+    internal = "Internal"
+    ext_enable = "Ext. Enable"
+    ext_trigger = "Ext. Trigger"
+    mult_trigger = "Mult. Trigger"
+    alignment = "Alignment"
+
+
+class PilatusDriver(ADBase):
+    def __init__(self, prefix: str) -> None:
+        self.trigger_mode = ad_rw(TriggerMode, prefix + "TriggerMode")
+        super().__init__(prefix)

ophyd_async/epics/areadetector/single_trigger_det.py

@@ -5,18 +5,18 @@ from bluesky.protocols import Triggerable
 
 from ophyd_async.core import AsyncStatus, SignalR, StandardReadable
 
-from .ad_driver import ADDriver
-from .nd_plugin import NDPlugin
+from .drivers.ad_base import ADBase
 from .utils import ImageMode
+from .writers.nd_plugin import NDPluginBase
 
 
 class SingleTriggerDet(StandardReadable, Triggerable):
     def __init__(
         self,
-        drv: ADDriver,
+        drv: ADBase,
         read_uncached: Sequence[SignalR] = (),
         name="",
-        **plugins: NDPlugin,
+        **plugins: NDPluginBase,
     ) -> None:
         self.drv = drv
         self.__dict__.update(plugins)

ophyd_async/epics/areadetector/utils.py

@@ -1,7 +1,8 @@
 from enum import Enum
-from typing import Type
+from typing import Optional, Type
+from xml.etree import cElementTree as ET
 
-from ophyd_async.core import SignalR, SignalRW, T
+from ophyd_async.core import DEFAULT_TIMEOUT, SignalR, SignalRW, T, wait_for_value
 
 from ..signal.signal import epics_signal_r, epics_signal_rw
 
@@ -20,7 +21,94 @@ class FileWriteMode(str, Enum):
     stream = "Stream"
 
 
-class ImageMode(Enum):
+class ImageMode(str, Enum):
     single = "Single"
     multiple = "Multiple"
     continuous = "Continuous"
+
+
+class NDAttributeDataType(str, Enum):
+    INT = "INT"
+    DOUBLE = "DOUBLE"
+    STRING = "STRING"
+
+
+class NDAttributesXML:
+    """Helper to make NDAttributesFile XML for areaDetector"""
+
+    _dbr_types = {
+        None: "DBR_NATIVE",
+        NDAttributeDataType.INT: "DBR_LONG",
+        NDAttributeDataType.DOUBLE: "DBR_DOUBLE",
+        NDAttributeDataType.STRING: "DBR_STRING",
+    }
+
+    def __init__(self):
+        self._root = ET.Element("Attributes")
+
+    def add_epics_pv(
+        self,
+        name: str,
+        pv: str,
+        datatype: Optional[NDAttributeDataType] = None,
+        description: str = "",
+    ):
+        """Add a PV to the attribute list
+
+        Args:
+            name: The attribute name
+            pv: The pv to get from
+            datatype: An override datatype, otherwise will use native EPICS type
+            description: A description that appears in the HDF file as an attribute
+        """
+        ET.SubElement(
+            self._root,
+            "Attribute",
+            name=name,
+            type="EPICS_PV",
+            source=pv,
+            datatype=self._dbr_types[datatype],
+            description=description,
+        )
+
+    def add_param(
+        self,
+        name: str,
+        param: str,
+        datatype: NDAttributeDataType,
+        addr: int = 0,
+        description: str = "",
+    ):
+        """Add a driver or plugin parameter to the attribute list
+
+        Args:
+            name: The attribute name
+            param: The parameter string as seen in the INP link of the record
+            datatype: The datatype of the parameter
+            description: A description that appears in the HDF file as an attribute
+        """
+        ET.SubElement(
+            self._root,
+            "Attribute",
+            name=name,
+            type="PARAM",
+            source=param,
+            addr=str(addr),
+            datatype=datatype.value,
+            description=description,
+        )
+
+    def __str__(self) -> str:
+        """Output the XML pretty printed"""
+        ET.indent(self._root, space="    ", level=0)
+        return ET.tostring(self._root, xml_declaration=True, encoding="utf-8").decode()
+
+
+async def stop_busy_record(
+    signal: SignalRW[T],
+    value: T,
+    timeout: float = DEFAULT_TIMEOUT,
+    status_timeout: Optional[float] = None,
+) -> None:
+    await signal.set(value, wait=False, timeout=status_timeout)
+    await wait_for_value(signal, value, timeout=timeout)

ophyd_async/epics/areadetector/writers/__init__.py

@@ -0,0 +1,5 @@
+from .hdf_writer import HDFWriter
+from .nd_file_hdf import NDFileHDF
+from .nd_plugin import NDPluginBase, NDPluginStats
+
+__all__ = ["HDFWriter", "NDFileHDF", "NDPluginBase", "NDPluginStats"]

ophyd_async/epics/areadetector/writers/_hdfdataset.py

@@ -0,0 +1,10 @@
+from dataclasses import dataclass
+from typing import Sequence
+
+
+@dataclass
+class _HDFDataset:
+    name: str
+    path: str
+    shape: Sequence[int]
+    multiplier: int

ophyd_async/epics/areadetector/writers/_hdffile.py

@@ -0,0 +1,54 @@
+from pathlib import Path
+from typing import Iterator, List
+
+from event_model import StreamDatum, StreamResource, compose_stream_resource
+
+from ophyd_async.core import DirectoryInfo
+
+from ._hdfdataset import _HDFDataset
+
+
+class _HDFFile:
+    """
+    :param directory_info: Contains information about how to construct a StreamResource
+    :param full_file_name: Absolute path to the file to be written
+    :param datasets: Datasets to write into the file
+    """
+
+    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=str(full_file_name.relative_to(directory_info.root)),
+                resource_kwargs={
+                    "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 = dict(
+                start=self._last_emitted,
+                stop=indices_written,
+            )
+            self._last_emitted = indices_written
+            for bundle in self._bundles:
+                yield bundle.compose_stream_datum(indices)
+        return None

ophyd_async/epics/areadetector/writers/hdf_writer.py

@@ -0,0 +1,133 @@
+import asyncio
+from pathlib import Path
+from typing import AsyncGenerator, AsyncIterator, Dict, List, Optional
+
+from bluesky.protocols import Descriptor, Hints, StreamAsset
+
+from ophyd_async.core import (
+    DEFAULT_TIMEOUT,
+    AsyncStatus,
+    DetectorWriter,
+    DirectoryProvider,
+    NameProvider,
+    ShapeProvider,
+    set_and_wait_for_value,
+    wait_for_value,
+)
+from ophyd_async.core.signal import observe_value
+
+from ._hdfdataset import _HDFDataset
+from ._hdffile import _HDFFile
+from .nd_file_hdf import FileWriteMode, NDFileHDF
+
+
+class HDFWriter(DetectorWriter):
+    def __init__(
+        self,
+        hdf: NDFileHDF,
+        directory_provider: DirectoryProvider,
+        name_provider: NameProvider,
+        shape_provider: ShapeProvider,
+        **scalar_datasets_paths: str,
+    ) -> None:
+        self.hdf = hdf
+        self._directory_provider = directory_provider
+        self._name_provider = name_provider
+        self._shape_provider = shape_provider
+        self._scalar_datasets_paths = scalar_datasets_paths
+        self._capture_status: Optional[AsyncStatus] = None
+        self._datasets: List[_HDFDataset] = []
+        self._file: Optional[_HDFFile] = None
+        self._multiplier = 1
+
+    async def open(self, multiplier: int = 1) -> Dict[str, Descriptor]:
+        self._file = None
+        info = self._directory_provider()
+        await asyncio.gather(
+            self.hdf.num_extra_dims.set(0),
+            self.hdf.lazy_open.set(True),
+            self.hdf.swmr_mode.set(True),
+            # See https://github.com/bluesky/ophyd-async/issues/122
+            self.hdf.file_path.set(str(info.root / info.resource_dir)),
+            self.hdf.file_name.set(f"{info.prefix}{self.hdf.name}{info.suffix}"),
+            self.hdf.file_template.set("%s/%s.h5"),
+            self.hdf.file_write_mode.set(FileWriteMode.stream),
+        )
+
+        assert (
+            await self.hdf.file_path_exists.get_value()
+        ), f"File path {self.hdf.file_path.get_value()} for hdf plugin does not exist"
+
+        # Overwrite num_capture to go forever
+        await self.hdf.num_capture.set(0)
+        # Wait for it to start, stashing the status that tells us when it finishes
+        self._capture_status = await set_and_wait_for_value(self.hdf.capture, True)
+        name = self._name_provider()
+        detector_shape = tuple(await self._shape_provider())
+        self._multiplier = multiplier
+        outer_shape = (multiplier,) if multiplier > 1 else ()
+        # Add the main data
+        self._datasets = [
+            _HDFDataset(name, "/entry/data/data", detector_shape, multiplier)
+        ]
+        # And all the scalar datasets
+        for ds_name, ds_path in self._scalar_datasets_paths.items():
+            self._datasets.append(
+                _HDFDataset(
+                    f"{name}-{ds_name}",
+                    f"/entry/instrument/NDAttributes/{ds_path}",
+                    (),
+                    multiplier,
+                )
+            )
+        describe = {
+            ds.name: Descriptor(
+                source=self.hdf.full_file_name.source,
+                shape=outer_shape + tuple(ds.shape),
+                dtype="array" if ds.shape else "number",
+                external="STREAM:",
+            )
+            for ds in self._datasets
+        }
+        return describe
+
+    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.hdf.num_captured, timeout):
+            yield num_captured // self._multiplier
+
+    async def get_indices_written(self) -> int:
+        num_captured = await self.hdf.num_captured.get_value()
+        return num_captured // self._multiplier
+
+    async def collect_stream_docs(
+        self, indices_written: int
+    ) -> AsyncIterator[StreamAsset]:
+        # TODO: fail if we get dropped frames
+        await self.hdf.flush_now.set(True)
+        if indices_written:
+            if not self._file:
+                self._file = _HDFFile(
+                    self._directory_provider(),
+                    # See https://github.com/bluesky/ophyd-async/issues/122
+                    Path(await self.hdf.full_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
+
+    async def close(self):
+        # Already done a caput callback in _capture_status, so can't do one here
+        await self.hdf.capture.set(0, wait=False)
+        await wait_for_value(self.hdf.capture, 0, DEFAULT_TIMEOUT)
+        if self._capture_status:
+            # We kicked off an open, so wait for it to return
+            await self._capture_status
+
+    @property
+    def hints(self) -> Hints:
+        return {"fields": [self._name_provider()]}

ophyd_async/epics/areadetector/{nd_file_hdf.py → writers/nd_file_hdf.py}

@@ -1,14 +1,30 @@
-from ophyd_async.core import Device
+from enum import Enum
 
-from ..signal.signal import epics_signal_rw
-from .utils import FileWriteMode, ad_r, ad_rw
+from ...signal.signal import epics_signal_rw
+from ..utils import FileWriteMode, ad_r, ad_rw
+from .nd_plugin import NDPluginBase
 
 
-class NDFileHDF(Device):
-    def __init__(self, prefix: str) -> None:
+class Compression(str, Enum):
+    none = "None"
+    nbit = "N-bit"
+    szip = "szip"
+    zlib = "zlib"
+    blosc = "Blosc"
+    bslz4 = "BSLZ4"
+    lz4 = "LZ4"
+    jpeg = "JPEG"
+
+
+class NDFileHDF(NDPluginBase):
+    def __init__(self, prefix: str, name="") -> None:
         # Define some signals
+        self.position_mode = ad_rw(bool, prefix + "PositionMode")
+        self.compression = ad_rw(Compression, prefix + "Compression")
+        self.num_extra_dims = ad_rw(int, prefix + "NumExtraDims")
         self.file_path = ad_rw(str, prefix + "FilePath")
         self.file_name = ad_rw(str, prefix + "FileName")
+        self.file_path_exists = ad_r(bool, prefix + "FilePathExists")
         self.file_template = ad_rw(str, prefix + "FileTemplate")
         self.full_file_name = ad_r(str, prefix + "FullFileName")
         self.file_write_mode = ad_rw(FileWriteMode, prefix + "FileWriteMode")
@@ -20,3 +36,4 @@ class NDFileHDF(Device):
         self.flush_now = epics_signal_rw(bool, prefix + "FlushNow")
         self.array_size0 = ad_r(int, prefix + "ArraySize0")
         self.array_size1 = ad_r(int, prefix + "ArraySize1")
+        super().__init__(prefix, name)

ophyd_async/epics/areadetector/writers/nd_plugin.py

@@ -0,0 +1,30 @@
+from enum import Enum
+
+from ophyd_async.core import Device
+from ophyd_async.epics.signal import epics_signal_rw
+
+from ..utils import ad_r, ad_rw
+
+
+class Callback(str, Enum):
+    Enable = "Enable"
+    Disable = "Disable"
+
+
+class NDArrayBase(Device):
+    def __init__(self, prefix: str, name: str = "") -> None:
+        self.unique_id = ad_r(int, prefix + "UniqueId")
+        self.nd_attributes_file = epics_signal_rw(str, prefix + "NDAttributesFile")
+        super().__init__(name)
+
+
+class NDPluginBase(NDArrayBase):
+    def __init__(self, prefix: str, name: str = "") -> None:
+        self.nd_array_port = ad_rw(str, prefix + "NDArrayPort")
+        self.enable_callback = ad_rw(Callback, prefix + "EnableCallbacks")
+        self.nd_array_address = ad_rw(int, prefix + "NDArrayAddress")
+        super().__init__(prefix, name)
+
+
+class NDPluginStats(NDPluginBase):
+    pass

ophyd_async/epics/demo/__init__.py

@@ -19,7 +19,7 @@ from ophyd_async.core import AsyncStatus, Device, StandardReadable, observe_valu
 from ..signal.signal import epics_signal_r, epics_signal_rw, epics_signal_x
 
 
-class EnergyMode(Enum):
+class EnergyMode(str, Enum):
     """Energy mode for `Sensor`"""
 
     #: Low energy mode
@@ -112,7 +112,8 @@ class Mover(StandardReadable, Movable, Stoppable):
 
     async def stop(self, success=True):
         self._set_success = success
-        await self.stop_.execute()
+        status = self.stop_.trigger()
+        await status
 
 
 class SampleStage(Device):

ophyd_async/epics/demo/demo_ad_sim_detector.py

@@ -0,0 +1,35 @@
+from typing import Sequence
+
+from ophyd_async.core import DirectoryProvider, SignalR, StandardDetector
+
+from ..areadetector.controllers import ADSimController
+from ..areadetector.drivers import ADBase, ADBaseShapeProvider
+from ..areadetector.writers import HDFWriter, NDFileHDF
+
+
+class DemoADSimDetector(StandardDetector):
+    _controller: ADSimController
+    _writer: HDFWriter
+
+    def __init__(
+        self,
+        drv: ADBase,
+        hdf: NDFileHDF,
+        directory_provider: DirectoryProvider,
+        name: str = "",
+        config_sigs: Sequence[SignalR] = (),
+    ):
+        self.drv = drv
+        self.hdf = hdf
+
+        super().__init__(
+            ADSimController(self.drv),
+            HDFWriter(
+                self.hdf,
+                directory_provider,
+                lambda: self.name,
+                ADBaseShapeProvider(self.drv),
+            ),
+            config_sigs=config_sigs,
+            name=name,
+        )

ophyd_async/epics/motion/motor.py

@@ -81,4 +81,5 @@ class Motor(StandardReadable, Movable, Stoppable):
         self._set_success = success
         # Put with completion will never complete as we are waiting for completion on
         # the move above, so need to pass wait=False
-        await self.stop_.execute(wait=False)
+        status = self.stop_.trigger(wait=False)
+        await status

ophyd_async/epics/pvi.py

@@ -0,0 +1,70 @@
+from typing import Callable, Dict, FrozenSet, Optional, Type, TypedDict, TypeVar
+
+from ophyd_async.core.signal import Signal
+from ophyd_async.core.signal_backend import SignalBackend
+from ophyd_async.core.utils import DEFAULT_TIMEOUT
+from ophyd_async.epics._backend._p4p import PvaSignalBackend
+from ophyd_async.epics.signal.signal import (
+    epics_signal_r,
+    epics_signal_rw,
+    epics_signal_w,
+    epics_signal_x,
+)
+
+T = TypeVar("T")
+
+
+_pvi_mapping: Dict[FrozenSet[str], Callable[..., Signal]] = {
+    frozenset({"r", "w"}): lambda dtype, read_pv, write_pv: epics_signal_rw(
+        dtype, read_pv, write_pv
+    ),
+    frozenset({"rw"}): lambda dtype, read_pv, write_pv: epics_signal_rw(
+        dtype, read_pv, write_pv
+    ),
+    frozenset({"r"}): lambda dtype, read_pv, _: epics_signal_r(dtype, read_pv),
+    frozenset({"w"}): lambda dtype, _, write_pv: epics_signal_w(dtype, write_pv),
+    frozenset({"x"}): lambda _, __, write_pv: epics_signal_x(write_pv),
+}
+
+
+class PVIEntry(TypedDict, total=False):
+    d: str
+    r: str
+    rw: str
+    w: str
+    x: str
+
+
+async def pvi_get(
+    read_pv: str, timeout: float = DEFAULT_TIMEOUT
+) -> Dict[str, PVIEntry]:
+    """Makes a PvaSignalBackend purely to connect to PVI information.
+
+    This backend is simply thrown away at the end of this method. This is useful
+    because the backend handles a CancelledError exception that gets thrown on
+    timeout, and therefore can be used for error reporting."""
+    backend: SignalBackend = PvaSignalBackend(None, read_pv, read_pv)
+    await backend.connect(timeout=timeout)
+    d: Dict[str, Dict[str, Dict[str, str]]] = await backend.get_value()
+    pv_info = d.get("pvi") or {}
+    result = {}
+
+    for attr_name, attr_info in pv_info.items():
+        result[attr_name] = PVIEntry(**attr_info)  # type: ignore
+
+    return result
+
+
+def make_signal(signal_pvi: PVIEntry, dtype: Optional[Type[T]] = None) -> Signal[T]:
+    """Make a signal.
+
+    This assumes datatype is None so it can be used to create dynamic signals.
+    """
+    operations = frozenset(signal_pvi.keys())
+    pvs = [signal_pvi[i] for i in operations]  # type: ignore
+    signal_factory = _pvi_mapping[operations]
+
+    write_pv = "pva://" + pvs[0]
+    read_pv = write_pv if len(pvs) < 2 else "pva://" + pvs[1]
+
+    return signal_factory(dtype, read_pv, write_pv)

ophyd_async/epics/signal/__init__.py

@@ -1,8 +1,6 @@
-from .pvi_get import pvi_get
 from .signal import epics_signal_r, epics_signal_rw, epics_signal_w, epics_signal_x
 
 __all__ = [
-    "pvi_get",
     "epics_signal_r",
     "epics_signal_rw",
     "epics_signal_w",