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

ophyd_async/sim/__init__.py

@@ -0,0 +1,11 @@
+from .pattern_generator import PatternGenerator
+from .sim_pattern_detector_control import SimPatternDetectorControl
+from .sim_pattern_detector_writer import SimPatternDetectorWriter
+from .sim_pattern_generator import SimPatternDetector
+
+__all__ = [
+    "PatternGenerator",
+    "SimPatternDetectorControl",
+    "SimPatternDetectorWriter",
+    "SimPatternDetector",
+]

ophyd_async/sim/demo/__init__.py

@@ -0,0 +1,3 @@
+from .sim_motor import SimMotor
+
+__all__ = ["SimMotor"]

ophyd_async/sim/demo/sim_motor.py

@@ -0,0 +1,103 @@
+import asyncio
+import contextlib
+import time
+
+from bluesky.protocols import Movable, Stoppable
+
+from ophyd_async.core import StandardReadable
+from ophyd_async.core.async_status import AsyncStatus, WatchableAsyncStatus
+from ophyd_async.core.signal import (
+    observe_value,
+    soft_signal_r_and_setter,
+    soft_signal_rw,
+)
+from ophyd_async.core.standard_readable import ConfigSignal, HintedSignal
+from ophyd_async.core.utils import WatcherUpdate
+
+
+class SimMotor(StandardReadable, Movable, Stoppable):
+    def __init__(self, name="", instant=True) -> None:
+        """
+        Simulated motor device
+
+        args:
+        - prefix: str: Signal names prefix
+        - name: str: name of device
+        - instant: bool: whether to move instantly, or with a delay
+        """
+        # Define some signals
+        with self.add_children_as_readables(HintedSignal):
+            self.user_readback, self._user_readback_set = soft_signal_r_and_setter(
+                float, 0
+            )
+        with self.add_children_as_readables(ConfigSignal):
+            self.velocity = soft_signal_rw(float, 0 if instant else 1.0)
+            self.units = soft_signal_rw(str, "mm")
+        self.user_setpoint = soft_signal_rw(float, 0)
+
+        # Whether set() should complete successfully or not
+        self._set_success = True
+        self._move_status: AsyncStatus | None = None
+
+        super().__init__(name=name)
+
+    async def _move(self, old_position: float, new_position: float, move_time: float):
+        start = time.monotonic()
+        distance = abs(new_position - old_position)
+        while True:
+            time_elapsed = round(time.monotonic() - start, 2)
+
+            # update position based on time elapsed
+            if time_elapsed >= move_time:
+                # successfully reached our target position
+                self._user_readback_set(new_position)
+                break
+            else:
+                current_position = old_position + distance * time_elapsed / move_time
+
+            self._user_readback_set(current_position)
+
+            # 10hz update loop
+            await asyncio.sleep(0.1)
+
+    @WatchableAsyncStatus.wrap
+    async def set(self, new_position: float):
+        """
+        Asynchronously move the motor to a new position.
+        """
+        # Make sure any existing move tasks are stopped
+        await self.stop()
+        old_position, units, velocity = await asyncio.gather(
+            self.user_setpoint.get_value(),
+            self.units.get_value(),
+            self.velocity.get_value(),
+        )
+        # If zero velocity, do instant move
+        move_time = abs(new_position - old_position) / velocity if velocity else 0
+        self._move_status = AsyncStatus(
+            self._move(old_position, new_position, move_time)
+        )
+        # If stop is called then this will raise a CancelledError, ignore it
+        with contextlib.suppress(asyncio.CancelledError):
+            async for current_position in observe_value(
+                self.user_readback, done_status=self._move_status
+            ):
+                yield WatcherUpdate(
+                    current=current_position,
+                    initial=old_position,
+                    target=new_position,
+                    name=self.name,
+                    unit=units,
+                )
+        if not self._set_success:
+            raise RuntimeError("Motor was stopped")
+
+    async def stop(self, success=True):
+        """
+        Stop the motor if it is moving
+        """
+        self._set_success = success
+        if self._move_status:
+            self._move_status.task.cancel()
+            self._move_status = None
+        await self.user_setpoint.set(await self.user_readback.get_value())

