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

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())

ophyd_async/sim/_pattern_generator.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from pathlib import Path
+
+import h5py
+import numpy as np
+
+# raw data path
+DATA_PATH = "/entry/data/data"
+
+# pixel sum path
+SUM_PATH = "/entry/sum"
+
+
+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, channel: int, offset: float
+) -> float:
+    """Return a float value in range 0..1.
+
+    Interesting in x and y in range -10..10
+    """
+    return (np.sin(x) ** channel + np.cos(x * y + offset) + 2) / 4
+
+
+class PatternFile:
+    def __init__(
+        self,
+        path: Path,
+        width: int = 320,
+        height: int = 240,
+    ):
+        self.file = h5py.File(path, "w", libver="latest")
+        self.data = self.file.create_dataset(
+            name=DATA_PATH,
+            shape=(0, height, width),
+            dtype=np.uint8,
+            maxshape=(None, height, width),
+        )
+        self.sum = self.file.create_dataset(
+            name=SUM_PATH,
+            shape=(0,),
+            dtype=np.int64,
+            maxshape=(None,),
+        )
+        # Once datasets written, can switch the model to single writer multiple reader
+        self.file.swmr_mode = True
+        self.blob = generate_gaussian_blob(height, width) * np.iinfo(np.uint8).max
+        self.image_counter = 0
+        self.e = asyncio.Event()
+
+    def write_image_to_file(self, intensity: float):
+        data = np.floor(self.blob * intensity)
+        for dset, value in ((self.data, data), (self.sum, np.sum(data))):
+            dset.resize(self.image_counter + 1, axis=0)
+            dset[self.image_counter] = value
+            dset.flush()
+        self.image_counter += 1
+        self.e.set()
+        self.e.clear()
+
+    def close(self):
+        self.file.close()
+
+
+class PatternGenerator:
+    """Generates pattern images in files."""
+
+    def __init__(self, sleep=asyncio.sleep):
+        self._x = 0.0
+        self._y = 0.0
+        self._file: PatternFile | None = None
+        self.sleep = sleep
+
+    def set_x(self, x: float):
+        self._x = x
+
+    def set_y(self, y: float):
+        self._y = y
+
+    def generate_point(self, channel: int = 1, high_energy: bool = False) -> float:
+        """Make a point between 0 and 1 based on x and y."""
+        offset = 100 if high_energy else 10
+        return generate_interesting_pattern(self._x, self._y, channel, offset)
+
+    def open_file(self, path: Path, width: int, height: int):
+        self._file = PatternFile(path, width, height)
+
+    def _get_file(self) -> PatternFile:
+        if not self._file:
+            raise RuntimeError("open_file not run")
+        return self._file
+
+    async def write_images_to_file(
+        self, exposure: float, period: float, number_of_frames: int
+    ):
+        file = self._get_file()
+        start = time.monotonic()
+        for i in range(1, number_of_frames + 1):
+            deadline = start + i * period
+            timeout = deadline - time.monotonic()
+            await self.sleep(timeout)
+            intensity = self.generate_point() * exposure
+            file.write_image_to_file(intensity)
+
+    async def wait_for_next_index(self, timeout: float):
+        await asyncio.wait_for(self._get_file().e.wait(), timeout)
+
+    def get_last_index(self) -> int:
+        return self._get_file().image_counter
+
+    def close_file(self):
+        if self._file:
+            self._file.close()
+            self._file = None

ophyd_async/sim/_point_detector.py

