Support for mod-360 omega

src/dodal/common/maths.py

@@ -1,4 +1,11 @@
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import Self
+
 import numpy as np
+from ophyd_async.core import (
+    Array1D,
+)
 
 
 def step_to_num(start: float, stop: float, step: float) -> tuple[float, float, int]:
@@ -146,3 +153,100 @@ def rotate_clockwise(theta: float, x: float, y: float) -> tuple[float, float]:
 
 def rotate_counter_clockwise(theta: float, x: float, y: float) -> tuple[float, float]:
     return rotate_clockwise(-theta, x, y)
+
+
+MotorOffsetAndPhase = Array1D[np.float32]
+
+
+@dataclass
+class AngleWithPhase:
+    """Represents a point in an absolute rotational space which is defined by a phase where 0<=phase<360
+     and an offset from an origin where the absolute coordinate is the sum of the phase and the offset.
+
+    Attributes:
+        offset: The offset of 0 phase from some other unwrapped rotational coordinate space
+        phase: The phase in degrees relative to this offset.
+    """
+
+    offset: float
+    phase: float = 0.0
+
+    def __post_init__(self) -> None:
+        correction = 360 * (self.phase // 360)
+        self.offset += correction
+        self.phase -= correction
+
+    @classmethod
+    def from_iterable(cls, values: Iterable[float]) -> Self:
+        """Construct a normalised representation of the offset and phase, such that
+        0 <= phase < 360.
+
+        Args:
+            values (Iterable[float]): the offset and phase as a list or other iterable
+        """
+        offset, phase = values
+        return cls(offset, phase)
+
+    @classmethod
+    def wrap(cls, unwrapped: float) -> "AngleWithPhase":
+        """Construct a representation such that offset = n * 360 and 0 <= phase < 360.
+
+        Args:
+             unwrapped (float): The unwrapped angle in degrees
+        """
+        offset = AngleWithPhase.offset_from_unwrapped(unwrapped)
+        return cls(offset, unwrapped - offset)
+
+    def rebase_to(self, other: Self) -> "AngleWithPhase":
+        """Return this angle with the offset adjusted such that the phases can be compared."""
+        correction = other.offset - self.offset
+        if correction % 360:
+            return AngleWithPhase.from_iterable(
+                [self.offset + correction, self.phase - correction]
+            )
+        else:
+            return self
+
+    def unwrap(self) -> float:
+        """Generate the unwrapped representation of this angle."""
+        return self.offset + self.phase
+
+    def phase_distance(self, phase: float) -> float:
+        """Determine the shortest distance between this angle and the specified phase.
+
+        Args:
+            phase (float): The phase angle to compare to
+        """
+        phase = phase % 360
+        max_theta = max(self.phase, phase)
+        min_theta = min(self.phase, phase)
+        return min(max_theta - min_theta, min_theta + 360 - max_theta)
+
+    @classmethod
+    def offset_from_unwrapped(cls, unwrapped_deg: float) -> float:
+        """Obtain the offset from the corresponding wrapped angle in degrees."""
+        return round(unwrapped_deg // 360) * 360
+
+    def nearest_with_phase(self, phase_deg: float) -> "AngleWithPhase":
+        """Return the nearest angle to this one with the specified phase."""
+        phase_deg = phase_deg % 360
+        if phase_deg > self.phase:
+            return (
+                AngleWithPhase(self.offset, phase_deg)
+                if phase_deg - self.phase <= 180
+                else AngleWithPhase(self.offset - 360, phase_deg)
+            )
+        else:
+            return (
+                AngleWithPhase(self.offset, phase_deg)
+                if self.phase - phase_deg <= 180
+                else AngleWithPhase(self.offset + 360, phase_deg)
+            )
+
+
+def reflect_phase(phase) -> float:
+    """Convert the phase angle as if the corresponding unwrapped angle were to
+    be reflected about the origin and then re-wrapped, this corresponds to
+    converting a clockwise angle to an anti-clockwise angle and vice-versa.
+    """
+    return (360 - phase) % 360

src/dodal/devices/aithre_lasershaping/goniometer.py

@@ -1,9 +1,9 @@
 from ophyd_async.epics.motor import Motor
 
-from dodal.devices.motors import XYZOmegaStage, create_axis_perp_to_rotation
+from dodal.devices.motors import XYZWrappedOmegaStage, create_axis_perp_to_rotation
 
 
-class Goniometer(XYZOmegaStage):
+class Goniometer(XYZWrappedOmegaStage):
     """The Aithre lab goniometer and the XYZ stage it sits on.
 
     `x`, `y` and `z` control the axes of the positioner at the base, while `sampy` and

src/dodal/devices/motors.py

@@ -14,6 +14,7 @@
 from ophyd_async.epics.motor import Motor
 
 from dodal.common.maths import rotate_clockwise, rotate_counter_clockwise
+from dodal.devices.wrapped_axis import WrappedAxis
 
 _X = "X"
 _Y = "Y"
@@ -114,9 +115,27 @@ def __init__(
     ) -> None:
         with self.add_children_as_readables():
             self.omega = Motor(prefix + omega_infix)
+
         super().__init__(prefix, name, x_infix, y_infix, z_infix)
 
 
+class XYZWrappedOmegaStage(XYZOmegaStage):
+    """Four-axis stage with x, y, z linear axes and an omega axis with unrestricted rotation."""
+
+    def __init__(
+        self,
+        prefix: str = "",
+        name: str = "",
+        x_infix: str = _X,
+        y_infix: str = _Y,
+        z_infix: str = _Z,
+        omega_infix: str = _OMEGA,
+    ):
+        super().__init__(prefix, name, x_infix, y_infix, z_infix, omega_infix)
+        with self.add_children_as_readables():
+            self.wrapped_omega = WrappedAxis(self.omega)
+
+
 class XYZAzimuthStage(XYZStage):
     """Four-axis stage with a standard xyz stage and one axis of rotation: azimuth."""
 

src/dodal/devices/smargon.py

@@ -14,7 +14,8 @@
 from ophyd_async.epics.core import epics_signal_r, epics_signal_rw
 from ophyd_async.epics.motor import Motor
 
-from dodal.devices.motors import XYZOmegaStage
+from dodal.common.maths import AngleWithPhase
+from dodal.devices.motors import XYZWrappedOmegaStage
 from dodal.devices.util.epics_util import SetWhenEnabled
 
 
@@ -78,7 +79,7 @@ class CombinedMove(TypedDict, total=False):
     chi: float | None
 
 
-class Smargon(XYZOmegaStage, Movable):
+class Smargon(XYZWrappedOmegaStage, Movable):
     """Real motors added to allow stops following pin load (e.g. real_x1.stop() )
     X1 and X2 real motors provide compound chi motion as well as the compound X travel,
     increasing the gap between x1 and x2 changes chi, moving together changes virtual x.
@@ -104,6 +105,15 @@ def __init__(self, prefix: str, name: str = ""):
 
         super().__init__(prefix, name)
 
+    async def _get_target_value(self, motor_name: str, value: float) -> float:
+        if motor_name == "omega":
+            current_angle = AngleWithPhase.from_iterable(
+                await self.wrapped_omega.offset_and_phase.get_value()
+            )
+            return current_angle.nearest_with_phase(value).unwrap()
+        else:
+            return value
+
     @AsyncStatus.wrap
     async def set(self, value: CombinedMove):
         """This will move all motion together in a deferred move.
@@ -114,11 +124,16 @@ async def set(self, value: CombinedMove):
         only come back after the motion on that axis finished.
         """
         await self.defer_move.set(DeferMoves.ON)
+        # TODO Hotfix required here until https://github.com/DiamondLightSource/dodal/issues/1998
+        # is implemented in separate PR
         try:
             finished_moving = []
             for motor_name, new_setpoint in value.items():
                 if new_setpoint is not None and isinstance(new_setpoint, int | float):
-                    axis: Motor = getattr(self, motor_name)
+                    new_setpoint = await self._get_target_value(
+                        motor_name, new_setpoint
+                    )
+                    axis = getattr(self, motor_name)
                     await axis.check_motor_limit(
                         await axis.user_setpoint.get_value(), new_setpoint
                     )

src/dodal/devices/wrapped_axis.py

@@ -0,0 +1,73 @@
+import numpy as np
+from ophyd_async.core import (
+    Reference,
+    StandardReadable,
+    StandardReadableFormat,
+    derived_signal_rw,
+)
+from ophyd_async.epics.motor import Motor
+
+from dodal.common.maths import AngleWithPhase, MotorOffsetAndPhase
+
+
+class WrappedAxis(StandardReadable):
+    """This device is a wrapper around a rotational Motor that presents the unwrapped coordinate system of the
+    underlying motor as a combination of a phase angle and an offset from the motor's origin coordinate.
+
+    Attributes:
+          phase (float): This is a read-write signal that corresponds to the motor's phase angle, relative to the offset.
+              The behaviour of the phase signal differs from that of the underlying motor setpoint/readback signal
+              in the following ways:
+                 * Readback values are constrained to 0 <= omega < 360
+                 * Write values are normalised to 0 <= omega < 360, and then mapped to the nearest unwrapped angle. The
+                   underlying motor will be moved via the shortest path to the required phase angle, thus the direction
+                   of an unwrapped axis move is not always the same as the direction of a move in the wrapped axis.
+                 * Write values are normalised so for un-normalised writes, the readback will differ.
+                 * Bluesky mvr operations greater of 180 degrees or more will not result in the expected moves.
+                 * set_and_wait_for_value() is unreliable close to the wrap-around (however in the common case this is
+                   where deadband is 0.001 and values are rounded to this level, the default equality comparison is sufficient).
+                 * Sequences of moves on the unwrapped axis that would not result in cumulative motion axis will
+                   result in a cumulative motion in the wrapped axis. e.g. 0 -> 120 -> 240 -> 0 is 0 degrees of real
+                   cumulative motion when performed in the unwrapped case, but is 360 degrees of real motion when performed
+                   on the wrapped axis.
+                 * The reverse is also true 0 -> 240 -> 360 ->  is 0 degrees of real motion when executed on the wrapped axis
+                   but 360 degrees when performed in the unwrapped axis.
+                 * The above means that use of phase to set motor position is unsuitable for axes
+                   where the underlying motor rotation is constrained.
+          offset_and_phase (Array1D[np.float32]): retrieve the offset and phase together, for use when
+              mapping to the underlying unwrapped axis. These values can be converted and manipulated using the
+              AngleWithPhase helper class.
+    """
+
+    def __init__(self, real_motor: Motor, name=""):
+        self._real_motor = Reference(real_motor)
+        with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL):
+            self.offset_and_phase = derived_signal_rw(
+                self._get_motor_offset_and_phase,
+                self._set_motor_offset_and_phase,
+                motor_pos=real_motor,
+            )
+        with self.add_children_as_readables():
+            self.phase = derived_signal_rw(
+                self._get_phase, self._set_phase, offset_and_phase=self.offset_and_phase
+            )
+        super().__init__(name=name)
+
+    def _get_motor_offset_and_phase(self, motor_pos: float) -> MotorOffsetAndPhase:
+        angle = AngleWithPhase.wrap(motor_pos)
+        return np.array([angle.offset, angle.phase])
+
+    async def _set_motor_offset_and_phase(self, value: MotorOffsetAndPhase):
+        await self._real_motor().set(AngleWithPhase.from_iterable(value).unwrap())
+
+    def _get_phase(self, offset_and_phase: MotorOffsetAndPhase) -> float:
+        return offset_and_phase[1].item()
+
+    async def _set_phase(self, value: float):
+        """Set the motor phase to the specified phase value in degrees.
+        The motor will travel via the shortest distance path.
+        """
+        offset_and_phase = await self.offset_and_phase.get_value()
+        current_position = AngleWithPhase.from_iterable(offset_and_phase)
+        target_value = current_position.nearest_with_phase(value).unwrap()
+        await self._real_motor().set(target_value)