ophyd_async/sim/pattern_generator.py

@@ -0,0 +1,318 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import (
+    Any,
+    AsyncGenerator,
+    AsyncIterator,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+)
+
+import h5py
+import numpy as np
+from bluesky.protocols import DataKey, StreamAsset
+from event_model import (
+    ComposeStreamResource,
+    ComposeStreamResourceBundle,
+    StreamDatum,
+    StreamRange,
+    StreamResource,
+)
+
+from ophyd_async.core import DirectoryInfo, DirectoryProvider
+from ophyd_async.core.mock_signal_backend import MockSignalBackend
+from ophyd_async.core.signal import SignalR, observe_value
+from ophyd_async.core.utils import DEFAULT_TIMEOUT
+
+# raw data path
+DATA_PATH = "/entry/data/data"
+
+# pixel sum path
+SUM_PATH = "/entry/sum"
+
+MAX_UINT8_VALUE = np.iinfo(np.uint8).max
+
+SLICE_NAME = "AD_HDF5_SWMR_SLICE"
+
+
+@dataclass
+class DatasetConfig:
+    name: str
+    shape: Sequence[int]
+    maxshape: tuple[Any, ...] = (None,)
+    path: Optional[str] = None
+    multiplier: Optional[int] = 1
+    dtype: Optional[Any] = None
+    fillvalue: Optional[int] = None
+
+
+def get_full_file_description(
+    datasets: List[DatasetConfig], outer_shape: tuple[int, ...]
+):
+    full_file_description: Dict[str, DataKey] = {}
+    for d in datasets:
+        source = f"soft://{d.name}"
+        shape = outer_shape + tuple(d.shape)
+        dtype = "number" if d.shape == [1] else "array"
+        descriptor = DataKey(
+            source=source, shape=shape, dtype=dtype, external="STREAM:"
+        )
+        key = d.name.replace("/", "_")
+        full_file_description[key] = descriptor
+    return full_file_description
+
+
+def generate_gaussian_blob(height: int, width: int) -> np.ndarray:
+    """Make a Gaussian Blob with float values in range 0..1"""
+    x, y = np.meshgrid(np.linspace(-1, 1, width), np.linspace(-1, 1, height))
+    d = np.sqrt(x * x + y * y)
+    blob = np.exp(-(d**2))
+    return blob
+
+
+def generate_interesting_pattern(x: float, y: float) -> float:
+    """This function is interesting in x and y in range -10..10, returning
+    a float value in range 0..1
+    """
+    z = 0.5 + (np.sin(x) ** 10 + np.cos(10 + y * x) * np.cos(x)) / 2
+    return z
+
+
+class HdfStreamProvider:
+    def __init__(
+        self,
+        directory_info: DirectoryInfo,
+        full_file_name: Path,
+        datasets: List[DatasetConfig],
+    ) -> None:
+        self._last_emitted = 0
+        self._bundles: List[ComposeStreamResourceBundle] = self._compose_bundles(
+            directory_info, full_file_name, datasets
+        )
+
+    def _compose_bundles(
+        self,
+        directory_info: DirectoryInfo,
+        full_file_name: Path,
+        datasets: List[DatasetConfig],
+    ) -> List[StreamAsset]:
+        path = str(full_file_name.relative_to(directory_info.root))
+        root = str(directory_info.root)
+        bundler_composer = ComposeStreamResource()
+
+        bundles: List[ComposeStreamResourceBundle] = []
+
+        bundles = [
+            bundler_composer(
+                spec=SLICE_NAME,
+                root=root,
+                resource_path=path,
+                data_key=d.name.replace("/", "_"),
+                resource_kwargs={
+                    "path": d.path,
+                    "multiplier": d.multiplier,
+                    "timestamps": "/entry/instrument/NDAttributes/NDArrayTimeStamp",
+                },
+            )
+            for d in datasets
+        ]
+        return bundles
+
+    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:
+            updated_stream_range = StreamRange(
+                start=self._last_emitted,
+                stop=indices_written,
+            )
+            self._last_emitted = indices_written
+            for bundle in self._bundles:
+                yield bundle.compose_stream_datum(indices=updated_stream_range)
+        return None
+
+    def close(self) -> None:
+        for bundle in self._bundles:
+            bundle.close()
+
+
+class PatternGenerator:
+    def __init__(
+        self,
+        saturation_exposure_time: float = 1,
+        detector_width: int = 320,
+        detector_height: int = 240,
+    ) -> None:
+        self.saturation_exposure_time = saturation_exposure_time
+        self.exposure = saturation_exposure_time
+        self.x = 0.0
+        self.y = 0.0
+        self.height = detector_height
+        self.width = detector_width
+        self.written_images_counter: int = 0
+
+        # it automatically initializes to 0
+        self.signal_backend = MockSignalBackend(int)
+        self.mock_signal = SignalR(self.signal_backend)
+        blob = np.array(
+            generate_gaussian_blob(width=detector_width, height=detector_height)
+            * MAX_UINT8_VALUE
+        )
+        self.STARTING_BLOB = blob
+        self._hdf_stream_provider: Optional[HdfStreamProvider] = None
+        self._handle_for_h5_file: Optional[h5py.File] = None
+        self.target_path: Optional[Path] = None
+
+    async def write_image_to_file(self) -> None:
+        assert self._handle_for_h5_file, "no file has been opened!"
+        # prepare - resize the fixed hdf5 data structure
+        # so that the new image can be written
+        new_layer = self.written_images_counter + 1
+        target_dimensions = (new_layer, self.height, self.width)
+
+        # generate the simulated data
+        intensity: float = generate_interesting_pattern(self.x, self.y)
+        detector_data: np.uint8 = np.uint8(
+            self.STARTING_BLOB
+            * intensity
+            * self.exposure
+            / self.saturation_exposure_time
+        )
+
+        self._handle_for_h5_file[DATA_PATH].resize(target_dimensions)
+
+        print(f"writing image {new_layer}")
+        assert self._handle_for_h5_file, "no file has been opened!"
+        self._handle_for_h5_file[DATA_PATH].resize(target_dimensions)
+
+        self._handle_for_h5_file[SUM_PATH].resize((new_layer,))
+
+        # write data to disc (intermediate step)
+        self._handle_for_h5_file[DATA_PATH][self.written_images_counter] = detector_data
+        self._handle_for_h5_file[SUM_PATH][self.written_images_counter] = np.sum(
+            detector_data
+        )
+
+        # save metadata - so that it's discoverable
+        self._handle_for_h5_file[DATA_PATH].flush()
+        self._handle_for_h5_file[SUM_PATH].flush()
+
+        # counter increment is last
+        # as only at this point the new data is visible from the outside
+        self.written_images_counter += 1
+        await self.signal_backend.put(self.written_images_counter)
+
+    def set_exposure(self, value: float) -> None:
+        self.exposure = value
+
+    def set_x(self, value: float) -> None:
+        self.x = value
+
+    def set_y(self, value: float) -> None:
+        self.y = value
+
+    async def open_file(
+        self, directory: DirectoryProvider, multiplier: int = 1
+    ) -> Dict[str, DataKey]:
+        await self.mock_signal.connect()
+
+        self.target_path = self._get_new_path(directory)
+
+        self._handle_for_h5_file = h5py.File(self.target_path, "w", libver="latest")
+
+        assert self._handle_for_h5_file, "not loaded the file right"
+
+        datasets = self._get_datasets()
+        for d in datasets:
+            self._handle_for_h5_file.create_dataset(
+                name=d.name,
+                shape=d.shape,
+                dtype=d.dtype,
+                maxshape=d.maxshape,
+            )
+
+        # once datasets written, can switch the model to single writer multiple reader
+        self._handle_for_h5_file.swmr_mode = True
+
+        outer_shape = (multiplier,) if multiplier > 1 else ()
+        full_file_description = get_full_file_description(datasets, outer_shape)
+
+        # cache state to self
+        self._datasets = datasets
+        self.multiplier = multiplier
+        self._directory_provider = directory
+        return full_file_description
+
+    def _get_new_path(self, directory: DirectoryProvider) -> Path:
+        info = directory()
+        filename = f"{info.prefix}pattern{info.suffix}.h5"
+        new_path: Path = info.root / info.resource_dir / filename
+        return new_path
+
+    def _get_datasets(self) -> List[DatasetConfig]:
+        raw_dataset = DatasetConfig(
+            # name=data_name,
+            name=DATA_PATH,
+            dtype=np.uint8,
+            shape=(1, self.height, self.width),
+            maxshape=(None, self.height, self.width),
+        )
+
+        sum_dataset = DatasetConfig(
+            name=SUM_PATH,
+            dtype=np.float64,
+            shape=(1,),
+            maxshape=(None,),
+            fillvalue=-1,
+        )
+
+        datasets: List[DatasetConfig] = [raw_dataset, sum_dataset]
+        return datasets
+
+    async def collect_stream_docs(
+        self, indices_written: int
+    ) -> AsyncIterator[StreamAsset]:
+        """
+        stream resource says "here is a dataset",
+        stream datum says "here are N frames in that stream resource",
+        you get one stream resource and many stream datums per scan
+        """
+        if self._handle_for_h5_file:
+            self._handle_for_h5_file.flush()
+        # when already something was written to the file
+        if indices_written:
+            # if no frames arrived yet, there's no file to speak of
+            # cannot get the full filename the HDF writer will write
+            # until the first frame comes in
+            if not self._hdf_stream_provider:
+                assert self.target_path, "open file has not been called"
+                datasets = self._get_datasets()
+                self._datasets = datasets
+                self._hdf_stream_provider = HdfStreamProvider(
+                    self._directory_provider(),
+                    self.target_path,
+                    self._datasets,
+                )
+                for doc in self._hdf_stream_provider.stream_resources():
+                    yield "stream_resource", doc
+            if self._hdf_stream_provider:
+                for doc in self._hdf_stream_provider.stream_data(indices_written):
+                    yield "stream_datum", doc
+
+    def close(self) -> None:
+        if self._handle_for_h5_file:
+            self._handle_for_h5_file.close()
+            print("file closed")
+            self._handle_for_h5_file = None
+
+    async def observe_indices_written(
+        self, timeout=DEFAULT_TIMEOUT
+    ) -> AsyncGenerator[int, None]:
+        async for num_captured in observe_value(self.mock_signal, timeout=timeout):
+            yield num_captured // self.multiplier

