Cache lower gonio positions before robot load and restore them afterw…

…ards (#1467)

* First draft of caching lower gonios before robot load

* Add test for returning lower_gonio to correct position

* Specify argument for readability

setup.cfg

@@ -40,7 +40,7 @@ install_requires =
     ophyd-async >= 0.3a5
     bluesky >= 1.13.0a3
     blueapi >= 0.4.3-rc1
-    dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git@95c475d17a0d60dd32e2186395520b01ca445cb4
+    dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git@fff4a4e8fdcf534de6768c2b45a0dfa7a2e2ccc4
 
 [options.entry_points]
 console_scripts =

src/hyperion/experiment_plans/robot_load_then_centre_plan.py

@@ -17,6 +17,7 @@
 from dodal.devices.fast_grid_scan import PandAFastGridScan, ZebraFastGridScan
 from dodal.devices.flux import Flux
 from dodal.devices.focusing_mirror import FocusingMirrorWithStripes, VFMMirrorVoltages
+from dodal.devices.motors import XYZPositioner
 from dodal.devices.oav.oav_detector import OAV
 from dodal.devices.oav.pin_image_recognition import PinTipDetection
 from dodal.devices.robot import BartRobot, SampleLocation
@@ -30,6 +31,7 @@
 from dodal.devices.xbpm_feedback import XBPMFeedback
 from dodal.devices.zebra import Zebra
 from dodal.devices.zocalo import ZocaloResults
+from dodal.plans.motor_util_plans import MoveTooLarge, home_and_reset_wrapper
 from ophyd_async.panda import HDFPanda
 
 from hyperion.device_setup_plans.utils import (
@@ -85,6 +87,7 @@ class RobotLoadThenCentreComposite:
     # RobotLoad fields
     robot: BartRobot
     webcam: Webcam
+    lower_gonio: XYZPositioner
 
 
 def create_devices(context: BlueskyContext) -> RobotLoadThenCentreComposite:
@@ -146,6 +149,36 @@ def prepare_for_robot_load(composite: RobotLoadThenCentreComposite):
     yield from bps.wait("prepare_robot_load")
 
 
+def robot_load_and_energy_change(
+    composite: RobotLoadThenCentreComposite,
+    sample_location: SampleLocation,
+    demand_energy_ev: float | None,
+):
+    yield from bps.abs_set(
+        composite.robot,
+        sample_location,
+        group="robot_load",
+    )
+
+    if demand_energy_ev:
+        yield from set_energy_plan(
+            demand_energy_ev / 1000,
+            cast(SetEnergyComposite, composite),
+        )
+
+    yield from bps.wait("robot_load")
+
+
+def raise_exception_if_moved_out_of_cryojet(exception):
+    yield from bps.null()
+    if isinstance(exception, MoveTooLarge):
+        raise Exception(
+            f"Moving {exception.axis} back to {exception.position} after \
+                        robot load would move it out of the cryojet. The max safe \
+                        distance is {exception.maximum_move}"
+        )
+
+
 def robot_load_then_centre_plan(
     composite: RobotLoadThenCentreComposite,
     params: RobotLoadThenCentre,
@@ -166,23 +199,32 @@ def robot_load_then_centre_plan(
             ],
         }
     )
-    def robot_load():
+    def robot_load_and_snapshots():
         # TODO: get these from one source of truth #1347
         assert params.sample_puck is not None
         assert params.sample_pin is not None
-        yield from bps.abs_set(
-            composite.robot,
+
+        robot_load_plan = robot_load_and_energy_change(
+            composite,
             SampleLocation(params.sample_puck, params.sample_pin),
-            group="robot_load",
+            params.demand_energy_ev,
         )
 
-        if params.demand_energy_ev:
-            yield from set_energy_plan(
-                params.demand_energy_ev / 1000,
-                cast(SetEnergyComposite, composite),
-            )
-
-        yield from bps.wait("robot_load")
+        # The lower gonio must be in the correct position for the robot load and we
+        # want to put it back afterwards. Note we don't wait the robot is interlocked
+        # to the lower gonio and the  move is quicker than the robot takes to get to the
+        # load position.
+        yield from bpp.contingency_wrapper(
+            home_and_reset_wrapper(
+                robot_load_plan,
+                composite.lower_gonio,
+                BartRobot.LOAD_TOLERANCE_MM,
+                CONST.HARDWARE.CRYOJET_MARGIN_MM,
+                "lower_gonio",
+                wait_for_all=False,
+            ),
+            except_plan=raise_exception_if_moved_out_of_cryojet,
+        )
 
         yield from bps.abs_set(composite.thawer.thaw_for_time_s, params.thawing_time)
         yield from wait_for_smargon_not_disabled(composite.smargon)
@@ -197,7 +239,9 @@ def robot_load():
         yield from bps.read(composite.webcam)
         yield from bps.save()
 
-    yield from robot_load()
+        yield from bps.wait("reset-lower_gonio")
+
+    yield from robot_load_and_snapshots()
 
     yield from pin_centre_then_xray_centre_plan(
         cast(GridDetectThenXRayCentreComposite, composite),

src/hyperion/parameters/constants.py

@@ -57,6 +57,7 @@ class DocDescriptorNames:
 class HardwareConstants:
     OAV_REFRESH_DELAY = 0.3
     PANDA_FGS_RUN_UP_DEFAULT = 0.17
+    CRYOJET_MARGIN_MM = 0.2
 
 
 @dataclass(frozen=True)

tests/conftest.py

@@ -40,6 +40,7 @@
 from dodal.devices.synchrotron import Synchrotron, SynchrotronMode
 from dodal.devices.thawer import Thawer
 from dodal.devices.undulator import Undulator
+from dodal.devices.util.test_utils import patch_motor as oa_patch_motor
 from dodal.devices.webcam import Webcam
 from dodal.devices.zebra import Zebra
 from dodal.log import LOGGER as dodal_logger
@@ -417,6 +418,17 @@ def vfm(RE):
     return vfm
 
 
+@pytest.fixture
+def lower_gonio(RE):
+    lower_gonio = i03.lower_gonio(fake_with_ophyd_sim=True)
+    with (
+        oa_patch_motor(lower_gonio.x),
+        oa_patch_motor(lower_gonio.y),
+        oa_patch_motor(lower_gonio.z),
+    ):
+        yield lower_gonio
+
+
 @pytest.fixture
 def vfm_mirror_voltages():
     voltages = i03.vfm_mirror_voltages(fake_with_ophyd_sim=True)

tests/unit_tests/experiment_plans/test_wait_for_robot_load_then_centre.py

@@ -1,3 +1,4 @@
+from functools import partial
 from pathlib import Path
 from unittest.mock import MagicMock, patch
 
@@ -28,7 +29,7 @@
 
 @pytest.fixture
 def robot_load_composite(
-    smargon, dcm, robot, aperture_scatterguard, oav, webcam, thawer
+    smargon, dcm, robot, aperture_scatterguard, oav, webcam, thawer, lower_gonio
 ) -> RobotLoadThenCentreComposite:
     composite: RobotLoadThenCentreComposite = MagicMock()
     composite.smargon = smargon
@@ -40,6 +41,7 @@ def robot_load_composite(
     composite.aperture_scatterguard.set = MagicMock(return_value=NullStatus())
     composite.oav = oav
     composite.webcam = webcam
+    composite.lower_gonio = lower_gonio
     composite.thawer = thawer
     return composite
 
@@ -356,6 +358,49 @@ async def test_when_take_snapshots_called_then_filename_and_directory_set_and_de
     assert (await webcam.directory.get_value()) == TEST_DIRECTORY
 
 
+@patch(
+    "hyperion.experiment_plans.robot_load_then_centre_plan.pin_centre_then_xray_centre_plan",
+    MagicMock(),
+)
+def test_given_lower_gonio_moved_when_robot_load_then_lower_gonio_moved_to_home_and_back(
+    robot_load_composite: RobotLoadThenCentreComposite,
+    robot_load_then_centre_params_no_energy: RobotLoadThenCentre,
+    sim_run_engine,
+):
+    initial_values = {"x": 0.11, "y": 0.12, "z": 0.13}
+
+    def get_read(axis, msg):
+        return {f"lower_gonio-{axis}": {"value": initial_values[axis]}}
+
+    for axis in initial_values.keys():
+        sim_run_engine.add_handler(
+            "read", f"lower_gonio-{axis}", partial(get_read, axis)
+        )
+
+    messages = sim_run_engine.simulate_plan(
+        robot_load_then_centre(
+            robot_load_composite,
+            robot_load_then_centre_params_no_energy,
+        )
+    )
+
+    for axis in initial_values.keys():
+        messages = sim_run_engine.assert_message_and_return_remaining(
+            messages,
+            lambda msg: msg.command == "set"
+            and msg.obj.name == f"lower_gonio-{axis}"
+            and msg.args == (0,),
+        )
+
+    for axis, initial in initial_values.items():
+        messages = sim_run_engine.assert_message_and_return_remaining(
+            messages,
+            lambda msg: msg.command == "set"
+            and msg.obj.name == f"lower_gonio-{axis}"
+            and msg.args == (initial,),
+        )
+
+
 @patch(
     "hyperion.experiment_plans.robot_load_then_centre_plan.pin_centre_then_xray_centre_plan"
 )