dls-dodal 1.29.4__py3-none-any.whl → 1.30.0__py3-none-any.whl

dodal/devices/motors.py

@@ -3,8 +3,34 @@ from ophyd_async.epics.motion import Motor
 
 
 class XYZPositioner(Device):
-    def __init__(self, prefix: str, name: str):
-        self.x = Motor(prefix + "X")
-        self.y = Motor(prefix + "Y")
-        self.z = Motor(prefix + "Z")
-        super().__init__(name)
+    """
+
+    Standard ophyd_async xyz motor stage, by combining 3 Motors,
+    with added infix for extra flexibliy to allow different axes other than x,y,z.
+
+    Parameters
+    ----------
+    prefix:
+        EPICS PV (Common part up to and including :).
+    name:
+        name for the stage.
+    infix:
+        EPICS PV, default is the ["X", "Y", "Z"].
+    Notes
+    -----
+    Example usage::
+        async with DeviceCollector():
+            xyz_stage = XYZPositioner("BLXX-MO-STAGE-XX:")
+    Or::
+        with DeviceCollector():
+            xyz_stage = XYZPositioner("BLXX-MO-STAGE-XX:", suffix = ["A", "B", "C"])
+
+    """
+
+    def __init__(self, prefix: str, name: str, infix: list[str] | None = None):
+        if infix is None:
+            infix = ["X", "Y", "Z"]
+        self.x = Motor(prefix + infix[0])
+        self.y = Motor(prefix + infix[1])
+        self.z = Motor(prefix + infix[2])
+        super().__init__(name=name)

dodal/devices/oav/grid_overlay.py

@@ -1,3 +1,4 @@
+# type: ignore # OAV will soon be ophyd-async, see https://github.com/DiamondLightSource/dodal/issues/716
 from enum import Enum
 from functools import partial
 from os.path import join as path_join

dodal/devices/oav/microns_for_zoom_levels.json

@@ -52,4 +52,4 @@
         "micronsPerXPixel": 0.227,
         "micronsPerYPixel": 0.314
     }
-}
+}

dodal/devices/oav/oav_detector.py

@@ -1,3 +1,4 @@
+# type: ignore # OAV will soon be ophyd-async, see https://github.com/DiamondLightSource/dodal/issues/716
 from functools import partial
 
 from ophyd import ADComponent as ADC
@@ -46,7 +47,7 @@ class ZoomController(Device):
     sxst = Component(EpicsSignal, "MP:SELECT.SXST")
 
     def set_flatfield_on_zoom_level_one(self, value):
