Move aperture scatterguard independently (#932)

* Use the prepare verb to move the ap_sg whilst out of the beam

* Put moving out logic into set

* Add tests for moving aperture_scatterguard independently

* Add some docs about different ways of moving the aperture scatterguard

* Decouple setting enum from name in GDA parameters

* Make sure no axes are already moving when we do any other mvoe

src/dodal/devices/aperturescatterguard.py

@@ -2,7 +2,7 @@
 
 import asyncio
 
-from bluesky.protocols import Movable
+from bluesky.protocols import Movable, Preparable
 from ophyd_async.core import (
     AsyncStatus,
     StandardReadable,
@@ -21,6 +21,15 @@ class InvalidApertureMove(Exception):
     pass
 
 
+class _GDAParamApertureValue(StrictEnum):
+    """Maps from a short usable name to the value name in the GDA Beamline parameters"""
+
+    ROBOT_LOAD = "ROBOT_LOAD"
+    SMALL = "SMALL_APERTURE"
+    MEDIUM = "MEDIUM_APERTURE"
+    LARGE = "LARGE_APERTURE"
+
+
 class AperturePosition(BaseModel):
     """
     Represents one of the available positions for the Aperture-Scatterguard.
@@ -65,7 +74,7 @@ def tolerances_from_gda_params(
 
     @staticmethod
     def from_gda_params(
-        name: ApertureValue,
+        name: _GDAParamApertureValue,
         radius: float,
         params: GDABeamlineParameters,
     ) -> AperturePosition:
@@ -80,12 +89,16 @@ def from_gda_params(
 
 
 class ApertureValue(StrictEnum):
-    """Maps from a short usable name to the value name in the GDA Beamline parameters"""
+    """The possible apertures that can be selected.
+
+    Changing these means changing the external paramter model of Hyperion.
+    See https://github.com/DiamondLightSource/mx-bluesky/issues/760
+    """
 
-    ROBOT_LOAD = "ROBOT_LOAD"
     SMALL = "SMALL_APERTURE"
     MEDIUM = "MEDIUM_APERTURE"
     LARGE = "LARGE_APERTURE"
+    OUT_OF_BEAM = "Out of beam"
 
     def __str__(self):
         return self.name.capitalize()
@@ -95,22 +108,53 @@ def load_positions_from_beamline_parameters(
     params: GDABeamlineParameters,
 ) -> dict[ApertureValue, AperturePosition]:
     return {
-        ApertureValue.ROBOT_LOAD: AperturePosition.from_gda_params(
-            ApertureValue.ROBOT_LOAD, 0, params
+        ApertureValue.OUT_OF_BEAM: AperturePosition.from_gda_params(
+            _GDAParamApertureValue.ROBOT_LOAD, 0, params
         ),
         ApertureValue.SMALL: AperturePosition.from_gda_params(
-            ApertureValue.SMALL, 20, params
+            _GDAParamApertureValue.SMALL, 20, params
         ),
         ApertureValue.MEDIUM: AperturePosition.from_gda_params(
-            ApertureValue.MEDIUM, 50, params
+            _GDAParamApertureValue.MEDIUM, 50, params
         ),
         ApertureValue.LARGE: AperturePosition.from_gda_params(
-            ApertureValue.LARGE, 100, params
+            _GDAParamApertureValue.LARGE, 100, params
         ),
     }
 
 
-class ApertureScatterguard(StandardReadable, Movable):
+class ApertureScatterguard(StandardReadable, Movable, Preparable):
+    """Move the aperture and scatterguard assembly in a safe way. There are two ways to
+    interact with the device depending on if you want simplicity or move flexibility.
+
+    Examples:
+        The simple interface is using::
+
+            await aperture_scatterguard.set(ApertureValue.LARGE)
+
+        This will move the assembly so that the large aperture is in the beam, regardless
+        of where the assembly currently is.
+
+        We may also want to move the assembly out of the beam with::
+
+            await aperture_scatterguard.set(ApertureValue.OUT_OF_BEAM)
+
+        Note, to make sure we do this as quickly as possible, the scatterguard will stay
+        in the same position relative to the aperture.
+
+        We may then want to keep the assembly out of the beam whilst asynchronously preparing
+        the other axes for the aperture that's to follow::
+
+            await aperture_scatterguard.prepare(ApertureValue.LARGE)
+
+        Then, at a later time, move back into the beam::
+
+            await aperture_scatterguard.set(ApertureValue.LARGE)
+
+        Given the prepare has been done this move will now be faster as only the y is
+        left to move.
+    """
+
     def __init__(
         self,
         loaded_positions: dict[ApertureValue, AperturePosition],
@@ -135,22 +179,93 @@ def __init__(
                 self.radius,
             ],
         )
+
         with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL):
             self.selected_aperture = create_hardware_backed_soft_signal(
                 ApertureValue, self._get_current_aperture_position
             )
 
         super().__init__(name)
 
-    def get_position_from_gda_aperture_name(
-        self, gda_aperture_name: str
-    ) -> ApertureValue:
-        return ApertureValue(gda_aperture_name)
-
     @AsyncStatus.wrap
     async def set(self, value: ApertureValue):
+        """This set will move the aperture into the beam or move the whole assembly out"""
+
         position = self._loaded_positions[value]
-        await self._safe_move_within_datacollection_range(position, value)
+        await self._check_safe_to_move(position.aperture_z)
+
+        if value == ApertureValue.OUT_OF_BEAM:
+            out_y = self._loaded_positions[ApertureValue.OUT_OF_BEAM].aperture_y
+            await self.aperture.y.set(out_y)
+        else:
+            await self._safe_move_whilst_in_beam(position)
+
+    async def _check_safe_to_move(self, expected_z_position: float):
+        """The assembly is moved (in z) to be under the table when the beamline is not
+        in use. If we try and move whilst in the incorrect Z position we will collide
+        with the table.
+
+        Additionally, because there are so many collision possibilities in the device we
+        throw an error if any of the axes are already moving.
+        """
+        current_ap_z = await self.aperture.z.user_readback.get_value()
+        diff_on_z = abs(current_ap_z - expected_z_position)
+        aperture_z_tolerance = self._tolerances.aperture_z
+        if diff_on_z > aperture_z_tolerance:
+            raise InvalidApertureMove(
+                f"Current aperture z ({current_ap_z}), outside of tolerance ({aperture_z_tolerance}) from target ({expected_z_position})."
+            )
+
+        all_axes = [
+            self.aperture.x,
+            self.aperture.y,
+            self.aperture.z,
+            self.scatterguard.x,
+            self.scatterguard.y,
+        ]
+        for axis in all_axes:
+            axis_stationary = await axis.motor_done_move.get_value()
+            if not axis_stationary:
+                raise InvalidApertureMove(
+                    f"{axis.name} is still moving. Wait for it to finish before"
+                    "triggering another move."
+                )
+
+    async def _safe_move_whilst_in_beam(self, position: AperturePosition):
+        """
+        Move the aperture and scatterguard combo safely to a new position.
+        See https://github.com/DiamondLightSource/hyperion/wiki/Aperture-Scatterguard-Collisions
+        for why this is required. TLDR is that we have a collision at the top of y so we need
+        to make sure we move the assembly down before we move the scatterguard up.
+        """
+        current_ap_y = await self.aperture.y.user_readback.get_value()
+
+        aperture_x, aperture_y, aperture_z, scatterguard_x, scatterguard_y = (
+            position.values
+        )
+
+        if aperture_y > current_ap_y:
+            # Assembly needs to move up so move the scatterguard down first
+            await asyncio.gather(
+                self.scatterguard.x.set(scatterguard_x),
+                self.scatterguard.y.set(scatterguard_y),
+            )
+            await asyncio.gather(
+                self.aperture.x.set(aperture_x),
+                self.aperture.y.set(aperture_y),
+                self.aperture.z.set(aperture_z),
+            )
+        else:
+            await asyncio.gather(
+                self.aperture.x.set(aperture_x),
+                self.aperture.y.set(aperture_y),
+                self.aperture.z.set(aperture_z),
+            )
+
+            await asyncio.gather(
+                self.scatterguard.x.set(scatterguard_x),
+                self.scatterguard.y.set(scatterguard_y),
+            )
 
     @AsyncStatus.wrap
     async def _set_raw_unsafe(self, position: AperturePosition):
@@ -167,80 +282,51 @@ async def _set_raw_unsafe(self, position: AperturePosition):
             self.scatterguard.y.set(scatterguard_y),
         )
 
+    async def _is_out_of_beam(self) -> bool:
+        current_ap_y = await self.aperture.y.user_readback.get_value()
+        out_ap_y = self._loaded_positions[ApertureValue.OUT_OF_BEAM].aperture_y
+        return current_ap_y <= out_ap_y + self._tolerances.aperture_y
+
     async def _get_current_aperture_position(self) -> ApertureValue:
         """
         Returns the current aperture position using readback values
         for SMALL, MEDIUM, LARGE. ROBOT_LOAD position defined when
         mini aperture y <= ROBOT_LOAD.location.aperture_y + tolerance.
         If no position is found then raises InvalidApertureMove.
         """
-        current_ap_y = await self.aperture.y.user_readback.get_value(cached=False)
-        robot_load_ap_y = self._loaded_positions[ApertureValue.ROBOT_LOAD].aperture_y
         if await self.aperture.large.get_value(cached=False) == 1:
             return ApertureValue.LARGE
         elif await self.aperture.medium.get_value(cached=False) == 1:
             return ApertureValue.MEDIUM
         elif await self.aperture.small.get_value(cached=False) == 1:
             return ApertureValue.SMALL
-        elif current_ap_y <= robot_load_ap_y + self._tolerances.aperture_y:
-            return ApertureValue.ROBOT_LOAD
+        elif await self._is_out_of_beam():
+            return ApertureValue.OUT_OF_BEAM
 
         raise InvalidApertureMove("Current aperture/scatterguard state unrecognised")
 
     async def _get_current_radius(self) -> float:
         current_value = await self._get_current_aperture_position()
         return self._loaded_positions[current_value].radius
 
-    async def _safe_move_within_datacollection_range(
-        self, position: AperturePosition, value: ApertureValue
-    ):
-        """
-        Move the aperture and scatterguard combo safely to a new position.
-        See https://github.com/DiamondLightSource/hyperion/wiki/Aperture-Scatterguard-Collisions
-        for why this is required.
-        """
-        assert self._loaded_positions is not None
-
-        ap_z_in_position = await self.aperture.z.motor_done_move.get_value()
-        if not ap_z_in_position:
-            raise InvalidApertureMove(
-                "ApertureScatterguard z is still moving. Wait for it to finish "
-                "before triggering another move."
-            )
+    @AsyncStatus.wrap
+    async def prepare(self, value: ApertureValue):
+        """Moves the assembly to the position for the specified aperture, whilst keeping
+        it out of the beam if it already is so.
 
-        current_ap_z = await self.aperture.z.user_readback.get_value()
-        diff_on_z = abs(current_ap_z - position.aperture_z)
-        if diff_on_z > self._tolerances.aperture_z:
-            raise InvalidApertureMove(
-                "ApertureScatterguard safe move is not yet defined for positions "
-                "outside of LARGE, MEDIUM, SMALL, ROBOT_LOAD. "
-                f"Current aperture z ({current_ap_z}), outside of tolerance ({self._tolerances.aperture_z}) from target ({position.aperture_z})."
+        Moving the assembly whilst out of the beam has no collision risk so we can just
+        move all the motors together.
+        """
+        if await self._is_out_of_beam():
+            aperture_x, _, aperture_z, scatterguard_x, scatterguard_y = (
+                self._loaded_positions[value].values
             )
 
-        current_ap_y = await self.aperture.y.user_readback.get_value()
-
-        aperture_x, aperture_y, aperture_z, scatterguard_x, scatterguard_y = (
-            position.values
-        )
-
-        if position.aperture_y > current_ap_y:
-            await asyncio.gather(
-                self.scatterguard.x.set(scatterguard_x),
-                self.scatterguard.y.set(scatterguard_y),
-            )
-            await asyncio.gather(
-                self.aperture.x.set(aperture_x),
-                self.aperture.y.set(aperture_y),
-                self.aperture.z.set(aperture_z),
-            )
-        else:
             await asyncio.gather(
                 self.aperture.x.set(aperture_x),
-                self.aperture.y.set(aperture_y),
                 self.aperture.z.set(aperture_z),
-            )
-
-            await asyncio.gather(
                 self.scatterguard.x.set(scatterguard_x),
                 self.scatterguard.y.set(scatterguard_y),
             )
+        else:
+            await self.set(value)