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

ophyd_async/sim/__main__.py

@@ -0,0 +1,43 @@
+"""Used for tutorial `Using Devices`."""
+
+# Import bluesky and ophyd
+from tempfile import mkdtemp
+
+import bluesky.plan_stubs as bps  # noqa: F401
+import bluesky.plans as bp  # noqa: F401
+import bluesky.preprocessors as bpp  # noqa: F401
+from bluesky.callbacks.best_effort import BestEffortCallback
+from bluesky.run_engine import RunEngine, autoawait_in_bluesky_event_loop
+
+from ophyd_async import sim
+from ophyd_async.core import StaticPathProvider, UUIDFilenameProvider, init_devices
+
+# Create a run engine and make ipython use it for `await` commands
+RE = RunEngine(call_returns_result=True)
+autoawait_in_bluesky_event_loop()
+
+# Add a callback for plotting
+bec = BestEffortCallback()
+RE.subscribe(bec)
+
+# Make a pattern generator that uses the motor positions
+# to make a test pattern. This simulates the real life process
+# of X-ray scattering off a sample
+pattern_generator = sim.PatternGenerator()
+
+# Make a path provider that makes UUID filenames within a static
+# temporary directory
+path_provider = StaticPathProvider(UUIDFilenameProvider(), mkdtemp())
+
+# All Devices created within this block will be
+# connected and named at the end of the with block
+with init_devices():
+    # Create a sample stage with X and Y motors that report their positions
+    # to the pattern generator
+    stage = sim.SimStage(pattern_generator)
+    # Make a detector device that gives the point value of the pattern generator
+    # when triggered
+    pdet = sim.SimPointDetector(pattern_generator)
+    # Make a detector device that gives a gaussian blob with intensity based
+    # on the point value of the pattern generator when triggered
+    bdet = sim.SimBlobDetector(path_provider, pattern_generator)

ophyd_async/sim/_blob_detector.py

@@ -0,0 +1,33 @@
+from collections.abc import Sequence
+
+from ophyd_async.core import PathProvider, SignalR, StandardDetector
+
+from ._blob_detector_controller import BlobDetectorController
+from ._blob_detector_writer import BlobDetectorWriter
+from ._pattern_generator import PatternGenerator
+
+
+class SimBlobDetector(StandardDetector):
+    """Simulates a detector and writes Blobs to file."""
+
+    def __init__(
+        self,
+        path_provider: PathProvider,
+        pattern_generator: PatternGenerator | None = None,
+        config_sigs: Sequence[SignalR] = (),
+        name: str = "",
+    ) -> None:
+        self.pattern_generator = pattern_generator or PatternGenerator()
+
+        super().__init__(
+            controller=BlobDetectorController(
+                pattern_generator=self.pattern_generator,
+            ),
+            writer=BlobDetectorWriter(
+                pattern_generator=self.pattern_generator,
+                path_provider=path_provider,
+                name_provider=lambda: self.name,
+            ),
+            config_sigs=config_sigs,
+            name=name,
+        )

ophyd_async/sim/_blob_detector_controller.py

@@ -0,0 +1,48 @@
+import asyncio
+from contextlib import suppress
+
+from ophyd_async.core import DetectorController, DetectorTrigger, TriggerInfo
+
+from ._pattern_generator import PatternGenerator
+
+
+class BlobDetectorController(DetectorController):
+    def __init__(self, pattern_generator: PatternGenerator):
+        self.pattern_generator = pattern_generator
+        self.trigger_info: TriggerInfo | None = None
+        self.task: asyncio.Task | None = None
+
+    def get_deadtime(self, exposure):
+        return 0.001
+
+    async def prepare(self, trigger_info: TriggerInfo):
+        # This is a simulation, so only support intenal triggering
+        if trigger_info.trigger != DetectorTrigger.INTERNAL:
+            raise RuntimeError(f"{trigger_info.trigger} not supported by {self}")
+        # Just hold onto the trigger info until we need it
+        self.trigger_info = trigger_info
+
+    async def arm(self):
+        if self.trigger_info is None:
+            raise RuntimeError(f"prepare() not called on {self}")
+        livetime = self.trigger_info.livetime or 0.1
+        # Start a background process off writing the images to file
+        coro = self.pattern_generator.write_images_to_file(
+            exposure=livetime,
+            period=livetime + self.trigger_info.deadtime,
+            number_of_frames=self.trigger_info.total_number_of_triggers,
+        )
+        self.task = asyncio.create_task(coro)
+
+    async def wait_for_idle(self):
+        # Wait for the background task to complete
+        if self.task:
+            await self.task
+
+    async def disarm(self):
+        # Stop the background task and wait for it to finish
+        if self.task:
+            self.task.cancel()
+            with suppress(asyncio.CancelledError):
+                await self.task
+            self.task = None

