add PolynomCompoundMotors class (#1471)

* first try doesn't work

* same

* first working version

* refactor input parameter to dict of motors and coefficients

* rename few variables

* small docstring change

* add tests

* few tests added

* added more tests

* remove motor set motor speed from tests

* order of variables change

* updated docs

* update tests

* rename few variables inside class

* fix tests

* finish merging with main

* rework docstring

* replace simmotor with motor

* remove name parameter

* rework docstring

---------

Co-authored-by: eir17846 <victor.rogalev@diamond.ac.uk>

Author: Villtord

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

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

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

@@ -0,0 +1,75 @@
+import asyncio
+
+import numpy as np
+from bluesky.protocols import (
+    Locatable,
+    Location,
+    Stoppable,
+)
+from ophyd_async.core import (
+    Array1D,
+    AsyncStatus,
+    Reference,
+    StandardReadable,
+    StandardReadableFormat,
+)
+from ophyd_async.epics.motor import Motor
+
+
+class PolynomCompoundMotors(
+    StandardReadable,
+    Locatable[float],
+    Stoppable,
+):
+    """Compound motor controller that synchronizes the movement of multiple motors
+    based on polynomial relationships.
+
+    When the master motor is moved, all driven (slave) motors are asynchronously moved
+    to positions calculated using polynomial coefficients from master motor position.
+    The master motor is always included with a default polynomial coefficient
+    of [0.0, 1.0], representing a direct mapping. Driven motors' positions are
+    calculated using NumPy's polynomial evaluation.
+
+    Args:
+        master(Motor): master motor.
+        driven_dict (dict[Motor, Array1D[np.float64]]): dictionary that defines mapping
+            of each driven motor to its polynomial coefficients (NumPy array).
+        name (str,optional): name of the device. Defaults to an empty string.
+    """
+
+    def __init__(
+        self,
+        master: Motor,
+        driven_dict: dict[Motor, Array1D[np.float64]],
+        name: str = "",
+    ) -> None:
+        self.motor_coeff_dict: dict[Reference[Motor], Array1D[np.float64]] = {}
+
+        # master motor added with polynomial coeff (0,1)
+        self.master_ref = Reference(master)
+        self.motor_coeff_dict[self.master_ref] = np.array([0.0, 1.0])
+
+        # slave motors added with coefficients from input parameters
+        for slave in driven_dict.keys():
+            self.motor_coeff_dict[Reference(slave)] = driven_dict[slave]
+
+        self.add_readables(
+            [ref().user_readback for ref in self.motor_coeff_dict.keys()],
+            StandardReadableFormat.HINTED_SIGNAL,
+        )
+        super().__init__(name=name)
+
+    @AsyncStatus.wrap
+    async def set(self, new_position: float) -> None:
+        await asyncio.gather(
+            *[
+                ref().set(float(np.polynomial.polynomial.polyval(new_position, coeff)))
+                for ref, coeff in self.motor_coeff_dict.items()
+            ]
+        )
+
+    async def stop(self, success=False):
+        await asyncio.gather(*[ref().stop() for ref in self.motor_coeff_dict.keys()])
+
+    async def locate(self) -> Location[float]:
+        return await self.master_ref().locate()

tests/devices/beamlines/i05/__init__.py

@@ -0,0 +1,87 @@
+import numpy as np
+import pytest
+from bluesky import plan_stubs as bps
+from bluesky.run_engine import RunEngine
+from ophyd_async.core import init_devices
+from ophyd_async.epics.motor import Motor
+from ophyd_async.testing import assert_configuration, assert_reading, partial_reading
+
+from dodal.devices.beamlines.i05 import PolynomCompoundMotors
+
+
+@pytest.fixture
+async def x_motor() -> Motor:
+    async with init_devices(mock=True):
+        x_motor = Motor("")
+    return x_motor
+
+
+@pytest.fixture
+async def y_motor() -> Motor:
+    async with init_devices(mock=True):
+        y_motor = Motor("")
+    return y_motor
+
+
+@pytest.fixture
+async def z_motor() -> Motor:
+    async with init_devices(mock=True):
+        z_motor = Motor("")
+    return z_motor
+
+
+@pytest.fixture
+async def mock_compound(
+    x_motor: Motor,
+    y_motor: Motor,
+    z_motor: Motor,
+) -> PolynomCompoundMotors:
+    async with init_devices(mock=True):
+        mock_compound = PolynomCompoundMotors(
+            x_motor,
+            {
+                y_motor: np.array([1.0, 2.0], dtype=np.float64),
+                z_motor: np.array([-1.0, 0.5], dtype=np.float64),
+            },
+        )
+    return mock_compound
+
+
+async def test_config_includes(mock_compound: PolynomCompoundMotors):
+    await assert_configuration(
+        mock_compound,
+        {},
+    )
+
+
+async def test_read_includes(mock_compound: PolynomCompoundMotors):
+    await assert_reading(
+        mock_compound,
+        {
+            "x_motor": partial_reading(0.0),
+            "y_motor": partial_reading(0.0),
+            "z_motor": partial_reading(0.0),
+        },
+    )
+
+
+async def test_move(
+    mock_compound: PolynomCompoundMotors,
+    run_engine: RunEngine,
+):
+    new_position = 10.0
+    run_engine(bps.mv(mock_compound, new_position))
+    for motor, coeff in mock_compound.motor_coeff_dict.items():
+        expected_position = float(np.polynomial.polynomial.polyval(new_position, coeff))
+        actual_position = await motor().user_readback.get_value()
+        assert actual_position == expected_position
+
+
+async def test_move_and_locate(
+    mock_compound: PolynomCompoundMotors,
+    run_engine: RunEngine,
+):
+    new_position = 1.23
+    run_engine(bps.mv(mock_compound, new_position))
+    located_position = await mock_compound.locate()
+    assert located_position == {"setpoint": new_position, "readback": new_position}