ophyd_async/sim/sim_pattern_detector_control.py

@@ -0,0 +1,55 @@
+import asyncio
+from typing import Optional
+
+from ophyd_async.core import DirectoryProvider
+from ophyd_async.core.async_status import AsyncStatus
+from ophyd_async.core.detector import DetectorControl, DetectorTrigger
+from ophyd_async.sim.pattern_generator import PatternGenerator
+
+
+class SimPatternDetectorControl(DetectorControl):
+    def __init__(
+        self,
+        pattern_generator: PatternGenerator,
+        directory_provider: DirectoryProvider,
+        exposure: float = 0.1,
+    ) -> None:
+        self.pattern_generator: PatternGenerator = pattern_generator
+        self.pattern_generator.set_exposure(exposure)
+        self.directory_provider: DirectoryProvider = directory_provider
+        self.task: Optional[asyncio.Task] = None
+        super().__init__()
+
+    async def arm(
+        self,
+        num: int,
+        trigger: DetectorTrigger = DetectorTrigger.internal,
+        exposure: Optional[float] = 0.01,
+    ) -> AsyncStatus:
+        assert exposure is not None
+        period: float = exposure + self.get_deadtime(exposure)
+        task = asyncio.create_task(
+            self._coroutine_for_image_writing(exposure, period, num)
+        )
+        self.task = task
+        return AsyncStatus(task)
+
+    async def disarm(self):
+        if self.task:
+            self.task.cancel()
+            try:
+                await self.task
+            except asyncio.CancelledError:
+                pass
+            self.task = None
+
+    def get_deadtime(self, exposure: float) -> float:
+        return 0.001
+
+    async def _coroutine_for_image_writing(
+        self, exposure: float, period: float, frames_number: int
+    ):
+        for _ in range(frames_number):
+            self.pattern_generator.set_exposure(exposure)
+            await asyncio.sleep(period)
+            await self.pattern_generator.write_image_to_file()