ophyd_async/sim/_blob_detector_writer.py

@@ -0,0 +1,105 @@
+from collections.abc import AsyncGenerator, AsyncIterator
+from pathlib import Path
+
+import numpy as np
+from bluesky.protocols import Hints, StreamAsset
+from event_model import DataKey
+
+from ophyd_async.core import (
+    DetectorWriter,
+    HDFDatasetDescription,
+    HDFDocumentComposer,
+    NameProvider,
+    PathProvider,
+)
+
+from ._pattern_generator import DATA_PATH, SUM_PATH, PatternGenerator
+
+WIDTH = 320
+HEIGHT = 240
+
+
+class BlobDetectorWriter(DetectorWriter):
+    def __init__(
+        self,
+        pattern_generator: PatternGenerator,
+        path_provider: PathProvider,
+        name_provider: NameProvider,
+    ) -> None:
+        self.pattern_generator = pattern_generator
+        self.path_provider = path_provider
+        self.name_provider = name_provider
+        self.path: Path | None = None
+        self.composer: HDFDocumentComposer | None = None
+        self.datasets: list[HDFDatasetDescription] = []
+
+    async def open(self, multiplier: int = 1) -> dict[str, DataKey]:
+        name = self.name_provider()
+        path_info = self.path_provider(name)
+        self.path = path_info.directory_path / f"{path_info.filename}.h5"
+        self.pattern_generator.open_file(self.path, WIDTH, HEIGHT)
+        # We know it will write data and sum, so emit those
+        self.datasets = [
+            HDFDatasetDescription(
+                data_key=name,
+                dataset=DATA_PATH,
+                shape=(HEIGHT, WIDTH),
+                dtype_numpy=np.dtype(np.uint8).str,
+                chunk_shape=(HEIGHT, WIDTH),
+                multiplier=multiplier,
+            ),
+            HDFDatasetDescription(
+                data_key=f"{name}-sum",
+                dataset=SUM_PATH,
+                shape=(),
+                dtype_numpy=np.dtype(np.int64).str,
+                multiplier=multiplier,
+                chunk_shape=(1024,),
+            ),
+        ]
+        self.composer = None
+        outer_shape = (multiplier,) if multiplier > 1 else ()
+        describe = {
+            ds.data_key: DataKey(
+                source="sim://pattern-generator-hdf-file",
+                shape=list(outer_shape) + list(ds.shape),
+                dtype="array" if ds.shape else "number",
+                external="STREAM:",
+            )
+            for ds in self.datasets
+        }
+        return describe
+
+    @property
+    def hints(self) -> Hints:
+        """The hints to be used for the detector."""
+        return {"fields": [self.name_provider()]}
+
+    async def get_indices_written(self) -> int:
+        return self.pattern_generator.get_last_index()
+
+    async def observe_indices_written(
+        self, timeout: float
+    ) -> AsyncGenerator[int, None]:
+        while True:
+            yield self.pattern_generator.get_last_index()
+            await self.pattern_generator.wait_for_next_index(timeout)
+
+    async def collect_stream_docs(
+        self, indices_written: int
+    ) -> AsyncIterator[StreamAsset]:
+        # When we have written something to the file
+        if indices_written:
+            # Only emit stream resource the first time we see frames in
+            # the file
+            if not self.composer:
+                if not self.path:
+                    raise RuntimeError(f"open() not called on {self}")
+                self.composer = HDFDocumentComposer(self.path, self.datasets)
+                for doc in self.composer.stream_resources():
+                    yield "stream_resource", doc
+            for doc in self.composer.stream_data(indices_written):
+                yield "stream_datum", doc
+
+    async def close(self) -> None:
+        self.pattern_generator.close_file()

