Add i05 sample stage

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,45 @@ 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]:
+    rotation_matrix = np.array(
+        [
+            [np.cos(theta), -np.sin(theta)],
+            [np.sin(theta), np.cos(theta)],
+        ]
+    )
+    return do_rotation(x, y, rotation_matrix)

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

@@ -1,3 +1,4 @@
-from dodal.devices.beamlines.i05.enums import Grating
+from .enums import Grating
+from .i05_motors import I05Goniometer
 
-__all__ = ["Grating"]
+__all__ = ["Grating", "I05Goniometer"]

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

@@ -0,0 +1,69 @@
+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-axis stage with a standard xyz stage and three axis of rotation: polar,
+    azimuth, and tilt.
+
+    In addition, it defines two virtual translational axes, `perp` and `long`, which
+    form a rotated Cartesian frame within the x-y plane.
+
+    - `long`: Translation along the longitudinal direction of the rotated in-plane
+        coordinate frame defined by ``rotation_angle_deg``.
+
+    - `perp`: Translation perpendicular to `long` within the x-y plane.
+
+    The `perp` and `long` axes are derived from the underlying x and y motors using a
+    fixed rotation angle (default 50 degrees). From the user's point of view, these
+    behave as ordinary orthogonal Cartesian translation axes aligned with physically
+    meaningful directions on the sample, while internally coordinating motion of the x
+    and y motors.
+
+    Unlike sample-frame axes that rotate with a live rotation motor, these axes are
+    defined at a constant orientation set by `rotation_angle_deg`.
+    """
+
+    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-axis stage with a standard xyz stage and three axis of rotation: polar,
+    azimuth, and 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 translation axes in the sample frame.
+        These axes are derived from the lab-frame x and y motors and rotate
+        with the azimuth angle, so that motion along `hor` and `vert`
+        remains aligned with the sample regardless of its azimuthal rotation.
+
+    - `long` and `perp`:
+        Longitudinal and perpendicular translation axes in the tilted sample
+        frame. These axes are derived from the lab-frame z motor and the
+        sample-frame `hor` axis, and rotate with the polar angle.
+        Motion along `long` follows the sample's longitudinal direction,
+        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="SMX",
+        y_infix="SMY",
+        z_infix="SMZ",
+        polar_infix="POL",
+        azimuth_infix="AZM",
+        defocus_infix="SMDF",
+        name="",
+    ):
+        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 signals 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),
+    )