from __future__ import annotations from collections.abc import Mapping from dataclasses import replace from types import SimpleNamespace import unittest from unittest.mock import patch from app.main import compile_reactflow_network, compile_system_xml_network from app.simulation.components.amesim.boundary.sources import AmesimPnpl01 from app.simulation.components.amesim.flow.orifices import ( AmesimPnor001, AmesimPnvo001FixedOpening, ) from app.simulation.components.amesim.flow.pipes import ( AmesimPnl00r, AmesimPnl0001, AmesimPnl0002, ) from app.simulation.components.amesim.storage.chambers import AmesimPnch023 from app.simulation.components.experimental.flow.orifice import Orifice from app.simulation.components.experimental.storage.cylinder import Cylinder from app.simulation.components.experimental.storage.tank import Tank from app.simulation.core.medium import IdealGasMedium from app.simulation.solvers.algebraic import ( AlgebraicSolveDiagnostics, AlgebraicSolveError, ) from app.simulation.solvers.algebraic_blocks import StreamBlockSolveResult from app.simulation.solvers.solver import SolveIVPConfig from app.simulation.systems.generic import ( GenericFluidSystem, simulation_preparation_issues, ) from app.simulation.systems.network import SimulationNetwork from app.system_xml import validate_system_xml_document from tests.test_amesim_pnvo001_signal_xml import high_pressure_helium_step_project from tests.test_amesim_mechanical_xml import elastic_contact_project from tests.test_generic_system_xml_simulation import chain_project from tests.test_high_stiffness_explicit_rk45 import short_explicit_rk45_xml class _UnclassifiedCustomOrifice(Orifice): """A custom subclass must not inherit the catalog purity declaration.""" def update_stream_outflows(self, connected_h: Mapping[str, float]) -> None: super().update_stream_outflows(connected_h) class _InvalidDependencyDeclarationOrifice(Orifice): PRESSURE_FLOW_DEPENDS_ON_STREAM = 1 class _CountingCylinder(Cylinder): def __init__(self, *args, **kwargs) -> None: self.refresh_count = 0 super().__init__(*args, **kwargs) def refresh_thermodynamic_ports(self): self.refresh_count += 1 return super().refresh_thermodynamic_ports() def _three_component_network( middle: Orifice, *, prefix: str = "independent", ) -> tuple[SimulationNetwork, _CountingCylinder, Tank]: medium = IdealGasMedium() source = _CountingCylinder( f"{prefix}_source", medium, V=0.02, p0=500_000.0, T0=320.0, ) sink = Tank( f"{prefix}_sink", medium, V=0.05, p0=100_000.0, T0=290.0, ) network = SimulationNetwork(prefix) for component in (source, middle, sink): network.add_component(component) network.connect(source.name, "port_b", middle.name, "port_a") network.connect(middle.name, "port_b", sink.name, "port_a") return network, source, sink def _mixed_island_network( *, independent_pressure: float = 450_000.0, ) -> SimulationNetwork: medium = IdealGasMedium() sensitive_source = Cylinder( "sensitive_source", medium, V=0.02, p0=600_000.0, T0=350.0, ) sensitive_sink = Tank( "sensitive_sink", medium, V=0.05, p0=100_000.0, T0=280.0, ) valve = AmesimPnvo001FixedOpening( "sensitive_valve", medium, area0=1.0e-5, opening=0.8, ) independent_source = Cylinder( "independent_source", medium, V=0.02, p0=independent_pressure, T0=310.0, ) independent_sink = Tank( "independent_sink", medium, V=0.05, p0=120_000.0, T0=295.0, ) independent_orifice = Orifice("independent_orifice", K=2.0e-5) network = SimulationNetwork("mixed-islands") for component in ( sensitive_source, sensitive_sink, valve, independent_source, independent_sink, independent_orifice, ): network.add_component(component) network.connect("sensitive_source", "port_b", "sensitive_valve", "port_2") network.connect("sensitive_valve", "port_3", "sensitive_sink", "port_a") network.connect("independent_source", "port_b", "independent_orifice", "port_a") network.connect("independent_orifice", "port_b", "independent_sink", "port_a") return network def _special_seed_islands_network() -> SimulationNetwork: medium = IdealGasMedium() series_source = AmesimPnch023("series_source", medium, p0=15.3e6) series_source_plug = AmesimPnpl01("series_source_plug") series_orifice = AmesimPnor001("series_orifice", medium) series_pipe = AmesimPnl0001("series_pipe", medium, p0=14.0e6) series_pipe_plug = AmesimPnpl01("series_pipe_plug") closed_pipe = AmesimPnl0002("closed_pipe", medium, p0=200_000.0) closed_pipe_left = AmesimPnpl01("closed_pipe_left") closed_pipe_right = AmesimPnpl01("closed_pipe_right") resistance_source = Cylinder( "resistance_source", medium, V=0.02, p0=500_000.0, T0=310.0, ) resistance = AmesimPnl00r("resistance", medium) resistance_plug = AmesimPnpl01("resistance_plug") network = SimulationNetwork("special-seed-islands") for component in ( series_source, series_source_plug, series_orifice, series_pipe, series_pipe_plug, closed_pipe, closed_pipe_left, closed_pipe_right, resistance_source, resistance, resistance_plug, ): network.add_component(component) network.connect("series_source_plug", "port_1", "series_source", "port_1") network.connect("series_source", "port_2", "series_orifice", "port_1") network.connect("series_orifice", "port_2", "series_pipe", "port_1") network.connect("series_pipe", "port_2", "series_pipe_plug", "port_1") network.connect("closed_pipe_left", "port_1", "closed_pipe", "port_1") network.connect("closed_pipe", "port_2", "closed_pipe_right", "port_1") network.connect("resistance_source", "port_b", "resistance", "port_1") network.connect("resistance", "port_2", "resistance_plug", "port_1") return network def _force_legacy_global_coupling(system: GenericFluidSystem) -> None: plan = system._thermofluid_closure_plan system._thermofluid_closure_plan = replace( plan, secondary_pressure_solvers=(system.pressure_flow_solver,), secondary_component_groups=(plan.global_component_group,), uses_conservative_global_solver=True, ) class ThermofluidClosurePlanTests(unittest.TestCase): def test_independent_network_solves_pressure_once_and_refreshes_once(self) -> None: network, source, _sink = _three_component_network( Orifice("independent_orifice", K=1.0e-5) ) system = GenericFluidSystem(network) source.refresh_count = 0 system.consistent_initial_state_vector() self.assertEqual(system._thermofluid_closure_plan.secondary_pressure_solvers, ()) self.assertEqual(system.algebraic_solve_count, 1) self.assertEqual(system.thermofluid_pressure_pass_count, 1) self.assertEqual(source.refresh_count, 1) def test_mixed_network_revisits_only_the_stream_sensitive_island(self) -> None: system = GenericFluidSystem(_mixed_island_network()) plan = system._thermofluid_closure_plan self.assertFalse(plan.uses_conservative_global_solver) self.assertEqual(len(plan.secondary_pressure_solvers), 1) self.assertEqual( set(plan.secondary_component_groups[0]), {"sensitive_source", "sensitive_sink", "sensitive_valve"}, ) self.assertEqual( set(plan.secondary_pressure_solvers[0].network.components), set(plan.secondary_component_groups[0]), ) self.assertNotIn( "independent_orifice", plan.secondary_pressure_solvers[0].network.components, ) state = system.consistent_initial_state_vector() first = system.rhs(0.0, state) second = system.rhs(0.0, state) for first_value, second_value in zip(first, second): self.assertAlmostEqual(first_value, second_value, delta=1.0e-9) def test_secondary_attempt_chain_is_counted_as_one_logical_solve(self) -> None: system = GenericFluidSystem(_mixed_island_network()) secondary = system._thermofluid_closure_plan.secondary_block_solvers[0] aggregate = AlgebraicSolveDiagnostics( success=True, message="accepted global fallback", evaluations=5, pressure_scale=600_000.0, flow_scale=0.01, max_scaled_residual=0.2, max_raw_residual=2.0, residual_evaluations=18, jacobian_mode="dense", block_fallback_used=True, block_fallback_reason="blockResidualNotConverged", ) fake_result = StreamBlockSolveResult( diagnostics=(aggregate,), scopes=(tuple(system.network.components),), used_global_fallback=True, ) initial_diagnostics: list[AlgebraicSolveDiagnostics] = [] original_initial_solve = system.pressure_flow_solver.solve def recorded_initial_solve(*args, **kwargs): result = original_initial_solve(*args, **kwargs) initial_diagnostics.append(result) return result with patch.object( system.pressure_flow_solver, "solve", side_effect=recorded_initial_solve, ), patch.object(secondary, "solve", return_value=fake_result): system.consistent_initial_state_vector() self.assertEqual(len(initial_diagnostics), 1) initial = initial_diagnostics[0] self.assertEqual(system.algebraic_solve_count, 2) self.assertEqual( system.algebraic_block_fallback_count, int(initial.block_fallback_used) + 1, ) self.assertEqual( system.algebraic_optimizer_evaluation_count, initial.evaluations + 5, ) self.assertEqual( system.algebraic_residual_evaluation_count, initial.residual_evaluations + 18, ) def test_secondary_island_failure_reports_its_physical_scope(self) -> None: system = GenericFluidSystem(_mixed_island_network()) plan = system._thermofluid_closure_plan secondary = plan.secondary_pressure_solvers[0] def failed_result(_fun, x0, **_kwargs): return SimpleNamespace( x=x0.copy(), success=False, status=-1, message="forced secondary-island failure", nfev=1, ) with patch.object( secondary, "_solve_explicit_flow_unknowns", return_value=None, ), patch("scipy.optimize.least_squares", side_effect=failed_result): with self.assertRaises(AlgebraicSolveError) as raised: secondary.solve(effort_variables=()) self.assertEqual(raised.exception.scope_kind, "physicalIsland") self.assertEqual( raised.exception.scope_components, plan.secondary_component_groups[0], ) def test_secondary_islands_preserve_special_pressure_seed_plans(self) -> None: network = _special_seed_islands_network() system = GenericFluidSystem(network) plan = system._thermofluid_closure_plan solvers = { frozenset(solver.network.components): solver for solver in plan.secondary_pressure_solvers } series_solver = solvers[ frozenset( { "series_source", "series_source_plug", "series_orifice", "series_pipe", "series_pipe_plug", } ) ] closed_pipe_solver = solvers[ frozenset( {"closed_pipe", "closed_pipe_left", "closed_pipe_right"} ) ] resistance_solver = solvers[ frozenset( { "resistance_source", "resistance", "resistance_plug", } ) ] self.assertEqual(len(series_solver._pnor_pnl0001_series_plan), 1) self.assertEqual(len(closed_pipe_solver._closed_resistance_pressure_plan), 2) self.assertEqual(len(resistance_solver._closed_resistance_pressure_plan), 1) closed_pipe = network.components["closed_pipe"] expected_pressure = closed_pipe.properties().p closed_pipe.port_1.p = 10_000.0 closed_pipe.port_2.p = 20_000.0 network.components["closed_pipe_left"].port_1.p = 30_000.0 network.components["closed_pipe_right"].port_1.p = 40_000.0 closed_pipe_solver._seed_closed_resistance_pressures() self.assertAlmostEqual(closed_pipe.port_1.p, expected_pressure) self.assertAlmostEqual(closed_pipe.port_2.p, expected_pressure) self.assertAlmostEqual( network.components["closed_pipe_left"].port_1.p, expected_pressure, ) self.assertAlmostEqual( network.components["closed_pipe_right"].port_1.p, expected_pressure, ) def test_unclassified_custom_stream_component_uses_legacy_global_solver(self) -> None: network, _source, _sink = _three_component_network( _UnclassifiedCustomOrifice("custom_orifice", K=1.0e-5), prefix="custom", ) system = GenericFluidSystem(network) plan = system._thermofluid_closure_plan self.assertTrue(plan.uses_conservative_global_solver) self.assertEqual(plan.secondary_pressure_solvers, (system.pressure_flow_solver,)) self.assertEqual(plan.secondary_component_groups, (plan.global_component_group,)) def test_invalid_dependency_declaration_uses_legacy_global_solver(self) -> None: network, _source, _sink = _three_component_network( _InvalidDependencyDeclarationOrifice("invalid_orifice", K=1.0e-5), prefix="invalid", ) plan = GenericFluidSystem(network)._thermofluid_closure_plan self.assertTrue(plan.uses_conservative_global_solver) self.assertEqual( plan.conservative_fallback_reason, "invalidDependencyDeclaration", ) def test_non_square_physical_island_metadata_forces_global_fallback(self) -> None: system = GenericFluidSystem(_mixed_island_network()) templates = list(system.pressure_flow_solver.equation_templates) moved = next( index for index, equation in enumerate(templates) if equation.owner == "component" and equation.owner_id == "independent_orifice" ) equation = templates[moved] templates[moved] = replace( equation, owner_id="sensitive_valve", variables=tuple( variable.replace("independent_orifice", "sensitive_valve") for variable in equation.variables ), ) system.pressure_flow_solver._equation_templates = tuple(templates) plan = system._build_thermofluid_closure_plan() self.assertTrue(plan.uses_conservative_global_solver) self.assertEqual( plan.conservative_fallback_reason, "nonSquarePhysicalIsland", ) def test_pruned_and_legacy_chain_results_are_numerically_equivalent(self) -> None: config = SolveIVPConfig( t_start=0.0, t_stop=0.01, method="BDF", max_step=0.001, ) optimized = GenericFluidSystem(compile_reactflow_network(chain_project())) legacy = GenericFluidSystem(compile_reactflow_network(chain_project())) _force_legacy_global_coupling(legacy) optimized_result = optimized.simulate(config, sample_step=0.005) legacy_result = legacy.simulate(config, sample_step=0.005) self.assertTrue(optimized_result.success) self.assertTrue(legacy_result.success) self.assertEqual(optimized_result.series.keys(), legacy_result.series.keys()) for key, optimized_values in optimized_result.series.items(): legacy_values = legacy_result.series[key] self.assertEqual(len(optimized_values), len(legacy_values), key) for optimized_value, legacy_value in zip( optimized_values, legacy_values, ): self.assertAlmostEqual( optimized_value, legacy_value, delta=1.0e-11 * max(abs(legacy_value), 1.0), msg=key, ) self.assertEqual(optimized_result.final.keys(), legacy_result.final.keys()) for key, optimized_value in optimized_result.final.items(): legacy_value = legacy_result.final[key] self.assertAlmostEqual( optimized_value, legacy_value, delta=1.0e-11 * max(abs(legacy_value), 1.0), msg=key, ) self.assertLess(optimized.algebraic_solve_count, legacy.algebraic_solve_count) def test_stream_dependent_rhs_is_history_independent_after_other_trial(self) -> None: network_a = compile_reactflow_network(high_pressure_helium_step_project()) network_fresh = compile_reactflow_network(high_pressure_helium_step_project()) system = GenericFluidSystem(network_a) fresh = GenericFluidSystem(network_fresh) state = system.consistent_initial_state_vector() fresh_state = fresh.consistent_initial_state_vector() perturbed = list(state) perturbed[0] *= 1.000001 perturbed[1] *= 0.999999 first = system.rhs(0.041, state) system.rhs(0.041, perturbed) repeated = system.rhs(0.041, state) reference = fresh.rhs(0.041, fresh_state) for expected, actual in zip(first, repeated): self.assertAlmostEqual(actual, expected, delta=1.0e-10 * max(abs(expected), 1.0)) for expected, actual in zip(reference, repeated): self.assertAlmostEqual(actual, expected, delta=1.0e-10 * max(abs(expected), 1.0)) def test_block_solve_reuses_global_scales_from_extreme_other_island(self) -> None: optimized = GenericFluidSystem( _mixed_island_network(independent_pressure=1.0e10) ) legacy = GenericFluidSystem( _mixed_island_network(independent_pressure=1.0e10) ) _force_legacy_global_coupling(legacy) optimized_state = optimized.initial_state_vector() legacy_state = legacy.initial_state_vector() optimized_rhs = optimized.rhs(0.0, optimized_state) legacy_rhs = legacy.rhs(0.0, legacy_state) for expected, actual in zip(legacy_rhs, optimized_rhs): self.assertAlmostEqual( actual, expected, delta=1.0e-10 * max(abs(expected), 1.0), ) self.assertLessEqual( optimized.max_algebraic_residual, max(legacy.max_algebraic_residual, 1.0e-14), ) def test_high_stiffness_contact_island_matches_legacy_global_closure(self) -> None: report = validate_system_xml_document(short_explicit_rk45_xml()) self.assertTrue(report.valid) assert report.document is not None document = report.document optimized = GenericFluidSystem(compile_system_xml_network(document)) legacy = GenericFluidSystem(compile_system_xml_network(document)) _force_legacy_global_coupling(legacy) config = SolveIVPConfig( t_start=document.simulation.t_start, t_stop=document.simulation.t_stop, method=document.simulation.method, rtol=1.0e-6, max_step=document.simulation.max_step, ) optimized_result = optimized.simulate( config, sample_step=document.simulation.sample_step, ) legacy_result = legacy.simulate( config, sample_step=document.simulation.sample_step, ) self.assertTrue(optimized_result.success) self.assertTrue(legacy_result.success) optimized_totals = optimized_result.diagnostics["integration"]["totals"] legacy_totals = legacy_result.diagnostics["integration"]["totals"] self.assertEqual( optimized_totals["stateTransitionCount"], legacy_totals["stateTransitionCount"], ) self.assertEqual(optimized_result.series.keys(), legacy_result.series.keys()) for key, optimized_values in optimized_result.series.items(): legacy_values = legacy_result.series[key] self.assertEqual(len(optimized_values), len(legacy_values), key) for optimized_value, legacy_value in zip( optimized_values, legacy_values, ): self.assertAlmostEqual( optimized_value, legacy_value, delta=2.0e-7 * max(abs(legacy_value), 1.0), msg=key, ) def test_secondary_fluid_island_does_not_clear_active_contact_state(self) -> None: network = compile_reactflow_network(elastic_contact_project()) medium = IdealGasMedium() source = Cylinder( "separate_source", medium, V=0.02, p0=600_000.0, T0=350.0, ) sink = Tank( "separate_sink", medium, V=0.05, p0=100_000.0, T0=280.0, ) valve = AmesimPnvo001FixedOpening( "separate_valve", medium, area0=1.0e-5, opening=0.8, ) for component in (source, sink, valve): network.add_component(component) network.connect("separate_source", "port_b", "separate_valve", "port_2") network.connect("separate_valve", "port_3", "separate_sink", "port_a") self.assertEqual(simulation_preparation_issues(network), ()) system = GenericFluidSystem(network) contact = network.components["contact_1"] secondary = system._thermofluid_closure_plan.secondary_block_solvers[0] original_solve = secondary.solve observed: list[tuple[float | None, ...]] = [] def checked_solve(*, scale_context=None): if contact._causal_penetration is None: contact.set_causal_contact(penetration=1.0e-4, force=10.0) before = ( contact._causal_penetration, contact._causal_contact_force, contact._causal_port_1_x, contact._causal_port_2_x, contact._causal_port_1_v, contact._causal_port_2_v, ) result = original_solve(scale_context=scale_context) after = ( contact._causal_penetration, contact._causal_contact_force, contact._causal_port_1_x, contact._causal_port_2_x, contact._causal_port_1_v, contact._causal_port_2_v, ) self.assertEqual(after, before) observed.append(before) return result secondary.solve = checked_solve system.consistent_initial_state_vector() self.assertTrue(observed) self.assertIsNotNone(observed[0][0]) self.assertIsNotNone(observed[0][1]) if __name__ == "__main__": unittest.main()