Merge branch 'main' into add_apple_knot_to_i05

Author: Villtord

src/dodal/beamlines/i05.py

@@ -1,19 +1,31 @@
 from dodal.beamlines.i05_shared import devices as i05_shared_devices
 from dodal.common.beamlines.beamline_utils import set_beamline as set_utils_beamline
 from dodal.device_manager import DeviceManager
+from dodal.devices.beamlines.i05 import I05Goniometer
 from dodal.devices.temperture_controller import Lakeshore336
 from dodal.log import set_beamline as set_log_beamline
 from dodal.utils import BeamlinePrefix, get_beamline_name
 
-devices = DeviceManager()
-devices.include(i05_shared_devices)
-
 BL = get_beamline_name("i05")
 PREFIX = BeamlinePrefix(BL)
 set_log_beamline(BL)
 set_utils_beamline(BL)
 
+devices = DeviceManager()
+devices.include(i05_shared_devices)
+
 
 @devices.factory()
 def sample_temperature_controller() -> Lakeshore336:
     return Lakeshore336(prefix=f"{PREFIX.beamline_prefix}-EA-TCTRL-02:")
+
+
+@devices.factory()
+def sa() -> I05Goniometer:
+    """Sample Manipulator."""
+    return I05Goniometer(
+        f"{PREFIX.beamline_prefix}-EA-SM-01:",
+        x_infix="SAX",
+        y_infix="SAY",
+        z_infix="SAZ",
+    )

src/dodal/beamlines/i05_1.py

@@ -1,6 +1,7 @@
 from dodal.beamlines.i05_shared import devices as i05_shared_devices
 from dodal.common.beamlines.beamline_utils import set_beamline as set_utils_beamline
 from dodal.device_manager import DeviceManager
+from dodal.devices.beamlines.i05_1 import XYZPolarAzimuthDefocusStage
 from dodal.log import set_beamline as set_log_beamline
 from dodal.utils import BeamlinePrefix, get_beamline_name
 
@@ -11,3 +12,9 @@
 
 devices = DeviceManager()
 devices.include(i05_shared_devices)
+
+
+@devices.factory
+def sm() -> XYZPolarAzimuthDefocusStage:
+    """Sample Manipulator."""
+    return XYZPolarAzimuthDefocusStage(prefix=f"{PREFIX.beamline_prefix}-EA-SM-01:")

src/dodal/common/maths.py

@@ -116,3 +116,31 @@ def contains(self, x: float, y: float) -> bool:
             self.get_min_x() <= x <= self.get_max_x()
             and self.get_min_y() <= y <= self.get_max_y()
         )
+
+
+"""
+| Rotation | Formula for X_rot | Formula for Y_rot |
+| -------- | ----------------- | ----------------- |
+| CW       | x cosθ + y sinθ   | -x sinθ + y cosθ  |
+| CCW      | x cosθ - y sinθ   | x sinθ + y cosθ   |
+"""
+
+
+def do_rotation(x: float, y: float, rotation_matrix: np.ndarray) -> tuple[float, float]:
+    positions = np.array([x, y])
+    rotation = rotation_matrix @ positions
+    return float(rotation[0]), float(rotation[1])
+
+
+def rotate_clockwise(theta: float, x: float, y: float) -> tuple[float, float]:
+    rotation_matrix = np.array(
+        [
+            [np.cos(theta), np.sin(theta)],
+            [-np.sin(theta), np.cos(theta)],
+        ]
+    )
+    return do_rotation(x, y, rotation_matrix)
+
+
+def rotate_counter_clockwise(theta: float, x: float, y: float) -> tuple[float, float]:
+    return rotate_clockwise(-theta, x, y)

src/dodal/devices/beamlines/i05/i05_motors.py

