Motors passed in as part of scanSpec stack and fix Complete()

src/ophyd_async/epics/pmac/_pmacTrajectory.py

@@ -1,8 +1,10 @@
 import time
 
 from bluesky.protocols import Flyable, Preparable
+from scanspec.specs import Frames
 
 from ophyd_async.core.async_status import AsyncStatus, WatchableAsyncStatus
+from ophyd_async.core.signal import observe_value
 from ophyd_async.core.utils import WatcherUpdate
 from ophyd_async.epics.pmac import Pmac, PmacCSMotor
 
@@ -12,13 +14,8 @@
 class PmacTrajectory(Pmac, Flyable, Preparable):
     """Device that moves a PMAC Motor record"""
 
-    def __init__(
-        self, prefix: str, cs: int, motors: list[PmacCSMotor], name=""
-    ) -> None:
+    def __init__(self, prefix: str, cs: int, name="") -> None:
         # Make a dict of which motors are for which cs axis
-        self.motors = {}
-        for motor in motors:
-            self.motors[motor.csAxis] = motor
         self._fly_start: float
         self.cs = cs
         super().__init__(prefix, cs, name=name)
@@ -34,45 +31,40 @@ async def _ramp_up_velocity_pos(
         return [disp, accl_time]
 
     @AsyncStatus.wrap
-    async def prepare(self, scanSpecStack):
+    async def prepare(self, stack: list[Frames[PmacCSMotor]]):
         # Which Axes are in use?
 
-        scanSize = len(scanSpecStack[0])
-        scanAxes = scanSpecStack[0].axes()
+        scanSize = len(stack[0])
+        scanAxes = stack[0].axes()
 
         self.profile = {}
         for axis in scanAxes:
-            self.profile[axis.lower()] = []
             if axis != "DURATION":
-                self.profile[axis.lower() + "_velocity"] = []
+                self.profile[axis.csAxis] = []
+                self.profile[axis.csAxis + "_velocity"] = []
+            else:
+                self.profile[axis.lower()] = []
 
         # Calc Velocity
 
         for axis in scanAxes:
             for i in range(scanSize - 1):
                 if axis != "DURATION":
-                    self.profile[axis.lower() + "_velocity"].append(
-                        (
-                            scanSpecStack[0].midpoints[axis][i + 1]
-                            - scanSpecStack[0].midpoints[axis][i]
-                        )
-                        / (scanSpecStack[0].midpoints["DURATION"][i])
-                    )
-                    self.profile[axis.lower()].append(
-                        scanSpecStack[0].midpoints[axis][i]
+                    self.profile[axis.csAxis + "_velocity"].append(
+                        (stack[0].midpoints[axis][i + 1] - stack[0].midpoints[axis][i])
+                        / (stack[0].midpoints["DURATION"][i])
                     )
+                    self.profile[axis.csAxis].append(stack[0].midpoints[axis][i])
                 else:
                     self.profile[axis.lower()].append(
-                        int(scanSpecStack[0].midpoints[axis][i] / TICK_S)
+                        int(stack[0].midpoints[axis][i] / TICK_S)
                     )
             if axis != "DURATION":
-                self.profile[axis.lower()].append(
-                    scanSpecStack[0].midpoints[axis][scanSize - 1]
-                )
-                self.profile[axis.lower() + "_velocity"].append(0)
+                self.profile[axis.csAxis].append(stack[0].midpoints[axis][scanSize - 1])
+                self.profile[axis.csAxis + "_velocity"].append(0)
             else:
                 self.profile[axis.lower()].append(
-                    int(scanSpecStack[0].midpoints[axis][scanSize - 1] / TICK_S)
+                    int(stack[0].midpoints[axis][scanSize - 1] / TICK_S)
                 )
 
         # Calculate Starting Position to allow ramp up to velocity
@@ -82,29 +74,31 @@ async def prepare(self, scanSpecStack):
             if axis != "DURATION":
                 run_up_disp, run_up_time = await self._ramp_up_velocity_pos(
                     0,
-                    self.motors[axis.lower()],
-                    self.profile[axis.lower() + "_velocity"][0],
+                    axis,
+                    self.profile[axis.csAxis + "_velocity"][0],
+                )
+                self.initial_pos[axis.csAxis] = (
+                    self.profile[axis.csAxis][0] - run_up_disp
                 )
-                self.initial_pos[axis] = self.profile[axis.lower()][0] - run_up_disp
         self.profile["duration"][0] += run_up_time / TICK_S
 
         # Send trajectory to brick
         for axis in scanAxes:
             if axis != "DURATION":
                 self.profile_cs_name.set(self.cs)
                 self.points_to_build.set(scanSize)
-                getattr(self, "use_" + axis.lower()).set(True)
-                getattr(self, axis.lower()).set(self.profile[axis.lower()])
-                getattr(self, axis.lower() + "_vel").set(
-                    self.profile[axis.lower() + "_velocity"]
+                getattr(self, "use_" + axis.csAxis).set(True)
+                getattr(self, axis.csAxis).set(self.profile[axis.csAxis])
+                getattr(self, axis.csAxis + "_vel").set(
+                    self.profile[axis.csAxis + "_velocity"]
                 )
             else:
                 self.timeArray.set(self.profile["duration"])
 
         # MOVE TO START
         for axis in scanAxes:
             if axis != "DURATION":
-                await self.motors[axis.lower()].set(self.initial_pos[axis])
+                await axis.set(self.initial_pos[axis.csAxis])
 
         # Set No Of Points
 
@@ -117,12 +111,13 @@ async def kickoff(self):
 
     @WatchableAsyncStatus.wrap
     async def complete(self):
-        yield WatcherUpdate(
-            name=self.name,
-            current=self.scan_percent,
-            initial=0,
-            target=100,
-            unit="%",
-            precision=0,
-            time_elapsed=time.monotonic() - self._fly_start,
-        )
+        async for percent in observe_value(self.scan_percent):
+            yield WatcherUpdate(
+                name=self.name,
+                current=percent,
+                initial=0,
+                target=100,
+                unit="%",
+                precision=0,
+                time_elapsed=time.monotonic() - self._fly_start,
+            )

tests/epics/pmac/test_pmac.py

@@ -23,13 +23,12 @@ async def sim_x_motor():
 
 async def test_sim_pmac_simple_trajectory(sim_x_motor) -> None:
     # Test the generated Trajectory profile from a scanspec
+    prefix = "BLxxI-MO-STEP-01"
     async with DeviceCollector(mock=True):
-        prefix = "BLxxI-MO-STEP-01"
-        motors = [sim_x_motor]
-        traj = PmacTrajectory(prefix, "BRICK1.CS3", motors, name="sim_pmac")
-        spec = fly(Line("z", 1, 5, 9), 1)
-        stack = spec.calculate()
-        await traj.prepare(stack)
+        traj = PmacTrajectory(prefix, "BRICK1.CS3", name="sim_pmac")
+    spec = fly(Line(sim_x_motor, 1, 5, 9), 1)
+    stack = spec.calculate()
+    await traj.prepare(stack)
     assert traj.profile == {
         "duration": [
             1050000,