from pathlib import Path import os import unittest from unittest.mock import patch import numpy as np from app.main import compile_system_xml_network from app.simulation.solvers.jacobian import ExactColumnsUnavailable from app.simulation.solvers.mechanical import MechanicalConstraintGroup from app.simulation.solvers.solver import ODESolution, SolveIVPConfig from app.simulation.solvers.tangent import compile_three_piston_tangent_provider from app.simulation.systems.generic import ( ODE_JACOBIAN_MODE_ENVIRONMENT_VARIABLE, GenericFluidSystem, _requested_ode_jacobian_mode, ) from app.system_xml import validate_system_xml_document TARGET_XML = Path("tests/data/test_mql-full-branches-01-04.xml") TARGET_MASS_NAMES = ( "mass_friction_endstops_10", "mass_friction_endstops_11", "mass_friction_endstops_12", ) class ThreePistonTangentCompilerTests(unittest.TestCase): @classmethod def setUpClass(cls) -> None: report = validate_system_xml_document(TARGET_XML.read_bytes()) assert report.valid with patch.dict(os.environ, {"SIMULATION_CAUSAL_FAST_PATH": "1"}): cls.system = GenericFluidSystem( compile_system_xml_network(report.document) ) cls.initial_state = np.asarray( cls.system.consistent_initial_state_vector(0.0), dtype=float, ) def setUp(self) -> None: self.system.mechanical_state_reducer.reset_constraint_modes() self.system.apply_state_vector(self.initial_state.tolist()) def _closed_primal( self, state: np.ndarray, ) -> tuple[dict[str, dict[str, float]], np.ndarray]: self.system.apply_state_vector(state.tolist()) connected_h = self.system._close_current_state(0.0) derivative = np.asarray( self.system._state_derivatives(connected_h), dtype=float, ) return connected_h, derivative def test_target_layout_is_resolved_from_owner_and_slot(self) -> None: compilation = compile_three_piston_tangent_provider(self.system) self.assertTrue(compilation.eligible, compilation.reason) owner_slots: dict[str, tuple[int, int]] = {} cursor = 0 for entry in self.system.mechanical_state_reducer.state_entries: if isinstance(entry, MechanicalConstraintGroup): owner_slots[entry.representative.name] = (cursor, cursor + 1) cursor += 2 else: cursor += entry.state_size resolved = tuple( sorted( index for name in TARGET_MASS_NAMES for index in owner_slots[name] ) ) self.assertEqual(compilation.columns, resolved) # This is a fixture drift guard, not the provider's lookup mechanism. self.assertEqual(resolved, (20, 21, 38, 39, 54, 55)) self.assertEqual(compilation.reached_assignment_count, 34) def test_stale_primal_context_requests_typed_fallback(self) -> None: compilation = compile_three_piston_tangent_provider(self.system) self.assertTrue(compilation.eligible, compilation.reason) assert compilation.provider is not None with self.assertRaises(ExactColumnsUnavailable) as captured: compilation.provider( 0.0, self.initial_state.copy(), compilation.columns, ) self.assertEqual(captured.exception.reason, "stalePrimalContext") def test_initial_contact_boundary_requests_numerical_columns(self) -> None: compilation = compile_three_piston_tangent_provider(self.system) self.assertTrue(compilation.eligible, compilation.reason) assert compilation.provider is not None connected_h, _derivative = self._closed_primal( self.initial_state.copy() ) compilation.provider.request_primal_capture() compilation.provider.record_primal( 0.0, self.initial_state, connected_h, ) with self.assertRaises(ExactColumnsUnavailable) as captured: compilation.provider( 0.0, self.initial_state.copy(), compilation.columns, ) self.assertEqual( captured.exception.reason, "contactMode:contact_mode_boundary", ) def test_smooth_six_columns_match_full_rhs_centered_difference(self) -> None: compilation = compile_three_piston_tangent_provider(self.system) self.assertTrue(compilation.eligible, compilation.reason) provider = compilation.provider assert provider is not None state = self.initial_state.copy() # Move all three contacts away from gap==0 and all three PNL0001 laws # away from equal-pressure/zero-flow. The selected columns themselves # remain the six mechanical [v, x] seeds. for branch in provider.branches: chamber_offset, chamber_size = provider.state_offsets[ branch.chamber.name ] self.assertEqual(chamber_size, 2) state[chamber_offset] *= 1.01 state[branch.position_index] -= 1.0e-3 connected_h, _base = self._closed_primal(state) provider.request_primal_capture() provider.record_primal(0.0, state, connected_h) exact = provider(0.0, state.copy(), compilation.columns) numerical = np.empty_like(exact) for local_column, state_index in enumerate(compilation.columns): step = 1.0e-7 * max(abs(state[state_index]), 1.0) lower = state.copy() upper = state.copy() lower[state_index] -= step upper[state_index] += step lower_rhs = np.asarray( self.system.rhs(0.0, lower.tolist()), dtype=float, ) upper_rhs = np.asarray( self.system.rhs(0.0, upper.tolist()), dtype=float, ) numerical[:, local_column] = ( upper_rhs - lower_rhs ) / (2.0 * step) np.testing.assert_allclose( exact, numerical, rtol=2.0e-6, atol=1.0e-5, ) for downstream_mass in ( "mass_friction_endstops_18", "mass_friction_endstops_19", ): offset, size = provider.state_offsets[downstream_mass] self.assertEqual(size, 2) # Causal force reach includes both remote masses, but this fixture # holds their discrete constraints fixed, so final ODE rows are zero. np.testing.assert_allclose( exact[offset : offset + size, :], numerical[offset : offset + size, :], rtol=0.0, atol=1.0e-12, ) np.testing.assert_allclose( exact[offset : offset + size, :], 0.0, rtol=0.0, atol=1.0e-12, ) def test_disabled_causal_path_is_not_eligible(self) -> None: solver = self.system.pressure_flow_solver original = solver._causal_fast_path_environment_enabled try: solver._causal_fast_path_environment_enabled = False compilation = compile_three_piston_tangent_provider(self.system) finally: solver._causal_fast_path_environment_enabled = original self.assertFalse(compilation.eligible) self.assertEqual(compilation.reason, "causalFastPathDisabled") def test_semi_analytic_mode_is_explicit_opt_in(self) -> None: with patch.dict( os.environ, {ODE_JACOBIAN_MODE_ENVIRONMENT_VARIABLE: "semi-analytic"}, ): self.assertEqual( _requested_ode_jacobian_mode(), "semi-analytic", ) with patch.dict(os.environ, {}, clear=True): self.assertEqual(_requested_ode_jacobian_mode(), "scipy") def test_generic_simulation_wires_exact_columns_and_cleans_provider( self, ) -> None: report = validate_system_xml_document(TARGET_XML.read_bytes()) assert report.valid and report.document is not None with patch.dict(os.environ, {"SIMULATION_CAUSAL_FAST_PATH": "1"}): system = GenericFluidSystem( compile_system_xml_network(report.document) ) captured: dict[str, object] = {} def fake_integrate_ode(**options) -> ODESolution: captured.update(options) captured["activeProvider"] = system._ode_tangent_provider initial = [float(value) for value in options["initial_state"]] stop = float(options["config"].t_stop) return ODESolution( t=[0.0, stop], y=[[value, value] for value in initial], success=True, message="test", ) with patch.dict( os.environ, {ODE_JACOBIAN_MODE_ENVIRONMENT_VARIABLE: "semi-analytic"}, ), patch( "app.simulation.systems.generic.integrate_ode", side_effect=fake_integrate_ode, ): result = system.simulate( SolveIVPConfig(t_stop=1.0e-3, method="BDF"), sample_step=1.0e-3, ) jacobian = captured["jac"] self.assertIsNotNone(jacobian) self.assertIsNotNone(captured["activeProvider"]) self.assertIsNone(system._ode_tangent_provider) runtime = result.diagnostics["integration"]["jacobian"] self.assertEqual(runtime["mode"], "semiAnalyticExactColumns") self.assertEqual(runtime["effectiveMode"], "notEvaluated") self.assertEqual(runtime["originalColorGroupCount"], 31) # Generic wiring uses topology discovery rather than the legacy # three-name wrapper. This fixture contains one additional supported # branch, so all four branches are compiled automatically. self.assertEqual(runtime["remainingColorGroupCount"], 24) self.assertEqual(runtime["exactColumnCount"], 8) def test_generic_simulation_ineligible_path_uses_native_scipy(self) -> None: report = validate_system_xml_document(TARGET_XML.read_bytes()) assert report.valid and report.document is not None with patch.dict(os.environ, {"SIMULATION_CAUSAL_FAST_PATH": "1"}): system = GenericFluidSystem( compile_system_xml_network(report.document) ) system.pressure_flow_solver._causal_fast_path_environment_enabled = False captured: dict[str, object] = {} def fake_integrate_ode(**options) -> ODESolution: captured.update(options) captured["activeProvider"] = system._ode_tangent_provider initial = [float(value) for value in options["initial_state"]] stop = float(options["config"].t_stop) return ODESolution( t=[0.0, stop], y=[[value, value] for value in initial], success=True, message="test", ) with patch.dict( os.environ, {ODE_JACOBIAN_MODE_ENVIRONMENT_VARIABLE: "semi-analytic"}, ), patch( "app.simulation.systems.generic.integrate_ode", side_effect=fake_integrate_ode, ): result = system.simulate( SolveIVPConfig(t_stop=1.0e-3, method="BDF"), sample_step=1.0e-3, ) self.assertIsNone(captured["jac"]) self.assertIsNone(captured["activeProvider"]) self.assertIsNone(system._ode_tangent_provider) runtime = result.diagnostics["integration"]["jacobian"] self.assertEqual(runtime["mode"], "scipySparseFiniteDifference") self.assertEqual( runtime["fallbackReason"], "semiAnalytic:causalFastPathDisabled", ) def test_generic_simulation_cleans_provider_after_integration_error( self, ) -> None: report = validate_system_xml_document(TARGET_XML.read_bytes()) assert report.valid and report.document is not None with patch.dict(os.environ, {"SIMULATION_CAUSAL_FAST_PATH": "1"}): system = GenericFluidSystem( compile_system_xml_network(report.document) ) def fail_integration(**_options): self.assertIsNotNone(system._ode_tangent_provider) raise RuntimeError("injected integration failure") with patch.dict( os.environ, {ODE_JACOBIAN_MODE_ENVIRONMENT_VARIABLE: "semi-analytic"}, ), patch( "app.simulation.systems.generic.integrate_ode", side_effect=fail_integration, ), self.assertRaisesRegex(RuntimeError, "injected integration failure"): system.simulate( SolveIVPConfig(t_stop=1.0e-3, method="BDF"), sample_step=1.0e-3, ) self.assertIsNone(system._ode_tangent_provider) if __name__ == "__main__": unittest.main()