@@ -0,0 +1,65 @@
+from dodal.devices.beamlines.i05_shared.rotation_signals import (
+    create_rotational_ij_component_signals,
+)
+from dodal.devices.motors import (
+    _AZIMUTH,
+    _POLAR,
+    _TILT,
+    _X,
+    _Y,
+    _Z,
+    XYZPolarAzimuthTiltStage,
+)
+
+
+class I05Goniometer(XYZPolarAzimuthTiltStage):
+    """Six-physical-axis stage with a standard xyz translational stage and three axis of
+    rotation: polar, azimuth, and tilt.
+
+    In addition, it defines two virtual translational axes derived signals, `perp` and
+    `long`, which form a rotated Cartesian frame within the x-y plane.
+        - `long`: Translation along the rotated X-axis.
+        - `perp`: Translation along the rotated Y-axis.
+
+    The `perp` and `long` axes are virtual axes derived from the underlying x and y
+    motors using a fixed rotation angle (default 50 degrees). Rotation angle corresponds
+    to an angle between analyser axis and X-ray beam axis. From the user's point of
+    view, these virtual axes behave as ordinary orthogonal Cartesian translation axes
+    aligned with the incoming X-ray beam (long) and perpendicular to it (perp),
+    while internally coordinating motion of the x (perpendicular to analyser axis) and y
+    (along analyser axis) motors.
+    """
+
+    def __init__(
+        self,
+        prefix: str,
+        x_infix: str = _X,
+        y_infix: str = _Y,
+        z_infix: str = _Z,
+        polar_infix: str = _POLAR,
+        azimuth_infix: str = _AZIMUTH,
+        tilt_infix: str = _TILT,
+        rotation_angle_deg: float = 50.0,
+        name: str = "",
+    ):
+        self.rotation_angle_deg = rotation_angle_deg
+
+        super().__init__(
+            prefix,
+            name,
+            x_infix=x_infix,
+            y_infix=y_infix,
+            z_infix=z_infix,
+            polar_infix=polar_infix,
+            azimuth_infix=azimuth_infix,
+            tilt_infix=tilt_infix,
+        )
+
+        with self.add_children_as_readables():
+            self.perp, self.long = create_rotational_ij_component_signals(
+                i_read=self.x.user_readback,
+                i_write=self.x,
+                j_read=self.y.user_readback,
+                j_write=self.y,
+                angle_deg=self.rotation_angle_deg,
+            )

src/dodal/devices/beamlines/i05_1/__init__.py

@@ -0,0 +1,3 @@
+from .i05_1_motors import XYZPolarAzimuthDefocusStage
+
+__all__ = ["XYZPolarAzimuthDefocusStage"]

src/dodal/devices/beamlines/i05_1/i05_1_motors.py

