from __future__ import annotations from dataclasses import replace from pathlib import Path import unittest from app.main import compile_reactflow_network, compile_system_xml_network from app.simulation.solvers.causal_ir import ( CausalIROpcode, compile_causal_numeric_ir, ) from app.simulation.systems.generic import GenericFluidSystem from app.system_xml import validate_system_xml_document from tests.test_amesim_mechanical_xml import zero_force_mass_project from tests.test_amesim_pnl0001_xml import amesim_pnl0001_project from tests.test_amesim_pnvo001_signal_xml import high_pressure_helium_step_project from tests.test_generic_system_xml_simulation import chain_project TARGET_XML = Path("tests/data/test-mql-8.xml") class _SyntheticDirectReaderError(Exception): pass def _system(project) -> GenericFluidSystem: return GenericFluidSystem(compile_reactflow_network(project)) class CausalNumericIRTests(unittest.TestCase): def _assert_direct_reader_exception_is_controlled( self, exception: Exception, ) -> None: system = _system(amesim_pnl0001_project()) solver = system.pressure_flow_solver compilation = compile_causal_numeric_ir(solver) self.assertTrue(compilation.supported, compilation.fallback_reason) assert compilation.ir is not None ir = compilation.ir direct_evaluation = next( evaluation for stage in ir.program.effort_stages for evaluation in stage.evaluations if evaluation.opcode is CausalIROpcode.EFFORT_DIRECT_RESIDUAL ) evaluators = list(ir.bindings.evaluators) victim = solver.unknowns[-1] before = tuple(unknown.read() for unknown in solver.unknowns) def failing_reader() -> float: victim.write(victim.read() + 123.0) raise exception evaluators[direct_evaluation.evaluator_slot] = failing_reader failing_ir = replace( ir, bindings=replace(ir.bindings, evaluators=tuple(evaluators)), ) result = failing_ir.execute( failing_ir.create_workspace(), transactional=True, ) self.assertFalse(result.success) self.assertEqual( result.fallback_reason, f"effortEvaluationFailed:{type(exception).__name__}", ) self.assertTrue(result.rolled_back) self.assertEqual( tuple(unknown.read() for unknown in solver.unknowns), before, ) def test_direct_reader_key_error_returns_controlled_failure(self) -> None: self._assert_direct_reader_exception_is_controlled( KeyError("synthetic direct-reader key failure") ) def test_direct_reader_custom_exception_returns_controlled_failure( self, ) -> None: self._assert_direct_reader_exception_is_controlled( _SyntheticDirectReaderError("synthetic direct-reader failure") ) def test_direct_reader_memory_error_remains_fatal(self) -> None: system = _system(amesim_pnl0001_project()) compilation = compile_causal_numeric_ir(system.pressure_flow_solver) self.assertTrue(compilation.supported, compilation.fallback_reason) assert compilation.ir is not None ir = compilation.ir direct_evaluation = next( evaluation for stage in ir.program.effort_stages for evaluation in stage.evaluations if evaluation.opcode is CausalIROpcode.EFFORT_DIRECT_RESIDUAL ) evaluators = list(ir.bindings.evaluators) def failing_reader() -> float: raise MemoryError("synthetic direct-reader allocation failure") evaluators[direct_evaluation.evaluator_slot] = failing_reader failing_ir = replace( ir, bindings=replace(ir.bindings, evaluators=tuple(evaluators)), ) with self.assertRaisesRegex(MemoryError, "allocation failure"): failing_ir.execute( failing_ir.create_workspace(), transactional=True, ) def test_structure_signature_is_stable_and_excludes_bindings(self) -> None: first = compile_causal_numeric_ir( _system(zero_force_mass_project()).pressure_flow_solver ) second = compile_causal_numeric_ir( _system(zero_force_mass_project()).pressure_flow_solver ) self.assertTrue(first.supported) self.assertTrue(second.supported) assert first.ir is not None and second.ir is not None self.assertEqual( first.ir.program.structural_signature, second.ir.program.structural_signature, ) self.assertEqual( first.ir.program.structural_signature, first.ir.program.calculate_structural_signature(), ) self.assertEqual(len(first.ir.program.structural_signature), 64) changed_binding = replace( first.ir, bindings=replace( first.ir.bindings, evaluators=tuple( (lambda evaluate=evaluate: evaluate()) for evaluate in first.ir.bindings.evaluators ), ), ) self.assertEqual( changed_binding.program.structural_signature, first.ir.program.structural_signature, ) def test_reference_interpreter_matches_existing_object_plan_bitwise(self) -> None: system = _system(high_pressure_helium_step_project()) state = system.initial_state_vector() system.rhs(0.041, state) solver = system.pressure_flow_solver compilation = compile_causal_numeric_ir(solver) self.assertTrue(compilation.supported, compilation.fallback_reason) assert compilation.ir is not None before = tuple(unknown.read() for unknown in solver.unknowns) self.assertTrue(solver._execute_causal_effort_plan(("p",))) for unknown in solver._explicit_flow_unknowns_by_variables[ frozenset(("f", "m_flow")) ]: unknown.write(0.0) expected_flow_stages: list[tuple[float, ...]] = [] for stage in solver._explicit_flow_plan: values = solver._evaluate_explicit_flow_stage(stage) expected_flow_stages.append(values) for assignment, value in zip(stage.assignments, values): assignment.unknown.write(value) expected = tuple(unknown.read() for unknown in solver.unknowns) for unknown, value in zip(solver.unknowns, before): unknown.write(value) workspace = compilation.ir.create_workspace() observed: list[tuple[str, int, tuple[int, ...], tuple[float, ...]]] = [] result = compilation.ir.execute( workspace, stage_observer=lambda phase, index, slots, values: observed.append( (phase, index, slots, values) ), ) actual = tuple(unknown.read() for unknown in solver.unknowns) self.assertTrue(result.success, result.fallback_reason) self.assertEqual(actual, expected) self.assertEqual( result.effort_assignment_count, len(compilation.ir.program.effort_stages[0].operations), ) self.assertEqual( result.flow_assignment_count, compilation.ir.program.flow_assignment_count, ) self.assertEqual(observed[0][0], "effort:p") self.assertEqual( sum(item[0] == "flow" for item in observed), len(compilation.ir.program.flow_stages), ) self.assertEqual( tuple(item[3] for item in observed if item[0] == "flow"), tuple(expected_flow_stages), ) def test_workspace_is_reused_and_normal_execution_does_not_snapshot(self) -> None: system = _system(zero_force_mass_project()) system.rhs(0.0, system.initial_state_vector()) compilation = compile_causal_numeric_ir(system.pressure_flow_solver) assert compilation.ir is not None workspace = compilation.ir.create_workspace() array_ids = ( id(workspace.canonical_values), id(workspace.effort_residuals), id(workspace.flow_values), id(workspace.transaction_values), ) workspace.transaction_values.fill(float("nan")) first = compilation.ir.execute(workspace) second = compilation.ir.execute(workspace) self.assertTrue(first.success) self.assertTrue(second.success) self.assertEqual( array_ids, ( id(workspace.canonical_values), id(workspace.effort_residuals), id(workspace.flow_values), id(workspace.transaction_values), ), ) self.assertTrue( all(value != value for value in workspace.transaction_values) ) def test_transactional_audit_rolls_back_partial_flow_failure(self) -> None: system = _system(zero_force_mass_project()) system.rhs(0.0, system.initial_state_vector()) solver = system.pressure_flow_solver compilation = compile_causal_numeric_ir(solver) assert compilation.ir is not None ir = compilation.ir before = tuple(unknown.read() for unknown in solver.unknowns) operation = ir.program.flow_stages[0].operations[0] evaluators = list(ir.bindings.evaluators) original = evaluators[operation.evaluator_slot] if operation.opcode is CausalIROpcode.FLOW_DIRECT: evaluators[operation.evaluator_slot] = lambda: float("nan") else: equation_index = operation.equation_indices[0] def one_nonfinite_component(): values = list(original()) values[equation_index] = float("nan") return tuple(values) evaluators[operation.evaluator_slot] = one_nonfinite_component failing_ir = replace( ir, bindings=replace(ir.bindings, evaluators=tuple(evaluators)), ) result = failing_ir.execute( failing_ir.create_workspace(), transactional=True, ) self.assertFalse(result.success) self.assertEqual(result.fallback_reason, "nonFiniteFlowAssignment") self.assertTrue(result.rolled_back) self.assertEqual( tuple(unknown.read() for unknown in solver.unknowns), before, ) def test_unsupported_plan_returns_existing_proof_reason(self) -> None: solver = _system(chain_project()).pressure_flow_solver compilation = compile_causal_numeric_ir(solver) self.assertFalse(compilation.supported) self.assertIsNone(compilation.ir) self.assertEqual( compilation.fallback_reason, solver._causal_fast_path_fallback_reason, ) class TargetCausalNumericIRStructureTests(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 cls.system = GenericFluidSystem( compile_system_xml_network(report.document) ) cls.compilation = compile_causal_numeric_ir( cls.system.pressure_flow_solver ) def test_target_has_expected_canonical_and_compatibility_coordinates( self, ) -> None: self.assertTrue( self.compilation.supported, self.compilation.fallback_reason, ) assert self.compilation.ir is not None program = self.compilation.ir.program self.assertEqual(len(program.compatibility_slots), 776) self.assertEqual(len(program.canonical_slots), 452) self.assertEqual(program.effort_group_count, 116) self.assertEqual(program.flow_assignment_count, 336) self.assertEqual(program.effort_scatter_count, 440) self.assertEqual(program.eliminated_effort_replica_count, 324) self.assertEqual( tuple(slot.slot for slot in program.canonical_slots), tuple(range(452)), ) def test_target_batches_efforts_and_preserves_flow_stage_layout(self) -> None: assert self.compilation.ir is not None program = self.compilation.ir.program pressure_stage = next( stage for stage in program.effort_stages if stage.variable == "p" ) self.assertEqual(len(pressure_stage.operations), 76) self.assertEqual(len(pressure_stage.evaluations), 40) self.assertEqual( sum( evaluation.opcode is CausalIROpcode.EFFORT_DIRECT_RESIDUAL for evaluation in pressure_stage.evaluations ), 20, ) self.assertEqual( sum( evaluation.opcode is CausalIROpcode.EFFORT_COMPONENT_RESIDUAL for evaluation in pressure_stage.evaluations ), 20, ) self.assertEqual( self.system.pressure_flow_solver.causal_execution_diagnostics()[ "directEffortAnchorCount" ], 20, ) self.assertEqual( [len(stage.target_slots) for stage in program.flow_stages], [114, 134, 49, 33, 5, 1], ) def test_target_preserves_exact_direct_sum_assignment_counts(self) -> None: solver = self.system.pressure_flow_solver self.assertEqual( solver.causal_execution_diagnostics()[ "directSumFlowAssignmentCount" ], 38, ) component_equation_ids = frozenset( equation.id for evaluation in solver._component_equation_plan for equation in evaluation.templates ) secondary = ( self.system._thermofluid_closure_plan.secondary_block_solvers[0] ) secondary_direct_count = sum( assignment.evaluate is not None and assignment.equation_id in component_equation_ids for stage in secondary._selected_explicit_flow_plan for assignment in stage.assignments ) self.assertEqual(secondary_direct_count, 36) def test_target_reference_execution_matches_all_compatibility_slots(self) -> None: assert self.compilation.ir is not None solver = self.system.pressure_flow_solver self.system.rhs(0.0, self.system.initial_state_vector()) before = tuple(unknown.read() for unknown in solver.unknowns) self.assertTrue(solver._execute_causal_effort_plan(("p",))) self.assertIsNone(solver._execute_compiled_causal_flow_plan()) expected = tuple(unknown.read() for unknown in solver.unknowns) for unknown, value in zip(solver.unknowns, before): unknown.write(value) result = self.compilation.ir.execute( self.compilation.ir.create_workspace() ) self.assertTrue(result.success, result.fallback_reason) self.assertEqual( tuple(unknown.read() for unknown in solver.unknowns), expected, ) if __name__ == "__main__": unittest.main()