优化仿真求解性能并修复流量闭合问题（初版）

tests/test_thermofluid_closure_plan.py

@@ -0,0 +1,642 @@
+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()