PythonModels/systems/test_mql.py

@@ -3837,6 +3837,93 @@ class TestMqlSimulationResult:
     series: dict[str, list[float]]
 
 
+@dataclass(frozen=True)
+class TestMqlFullStateSnapshot:
+    pneumatic: object
+    mechanical: object
+
+    @property
+    def state_count(self) -> int:
+        return self.pneumatic_state_count + self.mechanical_state_count
+
+    @property
+    def pneumatic_state_count(self) -> int:
+        return 112
+
+    @property
+    def mechanical_state_count(self) -> int:
+        return self.mechanical.state_count
+
+
+_PISTON_FORCE_BINDINGS = (
+    ("mass_friction_endstops_10", "pn_brp2_8", "p4_port3_remote_primary_chamber"),
+    ("mass_friction_endstops_11", "pn_brp2_9", "p4_primary_chamber"),
+    ("mass_friction_endstops_12", "pn_brp2_10", "p4_port1_remote_primary_chamber"),
+    ("mass_friction_endstops_13", "pn_brp2_11", "p4_port1_far_primary_chamber"),
+    ("mass_friction_endstops_14", "pn_brp2_12", "p4_port1_next_primary_chamber"),
+    ("mass_friction_endstops_15", "pn_brp2_13", "p4_bridge_primary_chamber"),
+    ("mass_friction_endstops_16", "pn_brp2_14", "p4_port3_next_primary_chamber"),
+    ("mass_friction_endstops_17", "pn_brp2_15", "p4_port3_far_primary_chamber"),
+)
+
+
+class TestMqlFullStateClosure:
+    def __init__(self, *, pneumatic_closure: object, mechanical_closure: object) -> None:
+        self.pneumatic_closure = pneumatic_closure
+        self.mechanical_closure = mechanical_closure
+        self.pneumatic_state_count = len(pneumatic_closure.initial_state_vector())
+        self.mechanical_state_count = len(mechanical_closure.initial_state_vector())
+
+    @property
+    def state_count(self) -> int:
+        return self.pneumatic_state_count + self.mechanical_state_count
+
+    def initial_state_vector(self) -> list[float]:
+        return [
+            *self.pneumatic_closure.initial_state_vector(),
+            *self.mechanical_closure.initial_state_vector(),
+        ]
+
+    def _split_state(self, state_vector: list[float]) -> tuple[list[float], list[float]]:
+        if len(state_vector) != self.state_count:
+            raise ValueError("full test_mql state vector requires 132 values")
+        return (
+            state_vector[: self.pneumatic_state_count],
+            state_vector[self.pneumatic_state_count :],
+        )
+
+    def snapshot(self, state_vector: list[float] | None = None) -> TestMqlFullStateSnapshot:
+        values = self.initial_state_vector() if state_vector is None else list(state_vector)
+        pneumatic_state, mechanical_state = self._split_state(values)
+        return TestMqlFullStateSnapshot(
+            pneumatic=self.pneumatic_closure.snapshot(pneumatic_state),
+            mechanical=self.mechanical_closure.snapshot(mechanical_state),
+        )
+
+    def rhs(self, state_vector: list[float]) -> list[float]:
+        pneumatic_state, mechanical_state = self._split_state(list(state_vector))
+        pneumatic_snapshot = self.pneumatic_closure.snapshot(pneumatic_state)
+        return [
+            *self.pneumatic_closure.rhs(pneumatic_state),
+            *self.mechanical_closure.rhs(
+                mechanical_state,
+                force_by_mass_alias=self._piston_force_by_mass_alias(pneumatic_snapshot),
+            ),
+        ]
+
+    def _piston_force_by_mass_alias(self, pneumatic_snapshot: object) -> dict[str, float]:
+        from PythonModels.systems.test_mql_pneumatic import pressure_to_amesim_gauge_pa
+
+        force_by_mass_alias: dict[str, float] = {}
+        for mass_alias, piston_alias, chamber_field in _PISTON_FORCE_BINDINGS:
+            chamber = getattr(pneumatic_snapshot, chamber_field)
+            piston = self.mechanical_closure.assembly.pistons[piston_alias]
+            force_by_mass_alias[mass_alias] = (
+                pressure_to_amesim_gauge_pa(chamber.p) * piston.annulus_area_m2
+            )
+        return force_by_mass_alias
+
+
 class AmesimStructuralComponent(Component):
     def __init__(self, spec: dict[str, object]) -> None:
         super().__init__(name=str(spec["alias"]))
@@ -4141,6 +4228,51 @@ class TestMqlSystem:
             t_eval=t_eval,
         )
 