@@ -0,0 +1,75 @@
+from ophyd_async.epics.motor import Motor
+
+from dodal.devices.beamlines.i05_shared.rotation_signals import (
+    create_rotational_ij_component_signals,
+)
+from dodal.devices.motors import XYZPolarAzimuthStage
+
+
+class XYZPolarAzimuthDefocusStage(XYZPolarAzimuthStage):
+    """Six-physical-axis stage with a standard xyz stage, 2 axis of rotation: polar,
+    azimuth and one extra tranlastional axis defocus.
+
+    This device exposes four virtual translational axes that are defined in frames
+    of reference attached to the sample:
+
+    - `hor` and `vert`:
+        Horizontal and vertical virtual translation axes of the rotated sample frame.
+        These axes are derived from X and Y axes rotated
+        with the azimuth angle, so that motion along `hor` and `vert`
+        remains aligned with the gravity direction regardless of its azimuthal rotation.
+
+    - `long` and `perp`:
+        Longitudinal and perpendicular virtual translation axes in the rotated sample
+        frame. These axes are derived from the Z-axis and the
+        virtual `hor` axis, and depend on the polar angle.
+        Motion along `long` aligned with the analyser axis,
+        while `perp` moves perpendicular to it within the polar rotation plane.
+
+    All four virtual axes behave as ordinary orthogonal Cartesian translations
+    from the user's point of view, while internally coordinating motion of the
+    underlying motors to account for the current rotation angles.
+
+    This allows users to position the sample in physically meaningful, sample-aligned
+    coordinates without needing to manually compensate for azimuth or polar rotations.
+    """
+
+    def __init__(
+        self,
+        prefix: str,
+        x_infix: str = "SMX",
+        y_infix: str = "SMY",
+        z_infix: str = "SMZ",
+        polar_infix: str = "POL",
+        azimuth_infix: str = "AZM",
+        defocus_infix: str = "SMDF",
+        name: str = "",
+    ):
+        super().__init__(
+            prefix,
+            name,
+            x_infix=x_infix,
+            y_infix=y_infix,
+            z_infix=z_infix,
+            polar_infix=polar_infix,
+            azimuth_infix=azimuth_infix,
+        )
+
+        with self.add_children_as_readables():
+            self.defocus = Motor(prefix + defocus_infix)
+            self.hor, self.vert = create_rotational_ij_component_signals(
+                i_read=self.x.user_readback,
+                j_read=self.y.user_readback,
+                i_write=self.x,
+                j_write=self.y,
+                angle_deg=self.azimuth.user_readback,
+                clockwise_frame=True,
+            )
+            self.long, self.perp = create_rotational_ij_component_signals(
+                i_read=self.z.user_readback,
+                i_write=self.z,
+                j_read=self.hor,
+                j_write=self.hor,
+                angle_deg=self.polar.user_readback,
+                clockwise_frame=False,
+            )

src/dodal/devices/beamlines/i05_shared/rotation_signals.py

@@ -0,0 +1,89 @@
+import asyncio
+from math import radians
+
+from bluesky.protocols import Movable
+from bluesky.utils import maybe_await
+from ophyd_async.core import (
+    SignalR,
+    SignalRW,
+    derived_signal_rw,
+)
+
+from dodal.common.maths import rotate_clockwise, rotate_counter_clockwise
+
+
+async def _get_angle_deg(angle_deg: SignalR[float] | float) -> float:
+    if isinstance(angle_deg, SignalR):
+        return await angle_deg.get_value()
+    return angle_deg
+
+
+def create_rotational_ij_component_signals(
+    i_read: SignalR[float],
+    j_read: SignalR[float],
+    i_write: Movable[float],
+    j_write: Movable[float],
+    angle_deg: float | SignalR[float],
+    clockwise_frame: bool = True,
+) -> tuple[SignalRW[float], SignalRW[float]]:
+    """Create virtual i/j axes representing a Cartesian coordinate frame
+    that is rotated by a given angle relative to the underlying equipment axes.
+
+    The returned signals expose the position of the system in a *rotated frame
+    of reference* (e.g. the sample or stage frame), while transparently mapping
+    reads and writes onto the real i/j signals in the fixed equipment (lab) frame.
+
+    From the user's point of view, i and j behave like ordinary orthogonal
+    Cartesian axes attached to the rotating object. Internally, all reads apply
+    a rotation to the real motor positions, and all writes apply the inverse
+    rotation so that the requested motion is achieved in the rotated frame.
+
+    Args:
+        i_read (SignalR): SignalR representing the i motor readback.
+        j_read (SignalR): representing the j motor readback.
+        i_write (Movable): object for setting the i position.
+        j_write (Movable): object for setting the j position.
+        angle_deg (float | SignalR): Rotation angle in degrees.
+        clockwise_frame (boolean): If True, the rotated frame is defined using a
+            clockwise rotation; otherwise, a counter-clockwise rotation is used.
+
+    Returns:
+        tuple[SignalRW[float], SignalRW[float]] Two virtual read/write signals
+        corresponding to the rotated i and j components.
+    """
+    rotate = rotate_clockwise if clockwise_frame else rotate_counter_clockwise
+    inverse_rotate = rotate_counter_clockwise if clockwise_frame else rotate_clockwise
+
+    async def _read_rotated() -> tuple[float, float, float]:
+        i, j, ang = await asyncio.gather(
+            i_read.get_value(),
+            j_read.get_value(),
+            _get_angle_deg(angle_deg),
+        )
+        return (*rotate(radians(ang), i, j), ang)
+
+    async def _write_rotated(i_rot: float, j_rot: float, ang: float) -> None:
+        i_new, j_new = inverse_rotate(radians(ang), i_rot, j_rot)
+        await asyncio.gather(
+            maybe_await(i_write.set(i_new)),
+            maybe_await(j_write.set(j_new)),
+        )
+
+    def _read_i(i: float, j: float, ang: float) -> float:
+        return rotate(radians(ang), i, j)[0]
+
+    async def _set_i(value: float) -> None:
+        i_rot, j_rot, ang = await _read_rotated()
+        await _write_rotated(value, j_rot, ang)
+
+    def _read_j(i: float, j: float, ang: float) -> float:
+        return rotate(radians(ang), i, j)[1]
+
+    async def _set_j(value: float) -> None:
+        i_rot, j_rot, ang = await _read_rotated()
+        await _write_rotated(i_rot, value, ang)
+
+    return (
+        derived_signal_rw(_read_i, _set_i, i=i_read, j=j_read, ang=angle_deg),
+        derived_signal_rw(_read_j, _set_j, i=i_read, j=j_read, ang=angle_deg),
+    )