ophyd_async/sim/_mirror_horizontal.py

@@ -0,0 +1,46 @@
+import asyncio
+from typing import TypedDict
+
+from bluesky.protocols import Movable
+
+from ophyd_async.core import AsyncStatus, DerivedSignalFactory, Device, soft_signal_rw
+
+from ._mirror_vertical import TwoJackDerived, TwoJackTransform
+from ._motor import SimMotor
+
+
+class HorizontalMirrorDerived(TypedDict):
+    x: float
+    roll: float
+
+
+class HorizontalMirror(Device, Movable):
+    def __init__(self, name=""):
+        # Raw signals
+        self.x1 = SimMotor()
+        self.x2 = SimMotor()
+        # Parameter
+        self.x1_x2_distance = soft_signal_rw(float, initial_value=1)
+        # Derived signals
+        self._factory = DerivedSignalFactory(
+            TwoJackTransform,
+            self._set_mirror,
+            jack1=self.x1,
+            jack2=self.x2,
+            distance=self.x1_x2_distance,
+        )
+        self.x = self._factory.derived_signal_rw(float, "height")
+        self.roll = self._factory.derived_signal_rw(float, "angle")
+        super().__init__(name=name)
+
+    async def _set_mirror(self, derived: TwoJackDerived) -> None:
+        transform = await self._factory.transform()
+        raw = transform.derived_to_raw(**derived)
+        await asyncio.gather(
+            self.x1.set(raw["jack1"]),
+            self.x2.set(raw["jack2"]),
+        )
+
+    @AsyncStatus.wrap
+    async def set(self, value: HorizontalMirrorDerived) -> None:
+        await self._set_mirror(TwoJackDerived(height=value["x"], angle=value["roll"]))

ophyd_async/sim/_mirror_vertical.py

@@ -0,0 +1,74 @@
+import asyncio
+import math
+from typing import TypedDict
+
+from bluesky.protocols import Movable
+
+from ophyd_async.core import (
+    AsyncStatus,
+    DerivedSignalFactory,
+    Device,
+    Transform,
+    soft_signal_rw,
+)
+
+from ._motor import SimMotor
+
+
+class TwoJackRaw(TypedDict):
+    jack1: float
+    jack2: float
+
+
+class TwoJackDerived(TypedDict):
+    height: float
+    angle: float
+
+
+class TwoJackTransform(Transform):
+    distance: float
+
+    def raw_to_derived(self, *, jack1: float, jack2: float) -> TwoJackDerived:
+        diff = jack2 - jack1
+        return TwoJackDerived(
+            height=jack1 + diff / 2,
+            # need the cast as returns numpy float rather than float64, but this
+            # is ok at runtime
+            angle=math.atan(diff / self.distance),
+        )
+
+    def derived_to_raw(self, *, height: float, angle: float) -> TwoJackRaw:
+        diff = math.tan(angle) * self.distance
+        return TwoJackRaw(
+            jack1=height - diff / 2,
+            jack2=height + diff / 2,
+        )
+
+
+class VerticalMirror(Device, Movable[TwoJackDerived]):
+    def __init__(self, name=""):
+        # Raw signals
+        self.y1 = SimMotor()
+        self.y2 = SimMotor()
+        # Parameter
+        self.y1_y2_distance = soft_signal_rw(float, initial_value=1)
+        # Derived signals
+        self._factory = DerivedSignalFactory(
+            TwoJackTransform,
+            self.set,
+            jack1=self.y1,
+            jack2=self.y2,
+            distance=self.y1_y2_distance,
+        )
+        self.height = self._factory.derived_signal_rw(float, "height")
+        self.angle = self._factory.derived_signal_rw(float, "angle")
+        super().__init__(name=name)
+
+    @AsyncStatus.wrap
+    async def set(self, derived: TwoJackDerived) -> None:  # type: ignore until bluesky 1.13.2
+        transform = await self._factory.transform()
+        raw = transform.derived_to_raw(**derived)
+        await asyncio.gather(
+            self.y1.set(raw["jack1"]),
+            self.y2.set(raw["jack2"]),
+        )