ophyd_async/sim/sim_pattern_detector_writer.py

@@ -0,0 +1,34 @@
+from typing import AsyncGenerator, AsyncIterator, Dict
+
+from bluesky.protocols import DataKey
+
+from ophyd_async.core import DirectoryProvider
+from ophyd_async.core.detector import DetectorWriter
+from ophyd_async.sim.pattern_generator import PatternGenerator
+
+
+class SimPatternDetectorWriter(DetectorWriter):
+    pattern_generator: PatternGenerator
+
+    def __init__(
+        self, pattern_generator: PatternGenerator, directoryProvider: DirectoryProvider
+    ) -> None:
+        self.pattern_generator = pattern_generator
+        self.directory_provider = directoryProvider
+
+    async def open(self, multiplier: int = 1) -> Dict[str, DataKey]:
+        return await self.pattern_generator.open_file(
+            self.directory_provider, multiplier
+        )
+
+    async def close(self) -> None:
+        self.pattern_generator.close()
+
+    def collect_stream_docs(self, indices_written: int) -> AsyncIterator:
+        return self.pattern_generator.collect_stream_docs(indices_written)
+
+    def observe_indices_written(self, timeout=...) -> AsyncGenerator[int, None]:
+        return self.pattern_generator.observe_indices_written()
+
+    async def get_indices_written(self) -> int:
+        return self.pattern_generator.written_images_counter