+    def full_state_closure_from_spec(
+        self,
+        spec,
+        *,
+        inlet_node_pressure_pa: float,
+        resistance_boundary_pressure_pa: float,
+        inlet_node_temperature_k: float = 293.15,
+        resistance_boundary_temperature_k: float = 293.15,
+    ) -> TestMqlFullStateClosure:
+        return TestMqlFullStateClosure(
+            pneumatic_closure=self.pn3_p4_node_chamber_segment_closure_from_spec(
+                spec,
+                inlet_node_pressure_pa=inlet_node_pressure_pa,
+                resistance_boundary_pressure_pa=resistance_boundary_pressure_pa,
+                inlet_node_temperature_k=inlet_node_temperature_k,
+                resistance_boundary_temperature_k=resistance_boundary_temperature_k,
+            ),
+            mechanical_closure=self.mechanical_mass_closure(),
+        )
+
+    def simulate_full_state_from_spec(
+        self,
+        spec,
+        *,
+        inlet_node_pressure_pa: float,
+        resistance_boundary_pressure_pa: float,
+        inlet_node_temperature_k: float = 293.15,
+        resistance_boundary_temperature_k: float = 293.15,
+        config: SolveIVPConfig | None = None,
+        t_eval: list[float] | None = None,
+    ):
+        closure = self.full_state_closure_from_spec(
+            spec,
+            inlet_node_pressure_pa=inlet_node_pressure_pa,
+            resistance_boundary_pressure_pa=resistance_boundary_pressure_pa,
+            inlet_node_temperature_k=inlet_node_temperature_k,
+            resistance_boundary_temperature_k=resistance_boundary_temperature_k,
+        )
+        return integrate_ode(
+            rhs=lambda t, state: closure.rhs(state),
+            initial_state=closure.initial_state_vector(),
+            config=config or SolveIVPConfig(t_stop=1.0e-4, max_step=1.0e-5),
+            t_eval=t_eval,
+        )
+
     @staticmethod
     def pneumatic_branch_spec(
         *,

PythonModels/systems/test_mql_mechanical.py

@@ -188,14 +188,21 @@ class TestMqlMechanicalMassClosure:
         )
         return TestMqlMechanicalMassSnapshot(states=states)
 
-    def rhs(self, state_vector: list[float]) -> list[float]:
+    def rhs(
+        self,
+        state_vector: list[float],
+        *,
+        force_by_mass_alias: dict[str, float] | None = None,
+    ) -> list[float]:
         snapshot = self.snapshot(state_vector)
+        force_by_mass_alias = force_by_mass_alias or {}
         derivatives: list[float] = []
         for state in snapshot.states:
             mass = self.assembly.masses[state.alias].endstop()
             acceleration, velocity = mass.derivatives(
                 velocity_m_s=state.velocity_m_s,
                 displacement_m=state.displacement_m,
+                port_1_force_n=force_by_mass_alias.get(state.alias, 0.0),
             )
             derivatives.extend([acceleration, velocity])
         return derivatives

tests/test_test_mql_pnl0001_segment.py

@@ -526,6 +526,48 @@ class TestMqlPn3P4NodeChamberSegmentTests(unittest.TestCase):
             snapshot.p4_primary_line.T,
         )
 
+    def test_full_state_closure_combines_pneumatic_and_mechanical_states(self) -> None:
+        closure = self.system.full_state_closure_from_spec(
+            self.spec,
+            inlet_node_pressure_pa=15.31e6,
+            resistance_boundary_pressure_pa=15.29e6,
+        )
+
+        state = closure.initial_state_vector()
+        snapshot = closure.snapshot(state)
+        rhs = closure.rhs(state)
+
+        self.assertEqual(closure.pneumatic_state_count, 112)
+        self.assertEqual(closure.mechanical_state_count, 20)
+        self.assertEqual(closure.state_count, 132)
+        self.assertEqual(len(state), 132)
+        self.assertEqual(snapshot.state_count, 132)
+        self.assertEqual(snapshot.mechanical.state_count, 20)
+        self.assertEqual(len(rhs), 132)
+        self.assertAlmostEqual(
+            snapshot.pneumatic.p4_balance.pressure_pa,
+            snapshot.pneumatic.p4_primary_line.p,
+        )
+        self.assertEqual(
+            snapshot.mechanical.states[0].alias,
+            "mass_friction_endstops_10",
+        )
+        self.assertAlmostEqual(rhs[112], -0.8167936695810917)
+
+    def test_simulates_full_132_state_closure(self) -> None:
+        solution = self.system.simulate_full_state_from_spec(
+            self.spec,
+            inlet_node_pressure_pa=15.31e6,
+            resistance_boundary_pressure_pa=15.29e6,
+            config=SolveIVPConfig(t_start=0.0, t_stop=1.0e-5, max_step=1.0e-6),
+            t_eval=[0.0, 1.0e-5],
+        )
+
+        self.assertTrue(solution.success)
+        self.assertEqual(len(solution.t), 2)
+        self.assertEqual(len(solution.y), 132)
+        self.assertEqual([len(row) for row in solution.y], [2] * 132)
+
     def test_p4_closure_uses_structured_neighborhood(self) -> None:
         closure = self.system.pn3_p4_node_chamber_segment_closure_from_spec(
             self.spec,