Add B07 and B07-1 sample stage (#1913)

* Add B07 sample stage

* Add doc strings

* Add b07_motors_test

* Remove mock=True

* Reduce number of parametrize

* Corrected test name

* Applied feedback

Author: oliwenmandiamond

src/dodal/beamlines/b07.py

@@ -1,9 +1,15 @@
 from dodal.beamlines.b07_shared import devices as b07_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.b07 import Grating, LensMode, PsuMode
+from dodal.devices.beamlines.b07 import (
+    B07SampleManipulator52B,
+    Grating,
+    LensMode,
+    PsuMode,
+)
 from dodal.devices.electron_analyser.base import EnergySource
 from dodal.devices.electron_analyser.specs import SpecsDetector
+from dodal.devices.motors import XYZPolarStage
 from dodal.devices.pgm import PlaneGratingMonochromator
 from dodal.log import set_beamline as set_log_beamline
 from dodal.utils import BeamlinePrefix, get_beamline_name
@@ -40,3 +46,16 @@ def analyser(energy_source: EnergySource) -> SpecsDetector[LensMode, PsuMode]:
         psu_mode_type=PsuMode,
         energy_source=energy_source,
     )
+
+
+@devices.factory()
+def sm52b() -> B07SampleManipulator52B:
+    return B07SampleManipulator52B(prefix=f"{B_PREFIX.beamline_prefix}-EA-SM-52:")
+
+
+@devices.factory()
+def sm21b() -> XYZPolarStage:
+    """Sample manipulator. NOTE: The polar attribute is equivalent to GDA roty."""
+    return XYZPolarStage(
+        prefix=f"{B_PREFIX.beamline_prefix}-EA-SM-21:", polar_infix="ROTY"
+    )

src/dodal/beamlines/b07_1.py

@@ -9,6 +9,7 @@
 )
 from dodal.devices.electron_analyser.base import EnergySource
 from dodal.devices.electron_analyser.specs import SpecsDetector
+from dodal.devices.motors import XYZPolarAzimuthStage
 from dodal.devices.pgm import PlaneGratingMonochromator
 from dodal.log import set_beamline as set_log_beamline
 from dodal.utils import BeamlinePrefix, get_beamline_name
@@ -50,3 +51,16 @@ def analyser(energy_source: EnergySource) -> SpecsDetector[LensMode, PsuMode]:
         psu_mode_type=PsuMode,
         energy_source=energy_source,
     )
+
+
+@devices.factory()
+def sm() -> XYZPolarAzimuthStage:
+    """Sample manipulator."""
+    return XYZPolarAzimuthStage(
+        f"{C_PREFIX.beamline_prefix}-EA-SM-01:",
+        x_infix="XP",
+        y_infix="YP",
+        z_infix="ZP",
+        polar_infix="ROTA",
+        azimuth_infix="ROTB",
+    )

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

@@ -1,3 +1,4 @@
-from dodal.devices.beamlines.b07.enums import Grating, LensMode, PsuMode
+from .b07_motors import B07SampleManipulator52B
+from .enums import Grating, LensMode, PsuMode
 
-__all__ = ["Grating", "LensMode", "PsuMode"]
+__all__ = ["B07SampleManipulator52B", "Grating", "LensMode", "PsuMode"]

src/dodal/devices/beamlines/b07/b07_motors.py