ophyd_async/sim/sim_pattern_generator.py

@@ -0,0 +1,37 @@
+from pathlib import Path
+from typing import Sequence
+
+from ophyd_async.core import DirectoryProvider, StaticDirectoryProvider
+from ophyd_async.core.detector import StandardDetector
+from ophyd_async.protocols import AsyncReadable
+from ophyd_async.sim.pattern_generator import PatternGenerator
+
+from .sim_pattern_detector_control import SimPatternDetectorControl
+from .sim_pattern_detector_writer import SimPatternDetectorWriter
+
+
+class SimPatternDetector(StandardDetector):
+    def __init__(
+        self,
+        path: Path,
+        config_sigs: Sequence[AsyncReadable] = [],
+        name: str = "sim_pattern_detector",
+        writer_timeout: float = 1,
+    ) -> None:
+        self.directory_provider: DirectoryProvider = StaticDirectoryProvider(path)
+        self.pattern_generator = PatternGenerator()
+        writer = SimPatternDetectorWriter(
+            pattern_generator=self.pattern_generator,
+            directoryProvider=self.directory_provider,
+        )
+        controller = SimPatternDetectorControl(
+            pattern_generator=self.pattern_generator,
+            directory_provider=self.directory_provider,
+        )
+        super().__init__(
+            controller=controller,
+            writer=writer,
+            config_sigs=config_sigs,
+            name=name,
+            writer_timeout=writer_timeout,
+        )