Simplify assert rotated function, add doc string, improve type checking

Author: oliwenmandiamond

tests/devices/i05/test_i05_motors.py

@@ -7,8 +7,8 @@
 
 from dodal.devices.i05 import I05Goniometer
 from tests.devices.i05_shared.rotation_signal_test_util import (
-    AxesTestConfig,
-    assert_set_axis_preserves_other,
+    RotatedCartesianFrameTestConfig,
+    assert_rotated_axes_are_orthogonal,
 )
 
 
@@ -21,7 +21,7 @@ def expected_long_read(x: float, y: float, theta: float) -> float:
 
 
 @pytest.fixture
-async def goniometer():
+async def goniometer() -> I05Goniometer:
     with init_devices(mock=True):
         goniometer = I05Goniometer("TEST:")
     return goniometer
@@ -50,26 +50,23 @@ async def test_goniometer_read(goniometer: I05Goniometer) -> None:
     )
 
 
-async def test_goniometer_perp_and_long_set(
-    goniometer: I05Goniometer,
-) -> None:
-    axes_test_config = AxesTestConfig(
+async def test_goniometer_perp_and_long_set(goniometer: I05Goniometer) -> None:
+    frame_config = RotatedCartesianFrameTestConfig(
         i_read=goniometer.x.user_readback,
         j_read=goniometer.y.user_readback,
-        i_write=goniometer.x,  # type: ignore
-        j_write=goniometer.y,  # type: ignore
+        i_write=goniometer.x,
+        j_write=goniometer.y,
         angle_deg=goniometer.rotation_angle_deg,
         expected_i_read_func=expected_perp_read,
         expected_j_read_func=expected_long_read,
         i_rotation_axis=goniometer.perp,
         j_rotation_axis=goniometer.long,
     )
-
-    await assert_set_axis_preserves_other(
+    await assert_rotated_axes_are_orthogonal(
         i_val=10,
         j_val=5,
         angle_deg_val=goniometer.rotation_angle_deg,
         new_i_axis_value=20,
         new_j_axis_value=20,
-        axes_config=axes_test_config,
+        frame_config=frame_config,
     )

tests/devices/i05_1/test_i05_1_motors.py

@@ -7,8 +7,8 @@
 
 from dodal.devices.i05_1 import XYZPolarAzimuthDefocusStage
 from tests.devices.i05_shared.rotation_signal_test_util import (
-    AxesTestConfig,
-    assert_set_axis_preserves_other,
+    RotatedCartesianFrameTestConfig,
+    assert_rotated_axes_are_orthogonal,
 )
 
 
@@ -55,7 +55,6 @@ async def test_xyzpad_stage_read(xyzpad_stage: XYZPolarAzimuthDefocusStage) -> N
         xyzpad_stage.azimuth.set(azimuth_angle_deg),
         xyzpad_stage.polar.set(polar_angle_deg),
     )
-
     await assert_reading(
         xyzpad_stage,
         {
@@ -76,7 +75,7 @@ async def test_xyzpad_stage_read(xyzpad_stage: XYZPolarAzimuthDefocusStage) -> N
 async def test_xyzpad_hor_and_vert_set(
     xyzpad_stage: XYZPolarAzimuthDefocusStage,
 ) -> None:
-    axes_test_config = AxesTestConfig(
+    frame_config = RotatedCartesianFrameTestConfig(
         i_read=xyzpad_stage.x.user_readback,
         j_read=xyzpad_stage.y.user_readback,
         i_write=xyzpad_stage.x,
@@ -87,21 +86,20 @@ async def test_xyzpad_hor_and_vert_set(
         i_rotation_axis=xyzpad_stage.hor,
         j_rotation_axis=xyzpad_stage.vert,
     )
-
-    await assert_set_axis_preserves_other(
+    await assert_rotated_axes_are_orthogonal(
         i_val=10,
         j_val=5,
         angle_deg_val=45,
         new_i_axis_value=20,
         new_j_axis_value=20,
-        axes_config=axes_test_config,
+        frame_config=frame_config,
     )
 
 
 async def test_xyzpad_long_and_perp_set(
     xyzpad_stage: XYZPolarAzimuthDefocusStage,
 ) -> None:
-    axes_test_config = AxesTestConfig(
+    frame_config = RotatedCartesianFrameTestConfig(
         i_read=xyzpad_stage.z.user_readback,
         j_read=xyzpad_stage.hor,
         i_write=xyzpad_stage.z,
@@ -112,12 +110,11 @@ async def test_xyzpad_long_and_perp_set(
         i_rotation_axis=xyzpad_stage.long,
         j_rotation_axis=xyzpad_stage.perp,
     )
-
-    await assert_set_axis_preserves_other(
+    await assert_rotated_axes_are_orthogonal(
         i_val=10,
         j_val=5,
         angle_deg_val=45,
         new_i_axis_value=20,
         new_j_axis_value=20,
-        axes_config=axes_test_config,
+        frame_config=frame_config,
     )

tests/devices/i05_shared/rotation_signal_test_util.py

@@ -10,7 +10,7 @@
 
 
 @dataclass
-class AxesTestConfig:
+class RotatedCartesianFrameTestConfig:
     i_read: SignalR[float]
     j_read: SignalR[float]
     i_write: Movable[float]
@@ -22,68 +22,118 @@ class AxesTestConfig:
     j_rotation_axis: SignalRW[float]
 
 
-async def assert_set_axis_preserves_other(
-    *,
+async def _set_initial_state(
     i_val: float,
     j_val: float,
     angle_deg_val: float,
-    new_i_axis_value: float,
-    new_j_axis_value: float,
-    axes_config: AxesTestConfig,
+    frame_config: RotatedCartesianFrameTestConfig,
 ) -> None:
     await asyncio.gather(
-        maybe_await(axes_config.i_write.set(i_val)),
-        maybe_await(axes_config.j_write.set(j_val)),
+        maybe_await(frame_config.i_write.set(i_val)),
+        maybe_await(frame_config.j_write.set(j_val)),
     )
-    angle_deg = axes_config.angle_deg
-    if isinstance(angle_deg, Movable):
-        await maybe_await(angle_deg.set(angle_deg_val))
-
-    theta = radians(angle_deg_val)
+    if isinstance(frame_config.angle_deg, Movable):
+        await maybe_await(frame_config.angle_deg.set(angle_deg_val))
 
-    # Capture invariant component before move
-    expected_i = new_i_axis_value
-    expected_j = axes_config.expected_j_read_func(i_val, j_val, theta)
 
-    # Move axis to new position
-    await axes_config.i_rotation_axis.set(expected_i)
-
-    # Read back motors
+async def _read_rotated_components(
+    frame_config: RotatedCartesianFrameTestConfig, theta: float
+) -> tuple[float, float]:
     i_new, j_new = await asyncio.gather(
-        axes_config.i_read.get_value(), axes_config.j_read.get_value()
+        frame_config.i_read.get_value(),
+        frame_config.j_read.get_value(),
     )
-
-    # Compute rotated components after move
-    i_val_after = axes_config.expected_i_read_func(i_new, j_new, theta)
-    j_val_after = axes_config.expected_j_read_func(i_new, j_new, theta)
-
-    assert i_val_after == pytest.approx(expected_i)
-    assert j_val_after == pytest.approx(expected_j)
-
-    await asyncio.gather(
-        maybe_await(axes_config.i_write.set(i_val)),
-        maybe_await(axes_config.j_write.set(j_val)),
+    return (
+        frame_config.expected_i_read_func(i_new, j_new, theta),
+        frame_config.expected_j_read_func(i_new, j_new, theta),
     )
-    angle_deg = axes_config.angle_deg
-    if isinstance(angle_deg, Movable):
-        await maybe_await(angle_deg.set(angle_deg_val))
-
-    theta = radians(angle_deg_val)
-
-    expected_j = new_j_axis_value
-    expected_i = axes_config.expected_i_read_func(i_val, j_val, theta)
 
-    # Move axis to new position
-    await axes_config.j_rotation_axis.set(expected_j)
 
-    # Read back motors
-    i_new, j_new = await asyncio.gather(
-        axes_config.i_read.get_value(), axes_config.j_read.get_value()
+async def _assert_setting_axis_preserves_other(
+    set_axis: SignalRW[float],
+    set_value: float,
+    expected_i: float,
+    expected_j: float,
+    frame_config: RotatedCartesianFrameTestConfig,
+    theta: float,
+) -> None:
+    await set_axis.set(set_value)
+    i_after, j_after = await _read_rotated_components(
+        frame_config=frame_config, theta=theta
     )
+    assert i_after == pytest.approx(expected_i)
+    assert j_after == pytest.approx(expected_j)
 
-    # Compute rotated components after move
-    i_val_after = axes_config.expected_i_read_func(i_new, j_new, theta)
-    j_val_after = axes_config.expected_j_read_func(i_new, j_new, theta)
 
-    assert i_val_after == pytest.approx(expected_i)
-    assert j_val_after == pytest.approx(expected_j)
+async def assert_rotated_axes_are_orthogonal(
+    i_val: float,
+    j_val: float,
+    angle_deg_val: float,
+    new_i_axis_value: float,
+    new_j_axis_value: float,
+    frame_config: RotatedCartesianFrameTestConfig,
+) -> None:
+    """Assert that the virtual rotated axes behave as independent Cartesian
+    coordinates.
+
+    This helper verifies that setting one virtual axis in a rotated frame
+    updates only that axis while preserving the orthogonal axis in the same
+    rotated frame.
+
+    Specifically, the test:
+    - Initializes the underlying i/j motors and rotation angle.
+    - Sets the virtual i' axis and asserts that:
+        * i' moves to the requested value
+        * j' remains unchanged
+    - Resets the system.
+    - Sets the virtual j' axis and asserts that:
+        * j' moves to the requested value
+        * i' remains unchanged
+
+    This confirms that the virtual axes form a proper orthogonal Cartesian
+    coordinate frame, and that inverse rotations are applied correctly when
+    writing to the underlying motors.
+
+    Args:
+        i_val: Initial value for the underlying i-axis.
+        j_val: Initial value for the underlying j-axis.
+        angle_deg_val: Rotation angle (in degrees) defining the rotated frame.
+        new_i_axis_value: Target value for the virtual i' axis.
+        new_j_axis_value: Target value for the virtual j' axis.
+        frame_config: Configuration describing how virtual rotated axes are derived from
+            the underlying motors, including expected readout functions.
+    """
+    # Test setting i′ preserves j′
+    await _set_initial_state(
+        i_val=i_val,
+        j_val=j_val,
+        angle_deg_val=angle_deg_val,
+        frame_config=frame_config,
+    )
+    theta = radians(angle_deg_val)
+    expected_j = frame_config.expected_j_read_func(i_val, j_val, theta)
+
+    await _assert_setting_axis_preserves_other(
+        set_axis=frame_config.i_rotation_axis,
+        set_value=new_i_axis_value,
+        expected_i=new_i_axis_value,
+        expected_j=expected_j,
+        frame_config=frame_config,
+        theta=theta,
+    )
+    # Test setting j′ preserves i′
+    await _set_initial_state(
+        i_val=i_val,
+        j_val=j_val,
+        angle_deg_val=angle_deg_val,
+        frame_config=frame_config,
+    )
+    expected_i = frame_config.expected_i_read_func(i_val, j_val, theta)
+    await _assert_setting_axis_preserves_other(
+        set_axis=frame_config.j_rotation_axis,
+        set_value=new_j_axis_value,
+        expected_i=expected_i,
+        expected_j=new_j_axis_value,
+        frame_config=frame_config,
+        theta=theta,
+    )