@@ -0,0 +1,100 @@
+from bluesky.protocols import Movable
+from ophyd_async.core import AsyncStatus, StandardReadable, StandardReadableFormat
+from ophyd_async.epics.core import epics_signal_rw_rbv
+from ophyd_async.epics.motor import Motor
+
+from dodal.devices.motors import _OMEGA, XYZStage
+
+
+class VirtualAxis(StandardReadable, Movable[float]):
+    """Represents the virtual axis coordinate system for the B07SampleManipulator52B.
+    The value signal is the read signal for this device.
+    """
+
+    def __init__(self, pv: str, name: str = ""):
+        """Initialise the device via pv and name configuration.
+
+        Args:
+            pv (str): The base PV to connect to the virtual axis.
+            name (str, optional): Name of this device.
+        """
+        with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL):
+            self.value = epics_signal_rw_rbv(float, pv, read_suffix=":RBV")
+        super().__init__(name)
+
+    @AsyncStatus.wrap
+    async def set(self, value: float):
+        await self.value.set(value)
+
+    def set_name(self, name: str, *, child_name_separator: str | None = None) -> None:
+        super().set_name(name, child_name_separator=child_name_separator)
+        # Value should be named the same as its parent in read()
+        self.value.set_name(name)
+
+
+class B07SampleManipulator52B(XYZStage):
+    """Sample Manipulator of several Motors and VirtaulAxis.
+
+    Attributes:
+        x (Motor): Controls x direction.
+        y (Motor): Controls y direction.
+        z (Motor): Controls z direction.
+        roty (Motor): Controls y rotation.
+        rotz (Motor): Controls z rotation.
+        xp (Motor): Compound motor for x.
+        yp (Motor): Compound motor for y.
+        zp (Motor): Compound motor for z.
+        omega (VirtualAxis): Virtal rotational axis for the x direction.
+        phi (VirtualAxis): Virtal rotational axis for the y direction.
+        kappa (VirtualAxis): Virtal rotational axis for the z direction.
+    """
+
+    def __init__(
+        self,
+        prefix: str,
+        x_infix: str = "XP",
+        y_infix: str = "YP",
+        z_infix: str = "ZP",
+        roty_infix: str = "ROTY",
+        rotz_infix: str = "ROTZ",
+        kappa_infix: str = "KAPPA",
+        phi_infix: str = "PHI",
+        omega_infix: str = _OMEGA,
+        name: str = "",
+    ):
+        """Initialise the device via prefix and infix PV configuration.
+
+        Args:
+            prefix (str): Base PV used for connecting signals.
+            x_infix (str, optional): Infix between base prefix and x motor record.
+            y_infix (str, optional): Infix between base prefix and y motor record.
+            z_infix (str, optional): Infix between base prefix and z motor record.
+            roty_infix (str, optional): Infix between base prefix and roty motor record.
+            rotz_infix (str, optional): Infix between base prefix and rotz motor record.
+            kappa_infix (str, optional): Infix between base prefix and kappa virtual axis.
+            phi_infix (str, optional): Infix between base prefix and phi virtual axis.
+            omega_infix (str, optional): Infix between base prefix and omega virtual axis.
+            name (str, optional): The name of this device.
+        """
+        with self.add_children_as_readables():
+            # Compound motors
+            self.xp = Motor(prefix + x_infix)
+            self.yp = Motor(prefix + y_infix)
+            self.zp = Motor(prefix + z_infix)
+
+            # Raw motors
+            self.roty = Motor(prefix + roty_infix)
+            self.rotz = Motor(prefix + rotz_infix)
+
+            # Not standard motors, virtual axes coordinate system.
+            self.kappa = VirtualAxis(prefix + kappa_infix)
+            self.phi = VirtualAxis(prefix + phi_infix)
+            self.omega = VirtualAxis(prefix + omega_infix)
+
+        super().__init__(
+            prefix=prefix,
+            x_infix=x_infix,
+            y_infix=y_infix,
+            z_infix=z_infix,
+            name=name,
+        )

src/dodal/devices/motors.py

@@ -275,7 +275,7 @@ def __init__(
         with self.add_children_as_readables():
             self.kappa = Motor(prefix + kappa_infix)
             self.phi = Motor(prefix + phi_infix)
-        super().__init__(prefix, name, x_infix, y_infix, z_infix)
+        super().__init__(prefix, name, x_infix, y_infix, z_infix, omega_infix)
 
         self.vertical_in_lab_space = create_axis_perp_to_rotation(
             self.omega, self.y, self.z

tests/devices/beamlines/b07/__init__.py

@@ -0,0 +1,66 @@
+import asyncio
+
+import pytest
+from ophyd_async.core import init_devices
+from ophyd_async.testing import assert_reading, partial_reading
+
+from dodal.devices.beamlines.b07 import B07SampleManipulator52B
+
+
+@pytest.fixture
+def sm52b() -> B07SampleManipulator52B:
+    with init_devices(mock=True):
+        sm52b = B07SampleManipulator52B(prefix="TEST:")
+    return sm52b
+
+
+@pytest.mark.parametrize(
+    "x, y, z, xp, yp, zp, roty, rotz, kappa, phi, omega",
+    [
+        (1.23, 2.40, 0.0, 0.0, 0.0, 0.0, 1.0, 1.2, 0, 2.0, 1.7),
+        (2.0, 1.2, 3.51, 24.0, 1.0, 2.0, 5.0, 21.0, 1.0, 25.0, 19.2),
+    ],
+)
+async def test_sm52b_read(
+    sm52b: B07SampleManipulator52B,
+    x: float,
+    y: float,
+    z: float,
+    xp: float,
+    yp: float,
+    zp: float,
+    roty: float,
+    rotz: float,
+    kappa: float,
+    phi: float,
+    omega: float,
+) -> None:
+    await asyncio.gather(
+        sm52b.x.set(x),
+        sm52b.y.set(y),
+        sm52b.z.set(z),
+        sm52b.xp.set(xp),
+        sm52b.yp.set(yp),
+        sm52b.zp.set(zp),
+        sm52b.roty.set(roty),
+        sm52b.rotz.set(rotz),
+        sm52b.kappa.set(kappa),
+        sm52b.phi.set(phi),
+        sm52b.omega.set(omega),
+    )
+    await assert_reading(
+        sm52b,
+        {
+            "sm52b-x": partial_reading(x),
+            "sm52b-y": partial_reading(y),
+            "sm52b-z": partial_reading(z),
+            "sm52b-xp": partial_reading(xp),
+            "sm52b-yp": partial_reading(yp),
+            "sm52b-zp": partial_reading(zp),
+            "sm52b-roty": partial_reading(roty),
+            "sm52b-rotz": partial_reading(rotz),
+            "sm52b-kappa": partial_reading(kappa),
+            "sm52b-phi": partial_reading(phi),
+            "sm52b-omega": partial_reading(omega),
+        },
+    )