ophyd_async/sim/_motor.py

@@ -0,0 +1,233 @@
+import asyncio
+import contextlib
+import time
+
+import numpy as np
+from bluesky.protocols import Location, Reading, Stoppable, Subscribable
+from pydantic import BaseModel, ConfigDict, Field
+
+from ophyd_async.core import (
+    AsyncStatus,
+    Callback,
+    StandardReadable,
+    WatchableAsyncStatus,
+    WatcherUpdate,
+    observe_value,
+    soft_signal_r_and_setter,
+    soft_signal_rw,
+)
+from ophyd_async.core import StandardReadableFormat as Format
+
+
+class FlySimMotorInfo(BaseModel):
+    """Minimal set of information required to fly a [](#SimMotor)."""
+
+    model_config = ConfigDict(frozen=True)
+
+    cv_start: float
+    """Absolute position of the motor once it finishes accelerating to desired
+    velocity, in motor EGUs"""
+
+    cv_end: float
+    """Absolute position of the motor once it begins decelerating from desired
+    velocity, in EGUs"""
+
+    cv_time: float = Field(gt=0)
+    """Time taken for the motor to get from start_position to end_position, excluding
+    run-up and run-down, in seconds."""
+
+    @property
+    def velocity(self) -> float:
+        """Calculate the velocity of the constant velocity phase."""
+        return (self.cv_end - self.cv_start) / self.cv_time
+
+    def start_position(self, acceleration_time: float) -> float:
+        """Calculate the start position with run-up distance added on."""
+        return self.cv_start - acceleration_time * self.velocity / 2
+
+    def end_position(self, acceleration_time: float) -> float:
+        """Calculate the end position with run-down distance added on."""
+        return self.cv_end + acceleration_time * self.velocity / 2
+
+
+class SimMotor(StandardReadable, Stoppable, Subscribable[float]):
+    """For usage when simulating a motor."""
+
+    def __init__(self, name="", instant=True) -> None:
+        """Simulate a motor, with optional velocity.
+
+        :param name: name of device
+        :param instant: whether to move instantly or calculate move time using velocity
+        """
+        # Define some signals
+        with self.add_children_as_readables(Format.HINTED_SIGNAL):
+            self.user_readback, self._user_readback_set = soft_signal_r_and_setter(
+                float, 0
+            )
+        with self.add_children_as_readables(Format.CONFIG_SIGNAL):
+            self.velocity = soft_signal_rw(float, 0 if instant else 1.0)
+            self.acceleration_time = soft_signal_rw(float, 0.5)
+            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
+        # Stored in prepare
+        self._fly_info: FlySimMotorInfo | None = None
+        # Set on kickoff(), complete when motor reaches end position
+        self._fly_status: WatchableAsyncStatus | None = None
+
+        super().__init__(name=name)
+
+    def set_name(self, name: str, *, child_name_separator: str | None = None) -> None:
+        super().set_name(name, child_name_separator=child_name_separator)
+        # Readback should be named the same as its parent in read()
+        self.user_readback.set_name(name)
+
+    @AsyncStatus.wrap
+    async def prepare(self, value: FlySimMotorInfo):
+        """Calculate run-up and move there, setting fly velocity when there."""
+        self._fly_info = value
+        # Move to start as fast as we can
+        await self.velocity.set(0)
+        await self.set(value.start_position(await self.acceleration_time.get_value()))
+        # Set the velocity for the actual move
+        await self.velocity.set(value.velocity)
+
+    async def locate(self) -> Location[float]:
+        """Return the current setpoint and readback of the motor."""
+        setpoint, readback = await asyncio.gather(
+            self.user_setpoint.get_value(), self.user_readback.get_value()
+        )
+        return Location(setpoint=setpoint, readback=readback)
+
+    def subscribe(self, function: Callback[dict[str, Reading[float]]]) -> None:
+        self.user_readback.subscribe(function)
+
+    def clear_sub(self, function: Callback[dict[str, Reading[float]]]) -> None:
+        self.user_readback.clear_sub(function)
+
+    @AsyncStatus.wrap
+    async def kickoff(self):
+        """Begin moving motor from prepared position to final position."""
+        if not self._fly_info:
+            msg = "Motor must be prepared before attempting to kickoff"
+            raise RuntimeError(msg)
+        acceleration_time = await self.acceleration_time.get_value()
+        self._fly_status = self.set(self._fly_info.end_position(acceleration_time))
+        # Wait for the acceleration time to ensure we are at velocity
+        await asyncio.sleep(acceleration_time)
+
+    def complete(self) -> WatchableAsyncStatus:
+        """Mark as complete once motor reaches completed position."""
+        if not self._fly_status:
+            msg = "kickoff not called"
+            raise RuntimeError(msg)
+        return self._fly_status
+
+    async def _move(self, old_position: float, new_position: float, velocity: float):
+        if old_position == new_position:
+            return
+        start = time.monotonic()
+        acceleration_time = abs(await self.acceleration_time.get_value())
+        sign = np.sign(new_position - old_position)
+        velocity = abs(velocity) * sign
+        # The total distance to move
+        total_distance = new_position - old_position
+        # The ramp distance is the distance taken to ramp up (the same distance
+        # is taken to ramp down). This is the area under the triangle of the
+        # velocity ramp up (base * height / 2)
+        ramp_distance = acceleration_time * velocity / 2
+        if abs(ramp_distance * 2) >= abs(total_distance):
+            # All time is ramp up and down, so recalculate the maximum velocity
+            # we get to. We know the area under the ramp up triangle is half the
+            # total distance, and we also know the ratio of velocity over
+            # acceleration_time is the same as the ration of max_velocity over
+            # ramp_time, so solve the simultaneous equations to get
+            # max_velocity and ramp_time.
+            max_velocity = np.sqrt(total_distance * velocity / acceleration_time) * sign
+            ramp_time = total_distance / max_velocity
+            # So move time is just the ramp up and ramp down with no constant
+            # velocity section
+            move_time = 2 * ramp_time
+        else:
+            # Middle segments of constant velocity
+            max_velocity = velocity
+            # Ramp up and down time is exactly the requested acceleration time
+            ramp_time = acceleration_time
+            # So move time is twice this, plus the time taken to move the
+            # remaining distance at constant velocity
+            move_time = ramp_time * 2 + (total_distance - ramp_distance * 2) / velocity
+        # Make an array of relative update times at 10Hz intervals
+        update_times = list(np.arange(0.1, move_time, 0.1, dtype=float))
+        # With the end position appended
+        if update_times and np.isclose(update_times[-1], move_time):
+            update_times[-1] = move_time
+        else:
+            update_times.append(move_time)
+        # Iterate through the update times, calculating new position for each
+        for t in update_times:
+            if t <= ramp_time:
+                # Ramp up phase, calculate area under the ramp up triangle
+                current_velocity = t / ramp_time * max_velocity
+                position = old_position + current_velocity * t / 2
+            elif t >= move_time - ramp_time:
+                # Ramp down phase, subtract area under the ramp down triangle
+                time_left = move_time - t
+                current_velocity = time_left / ramp_time * max_velocity
+                position = new_position - current_velocity * time_left / 2
+            else:
+                # Constant velocity phase
+                position = old_position + ramp_distance + (t - ramp_time) * max_velocity
+            # Calculate how long to wait to get there
+            relative_time = time.monotonic() - start
+            await asyncio.sleep(t - relative_time)
+            # Update the readback position
+            self._user_readback_set(position)
+
+    @WatchableAsyncStatus.wrap
+    async def set(self, value: float):
+        """Asynchronously move the motor to a new position."""
+        new_position = value
+        # Make sure any existing move tasks are stopped
+        if self._move_status:
+            self._move_status.task.cancel()
+            self._move_status = None
+        # work out where we were
+        old_position, units, velocity = await asyncio.gather(
+            self.user_setpoint.get_value(),
+            self.units.get_value(),
+            self.velocity.get_value(),
+        )
+        # update the setpoint to where we want to be
+        await self.user_setpoint.set(new_position)
+        # If zero velocity, do instant move
+        if velocity == 0:
+            self._user_readback_set(new_position)
+        else:
+            self._move_status = AsyncStatus(
+                self._move(old_position, new_position, velocity)
+            )
+            # 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())