-        self.parent: "OAV"
+        self.parent: OAV
         flat_applied = self.parent.proc.port_name.get()
         no_flat_applied = self.parent.cam.port_name.get()
         return self.parent.grid_snapshot.input_plugin.set(

dodal/devices/oav/oav_parameters.py

@@ -1,7 +1,8 @@
 import json
-import xml.etree.cElementTree as et
+import xml.etree.ElementTree as et
 from collections import ChainMap
-from typing import Any, Tuple
+from typing import Any
+from xml.etree.ElementTree import Element
 
 from dodal.devices.oav.oav_errors import (
     OAVError_BeamPositionNotFound,
@@ -20,6 +21,13 @@ OAV_CONFIG_JSON = (
 )
 
 
+def _get_element_as_float(node: Element, element_name: str) -> float:
+    element = node.find(element_name)
+    assert element is not None, f"{element_name} not found in {node}"
+    assert element.text
+    return float(element.text)
+
+
 class OAVParameters:
     """
     The parameters to set up the OAV depending on the context.
@@ -65,11 +73,11 @@ class OAVParameters:
             try:
                 param = param_type(param)
                 return param
-            except AssertionError:
+            except AssertionError as e:
                 raise TypeError(
                     f"OAV param {name} from the OAV centring params json file has the "
                     f"wrong type, should be {param_type} but is {type(param)}."
-                )
+                ) from e
 
         self.exposure: float = update("exposure", float)
         self.acquire_period: float = update("acqPeriod", float)
@@ -134,14 +142,14 @@ class OAVConfigParams:
         root = tree.getroot()
         levels = root.findall(".//zoomLevel")
         for node in levels:
-            if float(node.find("level").text) == zoom:
+            if _get_element_as_float(node, "level") == zoom:
                 self.micronsPerXPixel = (
-                    float(node.find("micronsPerXPixel").text)
+                    _get_element_as_float(node, "micronsPerXPixel")
                     * DEFAULT_OAV_WINDOW[0]
                     / xsize
                 )
                 self.micronsPerYPixel = (
-                    float(node.find("micronsPerYPixel").text)
+                    _get_element_as_float(node, "micronsPerYPixel")
                     * DEFAULT_OAV_WINDOW[1]
                     / ysize
                 )
@@ -155,7 +163,7 @@ class OAVConfigParams:
 
     def get_beam_position_from_zoom(
         self, zoom: float, xsize: int, ysize: int
-    ) -> Tuple[int, int]:
+    ) -> tuple[int, int]:
         """
         Extracts the beam location in pixels `xCentre` `yCentre`, for a requested zoom \
         level. The beam location is manually inputted by the beamline operator on GDA \
@@ -164,7 +172,7 @@ class OAVConfigParams:
         """
         crosshair_x_line = None
         crosshair_y_line = None
-        with open(self.display_config, "r") as f:
+        with open(self.display_config) as f:
             file_lines = f.readlines()
             for i in range(len(file_lines)):
                 if file_lines[i].startswith("zoomLevel = " + str(zoom)):
@@ -188,7 +196,7 @@ class OAVConfigParams:
 
     def calculate_beam_distance(
         self, horizontal_pixels: int, vertical_pixels: int
-    ) -> Tuple[int, int]:
+    ) -> tuple[int, int]:
         """
         Calculates the distance between the beam centre and the given (horizontal, vertical).
 

dodal/devices/oav/oav_to_redis_forwarder.py

@@ -0,0 +1,100 @@
+import asyncio
+import io
+import pickle
+import uuid
+
+import numpy as np
+from aiohttp import ClientResponse, ClientSession
+from bluesky.protocols import Flyable
+from ophyd_async.core import AsyncStatus, StandardReadable
+from ophyd_async.core.signal import soft_signal_r_and_setter
+from ophyd_async.epics.signal import epics_signal_r
+from PIL import Image
+from redis.asyncio import StrictRedis
+
+from dodal.log import LOGGER
+
+
+async def get_next_jpeg(response: ClientResponse) -> bytes:
+    JPEG_START_BYTE = b"\xff\xd8"
+    JPEG_STOP_BYTE = b"\xff\xd9"
+    while True:
+        line = await response.content.readline()
+        if line.startswith(JPEG_START_BYTE):
+            return line + await response.content.readuntil(JPEG_STOP_BYTE)
+
+
+class OAVToRedisForwarder(StandardReadable, Flyable):
+    """Forwards OAV image data to redis. To use call:
+
+    > bps.kickoff(oav_forwarder)
+    > bps.monitor(oav_forwarder.uuid)
+    > bps.complete(oav_forwarder)
+
+    """
+
+    def __init__(
+        self,
+        prefix: str,
+        redis_host: str,
+        redis_password: str,
+        redis_db: int = 0,
+        name: str = "",
+        redis_key: str = "test-image",
+    ) -> None:
+        """Reads image data from the MJPEG stream on an OAV and forwards it into a
+        redis database. This is currently only used for murko integration.
+
+        Arguments:
+            prefix: str             the PV prefix of the OAV
+            redis_host: str         the host where the redis database is running
+            redis_password: str     the password for the redis database
+            redis_db: int           which redis database to connect to, defaults to 0
+            name: str               the name of this device
+            redis_key: str          the key to store data in, defaults to "test-image"
+        """
+        self.stream_url = epics_signal_r(str, f"{prefix}-DI-OAV-01:MJPG:HOST_RBV")
+
+        with self.add_children_as_readables():
+            self.uuid, self.uuid_setter = soft_signal_r_and_setter(str)
+
+        self.forwarding_task = None
+        self.redis_client = StrictRedis(
+            host=redis_host, password=redis_password, db=redis_db
+        )
+
+        self.redis_key = redis_key
+
+        # The uuid that images are being saved under, this should be monitored for
+        # callbacks to correlate the data
+        self.uuid, self.uuid_setter = soft_signal_r_and_setter(str)
+
+        super().__init__(name=name)
+
+    async def _get_frame_and_put_to_redis(self, response: ClientResponse):
+        """Converts the data that comes in as a jpeg byte stream into a numpy array of
+        RGB values, pickles this array then writes it to redis.
+        """
+        jpeg_bytes = await get_next_jpeg(response)
+        self.uuid_setter(image_uuid := str(uuid.uuid4()))
+        img = Image.open(io.BytesIO(jpeg_bytes))
+        image_data = pickle.dumps(np.asarray(img))
+        await self.redis_client.hset(self.redis_key, image_uuid, image_data)  # type: ignore
+        LOGGER.debug(f"Sent frame to redis key {self.redis_key} with uuid {image_uuid}")
+
+    async def _stream_to_redis(self):
+        stream_url = await self.stream_url.get_value()
+        async with ClientSession() as session:
+            async with session.get(stream_url) as response:
+                while True:
+                    await self._get_frame_and_put_to_redis(response)
+                    await asyncio.sleep(0.01)
+
+    @AsyncStatus.wrap
+    async def kickoff(self):
+        self.forwarding_task = asyncio.create_task(self._stream_to_redis())
+
+    @AsyncStatus.wrap
+    async def complete(self):
+        assert self.forwarding_task, "Device not kicked off"
+        self.forwarding_task.cancel()

dodal/devices/oav/pin_image_recognition/__init__.py

@@ -5,6 +5,7 @@ import numpy as np
 from numpy.typing import NDArray
 from ophyd_async.core import (
     AsyncStatus,
+    HintedSignal,
     StandardReadable,
     observe_value,
     soft_signal_r_and_setter,
@@ -77,12 +78,13 @@ class PinTipDetection(StandardReadable):
         self.min_tip_height = soft_signal_rw(int, 5, name="min_tip_height")
         self.validity_timeout = soft_signal_rw(float, 5.0, name="validity_timeout")
 
-        self.set_readable_signals(
-            read=[
+        self.add_readables(
+            [
                 self.triggered_tip,
                 self.triggered_top_edge,
                 self.triggered_bottom_edge,
             ],
+            wrapper=HintedSignal,
         )
 
         super().__init__(name=name)
@@ -134,9 +136,7 @@ class PinTipDetection(StandardReadable):
         location = sample_detection.processArray(array_data)
         end_time = time.time()
         LOGGER.debug(
-            "Sample location detection took {}ms".format(
-                (end_time - start_time) * 1000.0
-            )
+            f"Sample location detection took {(end_time - start_time) * 1000.0}ms"
         )
         return location
 
@@ -154,7 +154,7 @@ class PinTipDetection(StandardReadable):
                     location = await self._get_tip_and_edge_data(value)
                     self._set_triggered_values(location)
                 except Exception as e:
-                    LOGGER.warn(
+                    LOGGER.warning(
                         f"Failed to detect pin-tip location, will retry with next image: {e}"
                     )
                 else:

dodal/devices/oav/pin_image_recognition/utils.py

@@ -1,6 +1,7 @@
+from collections.abc import Callable
 from dataclasses import dataclass
 from enum import Enum
-from typing import Callable, Final, Tuple
+from typing import Final
 
 import cv2
 import numpy as np
@@ -103,7 +104,7 @@ class SampleLocation:
     edge_bottom: np.ndarray
 
 
-class MxSampleDetect(object):
+class MxSampleDetect:
     def __init__(
         self,
         *,
@@ -161,7 +162,7 @@ class MxSampleDetect(object):
     @staticmethod
     def _first_and_last_nonzero_by_columns(
         arr: np.ndarray,
-    ) -> Tuple[np.ndarray, np.ndarray]:
+    ) -> tuple[np.ndarray, np.ndarray]:
         """
         Finds the indexes of the first & last non-zero values by column in a 2d array.
 
@@ -243,9 +244,7 @@ class MxSampleDetect(object):
             bottom[x + 1 :] = NONE_VALUE
 
         LOGGER.info(
-            "pin-tip detection: Successfully located pin tip at (x={}, y={})".format(
-                tip_x, tip_y
-            )
+            f"pin-tip detection: Successfully located pin tip at (x={tip_x}, y={tip_y})"
         )
         return SampleLocation(
             tip_x=tip_x, tip_y=tip_y, edge_bottom=bottom, edge_top=top

dodal/devices/oav/utils.py

@@ -1,5 +1,5 @@
+from collections.abc import Generator
 from enum import IntEnum
-from typing import Generator, Tuple
 
 import bluesky.plan_stubs as bps
 import numpy as np
@@ -10,7 +10,7 @@ from dodal.devices.oav.oav_detector import OAVConfigParams
 from dodal.devices.oav.pin_image_recognition import PinTipDetection
 from dodal.devices.smargon import Smargon
 
-Pixel = Tuple[int, int]
+Pixel = tuple[int, int]
 
 
 class PinNotFoundException(Exception):

dodal/devices/p99/sample_stage.py

@@ -0,0 +1,43 @@
+from enum import Enum
+
+from ophyd_async.core import Device
+from ophyd_async.epics.signal import epics_signal_rw
+
+
+class SampleAngleStage(Device):
+    def __init__(self, prefix: str, name: str):
+        self.theta = epics_signal_rw(
+            float, prefix + "WRITETHETA:RBV", prefix + "WRITETHETA"
+        )
+        self.roll = epics_signal_rw(
+            float, prefix + "WRITEROLL:RBV", prefix + "WRITEROLL"
+        )
+        self.pitch = epics_signal_rw(
+            float, prefix + "WRITEPITCH:RBV", prefix + "WRITEPITCH"
+        )
+        super().__init__(name=name)
+
+
+class p99StageSelections(str, Enum):
+    Empty = "Empty"
+    Mn5um = "Mn 5um"
+    Fe = "Fe (empty)"
+    Co5um = "Co 5um"
+    Ni5um = "Ni 5um"
+    Cu5um = "Cu 5um"
+    Zn5um = "Zn 5um"
+    Zr = "Zr (empty)"
+    Mo = "Mo (empty)"
+    Rh = "Rh (empty)"
+    Pd = "Pd (empty)"
+    Ag = "Ag (empty)"
+    Cd25um = "Cd 25um"
+    W = "W (empty)"
+    Pt = "Pt (empty)"
+    User = "User"
+
+
+class FilterMotor(Device):
+    def __init__(self, prefix: str, name: str):
+        self.user_setpoint = epics_signal_rw(p99StageSelections, prefix)
+        super().__init__(name=name)

dodal/devices/robot.py

@@ -1,5 +1,5 @@
 import asyncio
-from asyncio import FIRST_COMPLETED, Task
+from asyncio import FIRST_COMPLETED, CancelledError, Task
 from dataclasses import dataclass
 from enum import Enum
 
@@ -45,6 +45,9 @@ class BartRobot(StandardReadable, Movable):
     LOAD_TIMEOUT = 60
     NO_PIN_ERROR_CODE = 25
 
+    # How far the gonio position can be out before loading will fail
+    LOAD_TOLERANCE_MM = 0.02
+
     def __init__(
         self,
         name: str,
@@ -74,19 +77,28 @@ class BartRobot(StandardReadable, Movable):
             await wait_for_value(self.error_code, self.NO_PIN_ERROR_CODE, None)
             raise RobotLoadFailed(self.NO_PIN_ERROR_CODE, "Pin was not detected")
 
-        finished, unfinished = await asyncio.wait(
-            [
-                Task(raise_if_no_pin()),
-                Task(
-                    wait_for_value(self.gonio_pin_sensor, PinMounted.PIN_MOUNTED, None)
-                ),
-            ],
-            return_when=FIRST_COMPLETED,
-        )
-        for task in unfinished:
-            task.cancel()
-        for task in finished:
-            await task
+        async def wfv():
+            await wait_for_value(self.gonio_pin_sensor, PinMounted.PIN_MOUNTED, None)
+
+        tasks = [
+            (Task(raise_if_no_pin())),
+            (Task(wfv())),
+        ]
+        try:
+            finished, unfinished = await asyncio.wait(
+                tasks,
+                return_when=FIRST_COMPLETED,
+            )
+            for task in unfinished:
+                task.cancel()
+            for task in finished:
+                await task
+        except CancelledError:
+            # If the outer enclosing task cancels after LOAD_TIMEOUT, this causes CancelledError to be raised
+            # in the current task, when it propagates to here we should cancel all pending tasks before bubbling up
+            for task in tasks:
+                task.cancel()
+            raise
 
     async def _load_pin_and_puck(self, sample_location: SampleLocation):
         LOGGER.info(f"Loading pin {sample_location}")
@@ -108,12 +120,12 @@ class BartRobot(StandardReadable, Movable):
         await self.pin_mounted_or_no_pin_found()
 
     @AsyncStatus.wrap
-    async def set(self, sample_location: SampleLocation):
+    async def set(self, value: SampleLocation):
         try:
             await asyncio.wait_for(
-                self._load_pin_and_puck(sample_location), timeout=self.LOAD_TIMEOUT
+                self._load_pin_and_puck(value), timeout=self.LOAD_TIMEOUT
             )
-        except asyncio.TimeoutError:
+        except asyncio.TimeoutError as e:
             error_code = await self.error_code.get_value()
             error_string = await self.error_str.get_value()
-            raise RobotLoadFailed(error_code, error_string)
+            raise RobotLoadFailed(int(error_code), error_string) from e

dodal/devices/scintillator.py

@@ -1,7 +1,10 @@
-from ophyd import Component as Cpt
-from ophyd import Device, EpicsMotor
+from ophyd_async.core import StandardReadable
+from ophyd_async.epics.motion import Motor
 
 
-class Scintillator(Device):
-    y = Cpt(EpicsMotor, "-MO-SCIN-01:Y")
-    z = Cpt(EpicsMotor, "-MO-SCIN-01:Z")
+class Scintillator(StandardReadable):
+    def __init__(self, prefix: str, name: str = ""):
+        with self.add_children_as_readables():
+            self.y = Motor(prefix + "-MO-SCIN-01:Y")
+            self.z = Motor(prefix + "-MO-SCIN-01:Z")
+        super().__init__(name)

dodal/devices/smargon.py

@@ -1,8 +1,8 @@
-from collections.abc import Generator
+from collections.abc import Collection, Generator
 from dataclasses import dataclass
 from enum import Enum
 from math import isclose
-from typing import Collection, cast
+from typing import cast
 
 from bluesky import plan_stubs as bps
 from bluesky.utils import Msg
@@ -87,7 +87,7 @@ class XYZLimits:
     def position_valid(self, pos: Collection[float]) -> bool:
         return all(
             axis_limits.contains(value)
-            for axis_limits, value in zip([self.x, self.y, self.z], pos)
+            for axis_limits, value in zip([self.x, self.y, self.z], pos, strict=False)
         )
 
 

dodal/devices/status.py

@@ -1,41 +1,12 @@
-from typing import Any, TypeVar
+from typing import Any
 
 from ophyd.status import SubscriptionStatus
 
-T = TypeVar("T")
-
 
 def await_value(
-    subscribable: Any, expected_value: T, timeout: None | int = None
+    subscribable: Any, expected_value: object, timeout: None | int = None
 ) -> SubscriptionStatus:
     def value_is(value, **_):
         return value == expected_value
 
     return SubscriptionStatus(subscribable, value_is, timeout=timeout)
-
-
-def await_value_in_list(
-    subscribable: Any, expected_value: list, timeout: None | int = None
-) -> SubscriptionStatus:
-    """Returns a status which is completed when the subscriptable contains a value
-    within the expected_value list"""
-
-    def value_is(value, **_):
-        return value in expected_value
-
-    if not isinstance(expected_value, list):
-        raise TypeError(f"expected value {expected_value} is not a list")
-    else:
-        return SubscriptionStatus(subscribable, value_is, timeout=timeout)
-
-
-def await_approx_value(
-    subscribable: Any,
-    expected_value: T,
-    deadband: float = 1e-09,
-    timeout: None | int = None,
-) -> SubscriptionStatus:
-    def value_is_approx(value, **_):
-        return abs(value - expected_value) <= deadband
-
-    return SubscriptionStatus(subscribable, value_is_approx, timeout=timeout)

dodal/devices/tetramm.py

@@ -1,6 +1,6 @@
 import asyncio
+from collections.abc import Sequence
 from enum import Enum
-from typing import Sequence
 
 from bluesky.protocols import Hints
 from ophyd_async.core import (
@@ -115,7 +115,7 @@ class TetrammController(DetectorControl):
     async def arm(
         self,
         num: int,
-        trigger: DetectorTrigger,
+        trigger: DetectorTrigger = DetectorTrigger.edge_trigger,
         exposure: float | None = None,
     ) -> AsyncStatus:
         if exposure is None:
@@ -132,7 +132,7 @@ class TetrammController(DetectorControl):
             self._drv.averaging_time.set(exposure), self.set_exposure(exposure)
         )
 
-        status = await set_and_wait_for_value(self._drv.acquire, 1)
+        status = await set_and_wait_for_value(self._drv.acquire, True)
 
         return status
 
@@ -150,7 +150,7 @@ class TetrammController(DetectorControl):
             )
 
     async def disarm(self):
-        await stop_busy_record(self._drv.acquire, 0, timeout=1)
+        await stop_busy_record(self._drv.acquire, False, timeout=1)
 
     async def set_exposure(self, exposure: float):
         """Tries to set the exposure time of a single frame.

dodal/devices/thawer.py

@@ -15,7 +15,7 @@ class ThawerStates(str, Enum):
     ON = "On"
 
 
-class ThawingTimer(Device):
+class ThawingTimer(Device, Stoppable):
     def __init__(self, control_signal: SignalRW[ThawerStates]) -> None:
         self._control_signal = control_signal
         self._thawing_task: Task | None = None
@@ -32,7 +32,8 @@ class ThawingTimer(Device):
         finally:
             await self._control_signal.set(ThawerStates.OFF)
 
-    async def stop(self):
+    @AsyncStatus.wrap
+    async def stop(self, *args, **kwargs):
         if self._thawing_task:
             self._thawing_task.cancel()
 
@@ -43,6 +44,7 @@ class Thawer(StandardReadable, Stoppable):
         self.thaw_for_time_s = ThawingTimer(self.control)
         super().__init__(name)
 
-    async def stop(self):
+    @AsyncStatus.wrap
+    async def stop(self, *args, **kwargs):
         await self.thaw_for_time_s.stop()
         await self.control.set(ThawerStates.OFF)

dodal/devices/undulator_dcm.py

@@ -74,14 +74,12 @@ class UndulatorDCM(StandardReadable, Movable):
             daq_configuration_path + "/domain/beamlineParameters"
         )["DCM_Perp_Offset_FIXED"]
 
-    def set(self, value: float) -> AsyncStatus:
-        async def _set():
-            await asyncio.gather(
-                self._set_dcm_energy(value),
-                self._set_undulator_gap_if_required(value),
-            )
-
-        return AsyncStatus(_set())
+    @AsyncStatus.wrap
+    async def set(self, value: float):
+        await asyncio.gather(
+            self._set_dcm_energy(value),
+            self._set_undulator_gap_if_required(value),
+        )
 
     async def _set_dcm_energy(self, energy_kev: float) -> None:
         access_level = await self.undulator.gap_access.get_value()

dodal/devices/util/adjuster_plans.py

@@ -3,10 +3,10 @@ All the methods in this module return a bluesky plan generator that adjusts a va
 according to some criteria either via feedback, preset positions, lookup tables etc.
 """
 
-from typing import Callable, Generator
+from collections.abc import Callable, Generator
 
 from bluesky import plan_stubs as bps
-from bluesky.run_engine import Msg
+from bluesky.utils import Msg
 from ophyd.epics_motor import EpicsMotor
 from ophyd_async.epics.motion import Motor