@@ -0,0 +1,86 @@
+import asyncio
+import time
+
+import numpy as np
+
+from ophyd_async.core import (
+    AsyncStatus,
+    DeviceVector,
+    SignalR,
+    StandardReadable,
+    StrictEnum,
+    gather_dict,
+    soft_signal_r_and_setter,
+    soft_signal_rw,
+)
+from ophyd_async.core import StandardReadableFormat as Format
+
+from ._pattern_generator import PatternGenerator
+
+
+class EnergyMode(StrictEnum):
+    """Energy mode for `SimPointDetector`."""
+
+    LOW = "Low Energy"
+    """Low energy mode"""
+
+    HIGH = "High Energy"
+    """High energy mode"""
+
+
+class SimPointDetectorChannel(StandardReadable):
+    def __init__(self, value_signal: SignalR[int], name=""):
+        with self.add_children_as_readables(Format.HINTED_SIGNAL):
+            self.value = value_signal
+        with self.add_children_as_readables(Format.CONFIG_SIGNAL):
+            self.mode = soft_signal_rw(EnergyMode)
+        super().__init__(name)
+
+
+class SimPointDetector(StandardReadable):
+    """Simalutes a point detector with multiple channels."""
+
+    def __init__(
+        self, generator: PatternGenerator, num_channels: int = 3, name: str = ""
+    ) -> None:
+        self._generator = generator
+        self.acquire_time = soft_signal_rw(float, 0.1)
+        self.acquiring, self._set_acquiring = soft_signal_r_and_setter(bool)
+        self._value_signals = dict(
+            soft_signal_r_and_setter(int) for _ in range(num_channels)
+        )
+        with self.add_children_as_readables():
+            self.channel = DeviceVector(
+                {
+                    i + 1: SimPointDetectorChannel(value_signal)
+                    for i, value_signal in enumerate(self._value_signals)
+                }
+            )
+        super().__init__(name=name)
+
+    async def _update_values(self, acquire_time: float):
+        # Get the modes
+        modes = await gather_dict(
+            {channel: channel.mode.get_value() for channel in self.channel.values()}
+        )
+        start = time.monotonic()
+        # Make an array of relative update times at 10Hz intervals
+        update_times = np.arange(0.1, acquire_time, 0.1)
+        # With the end position appended
+        update_times = np.concatenate((update_times, [acquire_time]))
+        for update_time in update_times:
+            # Calculate how long to wait to get there
+            relative_time = time.monotonic() - start
+            await asyncio.sleep(update_time - relative_time)
+            # Update the channel value
+            for i, channel in self.channel.items():
+                high_energy = modes[channel] == EnergyMode.HIGH
+                point = self._generator.generate_point(i, high_energy)
+                setter = self._value_signals[channel.value]
+                setter(int(point * 10000 * update_time))
+
+    @AsyncStatus.wrap
+    async def trigger(self):
+        for setter in self._value_signals.values():
+            setter(0)
+        await self._update_values(await self.acquire_time.get_value())

ophyd_async/sim/_stage.py

@@ -0,0 +1,19 @@
+from ophyd_async.core import StandardReadable
+from ophyd_async.sim._pattern_generator import PatternGenerator
+
+from ._motor import SimMotor
+
+
+class SimStage(StandardReadable):
+    """A simulated sample stage with X and Y movables."""
+
+    def __init__(self, pattern_generator: PatternGenerator, name="") -> None:
+        # Define some child Devices
+        with self.add_children_as_readables():
+            self.x = SimMotor(instant=False)
+            self.y = SimMotor(instant=False)
+        # Tell the pattern generator about the motor positions
+        self.x.user_readback.subscribe_value(pattern_generator.set_x)
+        self.y.user_readback.subscribe_value(pattern_generator.set_y)
+        # Set name of device and child devices
+        super().__init__(name=name)

ophyd_async/tango/__init__.py

@@ -0,0 +1 @@
+"""Tango support for Signals, and Devices that use them."""

ophyd_async/tango/core/__init__.py

