from pathlib import Path import os import unittest from unittest.mock import patch import numpy as np from scipy.optimize._numdiff import group_columns from app.main import compile_system_xml_network from app.simulation.solvers.jacobian import ( ExactColumnsUnavailable, SparseSecantJacobian, ) import app.simulation.solvers.tangent as tangent_module from app.simulation.solvers.tangent import ( compile_supported_piston_tangent_provider, ) from app.simulation.systems.generic import GenericFluidSystem from app.simulation.systems.network import Endpoint from app.system_xml import validate_system_xml_document TARGET_XML = Path("tests/data/test-mql-8.xml") EXPECTED_COLUMNS = tuple(range(104, 120)) EXPECTED_BRANCH_NAMES = ( ( "amesim_mecmas21_1", "amesim_pnrp17_1", "amesim_pnch012_15", "amesim_pnl0001_13", "amesim_lstp00a_1", ), ( "amesim_mecmas21_2", "amesim_pnrp17_2", "amesim_pnch012_14", "amesim_pnl0001_14", "amesim_lstp00a_2", ), ( "amesim_mecmas21_3", "amesim_pnrp17_3", "amesim_pnch012_13", "amesim_pnl0001_15", "amesim_lstp00a_3", ), ( "amesim_mecmas21_4", "amesim_pnrp17_4", "amesim_pnch012_12", "amesim_pnl0001_16", "amesim_lstp00a_4", ), ( "amesim_mecmas21_5", "amesim_pnrp17_5", "amesim_pnch012_11", "amesim_pnl0001_17", "amesim_lstp00a_5", ), ( "amesim_mecmas21_6", "amesim_pnrp17_6", "amesim_pnch012_10", "amesim_pnl0001_18", "amesim_lstp00a_6", ), ( "amesim_mecmas21_7", "amesim_pnrp17_7", "amesim_pnch012_9", "amesim_pnl0001_19", "amesim_lstp00a_7", ), ( "amesim_mecmas21_8", "amesim_pnrp17_8", "amesim_pnch012_8", "amesim_pnl0001_20", "amesim_lstp00a_8", ), ) class SupportedPistonTangentCompilerTests(unittest.TestCase): @classmethod def setUpClass(cls) -> 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"}): 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_discovers_all_branches_without_legacy_names_or_offsets(self) -> None: with patch.object( tangent_module, "_TARGET_BRANCH_NAMES", (("not", "a", "real", "branch", "name"),), ): compilation = compile_supported_piston_tangent_provider( self.system ) self.assertTrue(compilation.eligible, compilation.reason) self.assertEqual(compilation.columns, EXPECTED_COLUMNS) self.assertEqual(compilation.reached_assignment_count, 84) provider = compilation.provider assert provider is not None self.assertEqual(len(provider.branches), 8) self.assertEqual( tuple( ( branch.mass.name, branch.piston.name, branch.chamber.name, branch.pipe.name, branch.contact.name, ) for branch in provider.branches ), EXPECTED_BRANCH_NAMES, ) self.assertEqual( {branch.chamber_connection_port for branch in provider.branches}, {"port_2"}, ) def test_exact_columns_reduce_seed_zero_pattern_to_36_colors(self) -> None: compilation = compile_supported_piston_tangent_provider(self.system) self.assertTrue(compilation.eligible, compilation.reason) provider = compilation.provider assert provider is not None pattern = self.system.jacobian_sparsity().tocsc().astype(bool) columns_only = pattern.tolil(copy=True) columns_only[:, list(compilation.columns)] = False columns_only = columns_only.tocsc() columns_only.eliminate_zeros() active_columns = np.flatnonzero( np.asarray(columns_only.getnnz(axis=0)).reshape(-1) > 0 ) groups = group_columns( columns_only[:, active_columns], order=0, ) self.assertEqual(pattern.nnz, 3280) self.assertEqual(columns_only.nnz, 2448) self.assertEqual(len(active_columns), 116) self.assertEqual(int(groups.max(initial=-1)) + 1, 36) jacobian = SparseSecantJacobian( lambda _time, state: np.zeros_like(state), pattern, 1.0e-8, exact_rows=self.system._exact_ode_jacobian_rows(), exact_columns=(compilation.columns, provider), max_consecutive_reuses=0, ) diagnostics = jacobian.diagnostics() self.assertEqual(diagnostics["originalColorGroupCount"], 52) self.assertEqual(diagnostics["remainingColorGroupCount"], 36) self.assertEqual(diagnostics["exactColumnCount"], 16) self.assertEqual(diagnostics["finiteDifferenceColumnCount"], 116) self.assertEqual( jacobian._remaining_finite_difference_sparsity.nnz, 2160, ) def test_partial_unsupported_branch_fails_with_stable_reason(self) -> None: piston_endpoint = Endpoint("amesim_pnrp17_1", "port_5") connection_index, connection = next( (index, connection) for index, connection in enumerate(self.system.network.connections) if piston_endpoint in connection.endpoints ) del self.system.network.connections[connection_index] try: compilation = compile_supported_piston_tangent_provider( self.system ) finally: self.system.network.connections.insert( connection_index, connection, ) self.assertFalse(compilation.eligible) self.assertEqual( compilation.reason, "unsupportedPistonBranchTopology:contact", ) def test_initial_contact_boundary_requests_full_numeric_columns(self) -> None: compilation = compile_supported_piston_tangent_provider(self.system) self.assertTrue(compilation.eligible, compilation.reason) provider = compilation.provider assert provider is not None connected_h, _derivative = self._closed_primal( self.initial_state.copy() ) provider.request_primal_capture() provider.record_primal( 0.0, self.initial_state, connected_h, ) with self.assertRaises(ExactColumnsUnavailable) as captured: provider( 0.0, self.initial_state.copy(), compilation.columns, ) self.assertEqual( captured.exception.reason, "contactMode:contact_mode_boundary", ) def test_smooth_sixteen_columns_match_centered_full_rhs(self) -> None: compilation = compile_supported_piston_tangent_provider(self.system) self.assertTrue(compilation.eligible, compilation.reason) provider = compilation.provider assert provider is not None state = self.initial_state.copy() 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=3.0e-6, atol=2.0e-5, ) if __name__ == "__main__": unittest.main()