Refactor plans to simplify

DiamondLightSource · Sep 13, 2024 · eeee9ac · eeee9ac
1 parent d1332b5
commit eeee9ac
Showing 1 changed file with 58 additions and 55 deletions.
diff --git a/src/mx_bluesky/hyperion/experiment_plans/robot_load_then_centre_plan.py b/src/mx_bluesky/hyperion/experiment_plans/robot_load_then_centre_plan.py
@@ -187,83 +187,86 @@ def raise_exception_if_moved_out_of_cryojet(exception):
         )
 
 
-def pin_already_loaded(
+def _pin_already_loaded(
     robot: BartRobot, pin_to_load: int, puck_to_load: int
 ) -> Generator[Msg, None, bool]:
     current_puck = yield from bps.rd(robot.current_puck)
     current_pin = yield from bps.rd(robot.current_pin)
     return int(current_puck) == puck_to_load and int(current_pin) == pin_to_load
 
 
-def robot_load_then_centre_plan(
+def robot_load_and_snapshots(
     composite: RobotLoadThenCentreComposite,
     params: RobotLoadThenCentre,
+    location: SampleLocation,
 ):
-    # TODO: get these from one source of truth #1347
-    assert params.sample_puck is not None
-    assert params.sample_pin is not None
+    robot_load_plan = do_robot_load(
+        composite,
+        location,
+        params.demand_energy_ev,
+        params.thawing_time,
+    )
 
-    @bpp.run_decorator(
-        md={
-            "subplan_name": CONST.PLAN.ROBOT_LOAD,
-            "metadata": {
-                "visit_path": str(params.visit_directory),
-                "sample_id": params.sample_id,
-                "sample_puck": params.sample_puck,
-                "sample_pin": params.sample_pin,
-            },
-            "activate_callbacks": [
-                "RobotLoadISPyBCallback",
-            ],
-        }
+    # 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,
     )
-    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
-
-        robot_load_plan = do_robot_load(
-            composite,
-            SampleLocation(params.sample_puck, params.sample_pin),
-            params.demand_energy_ev,
-            params.thawing_time,
-        )
 
-        # 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 take_robot_snapshots(
+        composite.oav, composite.webcam, params.snapshot_directory
+    )
 
-        yield from take_robot_snapshots(
-            composite.oav, composite.webcam, params.snapshot_directory
-        )
+    yield from bps.create(name=CONST.DESCRIPTORS.ROBOT_LOAD)
+    yield from bps.read(composite.robot.barcode)
+    yield from bps.read(composite.oav.snapshot)
+    yield from bps.read(composite.webcam)
+    yield from bps.save()
 
-        yield from bps.create(name=CONST.DESCRIPTORS.ROBOT_LOAD)
-        yield from bps.read(composite.robot.barcode)
-        yield from bps.read(composite.oav.snapshot)
-        yield from bps.read(composite.webcam)
-        yield from bps.save()
+    yield from bps.wait("reset-lower_gonio")
 
-        yield from bps.wait("reset-lower_gonio")
+
+def robot_load_then_centre_plan(
+    composite: RobotLoadThenCentreComposite,
+    params: RobotLoadThenCentre,
+):
+    # TODO: get these from one source of truth #254
+    assert params.sample_puck is not None
+    assert params.sample_pin is not None
 
     if not (
-        yield from pin_already_loaded(
+        yield from _pin_already_loaded(
             composite.robot, params.sample_pin, params.sample_puck
         )
     ):
         yield from prepare_for_robot_load(composite)
-        yield from robot_load_and_snapshots()
+        yield from bpp.run_wrapper(
+            robot_load_and_snapshots(
+                composite, params, SampleLocation(params.sample_puck, params.sample_pin)
+            ),
+            md={
+                "subplan_name": CONST.PLAN.ROBOT_LOAD,
+                "metadata": {
+                    "visit_path": str(params.visit_directory),
+                    "sample_id": params.sample_id,
+                    "sample_puck": params.sample_puck,
+                    "sample_pin": params.sample_pin,
+                },
+                "activate_callbacks": [
+                    "RobotLoadISPyBCallback",
+                ],
+            },
+        )
     else:
         LOGGER.info(
             f"Pin/puck {params.sample_pin}/{params.sample_puck} already loaded, will not reload."