@@ -1,4 +1,4 @@
-from ._base_device import TangoDevice, TangoPolling
+from ._base_device import TangoDevice, TangoDeviceConnector, TangoPolling
 from ._signal import (
     infer_python_type,
     infer_signal_type,
@@ -19,10 +19,12 @@ from ._tango_transport import (
     get_tango_trl,
     get_trl_descriptor,
 )
+from ._utils import DevStateEnum, get_device_trl_and_attr, get_full_attr_trl
 
 __all__ = [
     "AttributeProxy",
     "CommandProxy",
+    "DevStateEnum",
     "ensure_proper_executor",
     "TangoSignalBackend",
     "get_python_type",
@@ -39,4 +41,7 @@ __all__ = [
     "TangoDevice",
     "TangoReadable",
     "TangoPolling",
+    "TangoDeviceConnector",
+    "get_device_trl_and_attr",
+    "get_full_attr_trl",
 ]

ophyd_async/tango/core/_base_device.py

@@ -4,27 +4,23 @@ from dataclasses import dataclass
 from typing import Any, Generic, TypeVar
 
 from ophyd_async.core import Device, DeviceConnector, DeviceFiller, LazyMock
-from tango import DeviceProxy as DeviceProxy
+from tango import DeviceProxy
 from tango.asyncio import DeviceProxy as AsyncDeviceProxy
 
 from ._signal import TangoSignalBackend, infer_python_type, infer_signal_type
+from ._utils import get_full_attr_trl
 
 T = TypeVar("T")
 
 
 class TangoDevice(Device):
-    """
-    General class for TangoDevices. Extends Device to provide attributes for Tango
-    devices.
-
-    Parameters
-    ----------
-    trl: str
-        Tango resource locator, typically of the device server.
-    device_proxy: Optional[Union[AsyncDeviceProxy, SyncDeviceProxy]]
-        Asynchronous or synchronous DeviceProxy object for the device. If not provided,
-        an asynchronous DeviceProxy object will be created using the trl and awaited
-        when the device is connected.
+    """General class for TangoDevices.
+
+    Extends Device to provide attributes for Tango devices.
+
+    :param trl: Tango resource locator, typically of the device server.
+        An asynchronous DeviceProxy object will be created using the
+        trl and awaited when the device is connected.
     """
 
     trl: str = ""
@@ -32,13 +28,15 @@ class TangoDevice(Device):
 
     def __init__(
         self,
-        trl: str | None = None,
-        device_proxy: DeviceProxy | None = None,
+        trl: str | None,
         support_events: bool = False,
         name: str = "",
+        auto_fill_signals: bool = True,
     ) -> None:
         connector = TangoDeviceConnector(
-            trl=trl, device_proxy=device_proxy, support_events=support_events
+            trl=trl,
+            support_events=support_events,
+            auto_fill_signals=auto_fill_signals,
         )
         super().__init__(name=name, connector=connector)
 
@@ -74,12 +72,12 @@ class TangoDeviceConnector(DeviceConnector):
     def __init__(
         self,
         trl: str | None,
-        device_proxy: DeviceProxy | None,
         support_events: bool,
+        auto_fill_signals: bool = True,
     ) -> None:
         self.trl = trl
-        self.proxy = device_proxy
         self._support_events = support_events
+        self._auto_fill_signals = auto_fill_signals
 
     def create_children_from_annotations(self, device: Device):
         if not hasattr(self, "filler"):
@@ -87,7 +85,7 @@ class TangoDeviceConnector(DeviceConnector):
                 device=device,
                 signal_backend_factory=TangoSignalBackend,
                 device_connector_factory=lambda: TangoDeviceConnector(
-                    None, None, self._support_events
+                    None, self._support_events
                 ),
             )
             list(self.filler.create_devices_from_annotations(filled=False))
@@ -105,28 +103,38 @@ class TangoDeviceConnector(DeviceConnector):
         return await super().connect_mock(device, mock)
 
     async def connect_real(self, device: Device, timeout: float, force_reconnect: bool):
-        if self.trl and self.proxy is None:
-            self.proxy = await AsyncDeviceProxy(self.trl)
-        elif self.proxy and not self.trl:
-            self.trl = self.proxy.name()
-        else:
-            raise TypeError("Neither proxy nor trl supplied")
-
+        if not self.trl:
+            raise RuntimeError(f"Could not created Device Proxy for TRL {self.trl}")
+        self.proxy = await AsyncDeviceProxy(self.trl)
         children = sorted(
             set()
             .union(self.proxy.get_attribute_list())
             .union(self.proxy.get_command_list())
         )
+
+        children = [
+            child for child in children if child not in self.filler.ignored_signals
+        ]
+
+        not_filled = {unfilled for unfilled, _ in device.children()}
+
+        # If auto_fill_signals is True, fill all children inferred from the device
+        # else fill only the children that are annotated
         for name in children:
-            # TODO: strip attribute name
-            full_trl = f"{self.trl}/{name}"
-            signal_type = await infer_signal_type(full_trl, self.proxy)
-            if signal_type:
-                backend = self.filler.fill_child_signal(name, signal_type)
-                backend.datatype = await infer_python_type(full_trl, self.proxy)
-                backend.set_trl(full_trl)
+            if self._auto_fill_signals or name in not_filled:
+                # TODO: strip attribute name
+                full_trl = get_full_attr_trl(self.trl, name)
+                signal_type = await infer_signal_type(full_trl, self.proxy)
+                if signal_type:
+                    backend = self.filler.fill_child_signal(name, signal_type)
+                    # don't overlaod datatype if provided by annotation
+                    if backend.datatype is None:
+                        backend.datatype = await infer_python_type(full_trl, self.proxy)
+                    backend.set_trl(full_trl)
+
         # Check that all the requested children have been filled
         self.filler.check_filled(f"{self.trl}: {children}")
+
         # Set the name of the device to name all children
         device.set_name(device.name)
         return await super().connect_real(device, timeout, force_reconnect)