src/dodal/devices/motors.py

@@ -2,9 +2,11 @@
 import math
 from abc import ABC
 
-from ophyd_async.core import StandardReadable, derived_signal_rw
+from ophyd_async.core import SignalRW, StandardReadable, derived_signal_rw
 from ophyd_async.epics.motor import Motor
 
+from dodal.common.maths import rotate_clockwise, rotate_counter_clockwise
+
 _X, _Y, _Z = "X", "Y", "Z"
 
 _OMEGA = "OMEGA"
@@ -315,7 +317,9 @@ def __init__(
         super().__init__(name)
 
 
-def create_axis_perp_to_rotation(motor_theta: Motor, motor_i: Motor, motor_j: Motor):
+def create_axis_perp_to_rotation(
+    motor_theta: Motor, motor_i: Motor, motor_j: Motor
+) -> SignalRW[float]:
     """Given a signal that controls a motor in a rotation axis and two other
     signals controlling motors on a pair of orthogonal axes, these axes being in the
     rotating frame of reference created by the first axis, create a derived signal
@@ -337,25 +341,24 @@ def create_axis_perp_to_rotation(motor_theta: Motor, motor_i: Motor, motor_j: Mo
             a move here is entirely parallel with the derived axis.
     """
 
-    def _get(j_val: float, i_val: float, rot_value: float) -> float:
-        i_component = i_val * math.cos(math.radians(rot_value))
-        j_component = j_val * math.sin(math.radians(rot_value))
-        return i_component + j_component
+    def _get(j_val: float, i_val: float, rot_deg_value: float) -> float:
+        x, y = rotate_clockwise(math.radians(rot_deg_value), i_val, j_val)
+        return x
 
     async def _set(vertical_value: float) -> None:
-        rot_value = await motor_theta.user_readback.get_value()
-        i_component = vertical_value * math.cos(math.radians(rot_value))
-        j_component = vertical_value * math.sin(math.radians(rot_value))
+        rot_deg_value = await motor_theta.user_readback.get_value()
+        theta = math.radians(rot_deg_value)
+        i_component, j_component = rotate_counter_clockwise(theta, vertical_value, 0.0)
         await asyncio.gather(
             motor_i.set(i_component),
             motor_j.set(j_component),
-            motor_theta.set(rot_value),
+            motor_theta.set(rot_deg_value),
         )
 
     return derived_signal_rw(
         _get,
         _set,
         i_val=motor_i,
         j_val=motor_j,
-        rot_value=motor_theta,
+        rot_deg_value=motor_theta,
     )