from __future__ import annotations from types import SimpleNamespace import unittest from unittest.mock import patch import numpy as np from app.simulation.components.amesim.boundary.sources import AmesimPnpl01 from app.simulation.components.experimental.storage.cylinder import Cylinder from app.simulation.core.base import AlgebraicComponent from app.simulation.core.equations import EquationResidual from app.simulation.core.medium import IdealGasMedium from app.simulation.core.ports import PortDefinition from app.simulation.solvers.algebraic import PressureFlowSolver from app.simulation.systems.network import SimulationNetwork class _IncompleteDependencyComponent(AlgebraicComponent): """Deliberately violate the residual dependency authoring contract.""" PORTS = (PortDefinition.pneumatic("port_1", nominal_role="bidirectional"),) def __init__(self, name: str) -> None: super().__init__(name) self.port_1 = self.register_declared_port("port_1") self.port_1.p = 90_000.0 self.port_1.m_flow = 1.0 def pressure_flow_equation_values(self) -> tuple[float, ...]: return self.port_1.p - 100_000.0, self.port_1.m_flow def pressure_flow_equation_residuals(self) -> tuple[EquationResidual, ...]: return ( EquationResidual( id=f"{self.name}:pressure", owner="component", owner_id=self.name, relation="state", # The live residual reads p, but the declaration omits it on # purpose. A zero structural row must disable sparse FD. variables=(), role="effort", value=self.port_1.p - 100_000.0, ), EquationResidual( id=f"{self.name}:flow", owner="component", owner_id=self.name, relation="constitutive", variables=(f"{self.name}.port_1.m_flow",), role="flow", value=self.port_1.m_flow, ), ) def _storage_boundary_solver(name: str) -> PressureFlowSolver: medium = IdealGasMedium() storage = Cylinder(f"{name}_storage", medium, V=0.1, p0=500_000.0) boundary = AmesimPnpl01(f"{name}_boundary") network = SimulationNetwork(name) network.add_component(storage) network.add_component(boundary) network.connect(storage.name, "port_b", boundary.name, "port_1") storage.refresh_thermodynamic_ports() boundary.port_1.m_flow = 1.0 return PressureFlowSolver(network) class PressureFlowSolverSparsityTests(unittest.TestCase): def test_sparse_and_dense_paths_produce_the_same_solution(self) -> None: sparse_solver = _storage_boundary_solver("sparse") dense_solver = _storage_boundary_solver("dense") dense_solver._jacobian_sparsity_is_trusted = False with patch.object( sparse_solver, "_solve_explicit_flow_unknowns", return_value=set(), ): sparse = sparse_solver.solve() with patch.object( dense_solver, "_solve_explicit_flow_unknowns", return_value=set(), ): dense = dense_solver.solve() self.assertEqual(sparse.jacobian_mode, "blockSparse") self.assertEqual(sparse.nonlinear_block_count, 1) self.assertLess( sparse.nonlinear_block_unknown_count, len(sparse_solver.unknowns), ) self.assertEqual(dense.jacobian_mode, "dense") self.assertFalse(sparse.dense_fallback_used) self.assertFalse(dense.dense_fallback_used) self.assertGreater(sparse.residual_evaluations, sparse.evaluations) self.assertLess(sparse.residual_evaluations, dense.residual_evaluations) sparse_values = { unknown.id.split("_", maxsplit=1)[-1]: unknown.read() for unknown in sparse_solver.unknowns } dense_values = { unknown.id.split("_", maxsplit=1)[-1]: unknown.read() for unknown in dense_solver.unknowns } self.assertEqual(sparse_values.keys(), dense_values.keys()) for variable_id in sparse_values: self.assertAlmostEqual( sparse_values[variable_id], dense_values[variable_id], delta=1.0e-9, ) self.assertEqual( sparse.as_dict()["residualEvaluations"], sparse.residual_evaluations, ) def test_incomplete_dependency_metadata_uses_dense_finite_differences( self, ) -> None: component = _IncompleteDependencyComponent("incomplete") network = SimulationNetwork("incomplete-dependency") network.add_component(component) solver = PressureFlowSolver(network) diagnostics = solver.solve() self.assertFalse(solver.jacobian_sparsity_is_trusted) self.assertEqual( solver.jacobian_sparsity_fallback_reason, "equationWithoutDeclaredUnknown", ) self.assertEqual(diagnostics.jacobian_mode, "dense") self.assertFalse(diagnostics.dense_fallback_used) self.assertAlmostEqual(component.port_1.p, 100_000.0, places=6) self.assertAlmostEqual(component.port_1.m_flow, 0.0, places=12) def test_failed_sparse_candidate_restarts_dense_from_original_x0(self) -> None: solver = _storage_boundary_solver("retry") # This regression specifically exercises the final whole-network # sparse -> dense compatibility fallback, not the preceding block # optimization. solver._equation_blocks_are_trusted = False original_x0: list[np.ndarray] = [] calls: list[str] = [] flow_indices = tuple( index for index, unknown in enumerate(solver.unknowns) if unknown.variable == "m_flow" ) def fake_least_squares(_fun, x0, **kwargs): if "jac_sparsity" in kwargs: calls.append("sparse") original_x0.append(x0.copy()) failed = x0.copy() failed[list(flow_indices)] = 123.0 return SimpleNamespace( x=failed, success=False, status=-1, message="forced sparse failure", nfev=3, ) calls.append("dense") np.testing.assert_array_equal(x0, original_x0[0]) solved = x0.copy() solved[list(flow_indices)] = 0.0 return SimpleNamespace( x=solved, success=True, status=1, message="dense fallback solved", nfev=4, ) with patch.object( solver, "_solve_explicit_flow_unknowns", return_value=set(), ), patch( "scipy.optimize.least_squares", side_effect=fake_least_squares, ): diagnostics = solver.solve() self.assertEqual(calls, ["sparse", "dense"]) self.assertTrue(diagnostics.success) self.assertEqual(diagnostics.evaluations, 7) self.assertEqual(diagnostics.jacobian_mode, "sparseThenDense") self.assertTrue(diagnostics.dense_fallback_used) for index in flow_indices: self.assertAlmostEqual(solver.unknowns[index].read(), 0.0, places=12) def test_compiled_pattern_matches_declared_storage_boundary_structure( self, ) -> None: solver = _storage_boundary_solver("pattern") self.assertTrue(solver.jacobian_sparsity_is_trusted) self.assertEqual(solver.jacobian_sparsity.shape, (4, 4)) self.assertEqual(solver.jacobian_sparsity.nnz, 6) if __name__ == "__main__": unittest.main()