PythonModels/systems/test_mql_mechanical.py

@@ -210,10 +210,20 @@ class TestMqlMechanicalNodeKinematics:
     displacements_m: dict[int, float]
 
 
+@dataclass(frozen=True)
+class TestMqlPistonKinematics:
+    alias: str
+    port_2_velocity_m_s: float
+    port_2_displacement_m: float
+    port_3_velocity_m_s: float
+    port_3_displacement_m: float
+
+
 @dataclass(frozen=True)
 class TestMqlMechanicalMassSnapshot:
     states: tuple[TestMqlMechanicalMassState, ...]
     node_kinematics_by_alias: dict[str, TestMqlMechanicalNodeKinematics]
+    piston_kinematics_by_alias: dict[str, TestMqlPistonKinematics]
 
     @property
     def state_count(self) -> int:
@@ -244,9 +254,14 @@ class TestMqlMechanicalMassClosure:
             )
             for index, alias in enumerate(self.mass_aliases)
         )
+        node_kinematics = self._node_kinematics_by_alias(states)
         return TestMqlMechanicalMassSnapshot(
             states=states,
-            node_kinematics_by_alias=self._node_kinematics_by_alias(states),
+            node_kinematics_by_alias=node_kinematics,
+            piston_kinematics_by_alias=self._piston_kinematics_by_alias(
+                states,
+                node_kinematics,
+            ),
         )
 
     def _node_kinematics_by_alias(
@@ -269,6 +284,34 @@ class TestMqlMechanicalMassClosure:
             ),
         }
 
+    def _piston_kinematics_by_alias(
+        self,
+        states: tuple[TestMqlMechanicalMassState, ...],
+        node_kinematics_by_alias: dict[str, TestMqlMechanicalNodeKinematics],
+    ) -> dict[str, TestMqlPistonKinematics]:
+        state_by_alias = {state.alias: state for state in states}
+        rear_node = node_kinematics_by_alias["dynamic_mechanical_node_alternative_3"]
+        piston_bindings = (
+            ("pn_brp2_8", "mass_friction_endstops_10", 8),
+            ("pn_brp2_9", "mass_friction_endstops_11", 7),
+            ("pn_brp2_10", "mass_friction_endstops_12", 6),
+            ("pn_brp2_11", "mass_friction_endstops_13", 5),
+            ("pn_brp2_12", "mass_friction_endstops_14", 4),
+            ("pn_brp2_13", "mass_friction_endstops_15", 3),
+            ("pn_brp2_14", "mass_friction_endstops_16", 2),
+            ("pn_brp2_15", "mass_friction_endstops_17", 1),
+        )
+        return {
+            piston_alias: TestMqlPistonKinematics(
+                alias=piston_alias,
+                port_2_velocity_m_s=state_by_alias[mass_alias].velocity_m_s,
+                port_2_displacement_m=state_by_alias[mass_alias].displacement_m,
+                port_3_velocity_m_s=rear_node.velocities_m_s[rear_node_port],
+                port_3_displacement_m=rear_node.displacements_m[rear_node_port],
+            )
+            for piston_alias, mass_alias, rear_node_port in piston_bindings
+        }
+
     def rhs(
         self,
         state_vector: list[float],

tests/test_test_mql_mechanical.py

@@ -120,6 +120,30 @@ class TestMqlMechanicalAssemblyTests(unittest.TestCase):
         self.assertAlmostEqual(rear.velocities_m_s[1], -0.3)
         self.assertAlmostEqual(rear.displacements_m[8], -0.72)
 
+    def test_mechanical_mass_snapshot_exposes_piston_kinematics(self) -> None:
+        closure = TestMqlMechanicalMassClosure(self.assembly)
+        state = closure.initial_state_vector()
+        state[0] = 0.4
+        state[1] = 0.37
+        state[18] = -0.3
+        state[19] = -0.72
+
+        snapshot = closure.snapshot(state)
+        piston = snapshot.piston_kinematics_by_alias["pn_brp2_8"]
+
+        self.assertEqual(set(snapshot.piston_kinematics_by_alias), set(self.assembly.pistons))
+        self.assertAlmostEqual(piston.port_2_velocity_m_s, 0.4)
+        self.assertAlmostEqual(piston.port_2_displacement_m, 0.37)
+        self.assertAlmostEqual(piston.port_3_velocity_m_s, -0.3)
+        self.assertAlmostEqual(piston.port_3_displacement_m, -0.72)
+        self.assertAlmostEqual(
+            self.assembly.pistons["pn_brp2_8"].geometry().chamber_length_mm(
+                piston.port_3_displacement_m,
+                piston.port_2_displacement_m,
+            ),
+            1090.0,
+        )
+
     def test_system_simulates_mechanical_mass_state_closure(self) -> None:
         system = TestMqlSystem()