"""Executable reference IR for compile-proven causal algebraic programs. The IR eliminates duplicate *logical* effort coordinates, but intentionally keeps a compatibility scatter map to existing ``PortState`` objects. Stream propagation, derivatives, and result collection still consume those objects; this is a reference for a future flat backend, not physical slot deletion. """ from __future__ import annotations from collections.abc import Callable, Iterable from dataclasses import dataclass, replace from enum import StrEnum from hashlib import sha256 import json from math import isfinite from typing import TYPE_CHECKING, Any if TYPE_CHECKING: import numpy as np CAUSAL_NUMERIC_IR_SCHEMA_VERSION = 1 PRESSURE_LOWER_BOUND_PA = 0.0 class CausalIROpcode(StrEnum): EFFORT_BROADCAST = "effort_broadcast" EFFORT_DIRECT_RESIDUAL = "effort_direct_residual" EFFORT_COMPONENT_RESIDUAL = "effort_component_residual" FLOW_DIRECT = "flow_direct" FLOW_COMPONENT_RESIDUAL = "flow_component_residual" @dataclass(frozen=True, slots=True) class CausalIRCompatibilitySlot: slot: int id: str variable: str @dataclass(frozen=True, slots=True) class CausalIRCanonicalSlot: slot: int id: str variable: str kind: str @dataclass(frozen=True, slots=True) class CausalIREffortOperation: opcode: CausalIROpcode variable: str result_slot: int anchor_compatibility_slot: int scatter_compatibility_slots: tuple[int, ...] equation_id: str @dataclass(frozen=True, slots=True) class CausalIREffortEvaluation: opcode: CausalIROpcode output_indices: tuple[int, ...] equation_indices: tuple[int, ...] equation_ids: tuple[str, ...] evaluator_slot: int @dataclass(frozen=True, slots=True) class CausalIREffortStage: variable: str operations: tuple[CausalIREffortOperation, ...] evaluations: tuple[CausalIREffortEvaluation, ...] @dataclass(frozen=True, slots=True) class CausalIRFlowOperation: opcode: CausalIROpcode output_indices: tuple[int, ...] equation_indices: tuple[int, ...] equation_ids: tuple[str, ...] evaluator_slot: int @dataclass(frozen=True, slots=True) class CausalIRFlowStage: target_slots: tuple[int, ...] scatter_compatibility_slots: tuple[int, ...] equation_ids: tuple[str, ...] operations: tuple[CausalIRFlowOperation, ...] @dataclass(frozen=True, slots=True) class CausalIRProgram: """Immutable callback-free structure used as the backend cache key.""" schema_version: int canonical_slots: tuple[CausalIRCanonicalSlot, ...] compatibility_slots: tuple[CausalIRCompatibilitySlot, ...] reset_compatibility_slots: tuple[int, ...] external_effort_compatibility_slots: tuple[int, ...] effort_stages: tuple[CausalIREffortStage, ...] flow_stages: tuple[CausalIRFlowStage, ...] structural_signature: str @property def assignment_count(self) -> int: return len(self.canonical_slots) @property def effort_group_count(self) -> int: return sum(len(stage.operations) for stage in self.effort_stages) @property def flow_assignment_count(self) -> int: return sum(len(stage.target_slots) for stage in self.flow_stages) @property def effort_scatter_count(self) -> int: return sum( len(operation.scatter_compatibility_slots) for stage in self.effort_stages for operation in stage.operations ) @property def eliminated_effort_replica_count(self) -> int: return self.effort_scatter_count - self.effort_group_count @property def maximum_effort_stage_width(self) -> int: return max((len(stage.operations) for stage in self.effort_stages), default=0) @property def maximum_flow_stage_width(self) -> int: return max((len(stage.target_slots) for stage in self.flow_stages), default=0) def structural_dict(self) -> dict[str, object]: return { "schemaVersion": self.schema_version, "canonicalSlots": [ { "slot": item.slot, "id": item.id, "variable": item.variable, "kind": item.kind, } for item in self.canonical_slots ], "compatibilitySlots": [ {"slot": item.slot, "id": item.id, "variable": item.variable} for item in self.compatibility_slots ], "resetCompatibilitySlots": list(self.reset_compatibility_slots), "externalEffortCompatibilitySlots": list( self.external_effort_compatibility_slots ), "effortStages": [ { "variable": stage.variable, "operations": [ { "opcode": operation.opcode.value, "resultSlot": operation.result_slot, "anchorCompatibilitySlot": ( operation.anchor_compatibility_slot ), "scatterCompatibilitySlots": list( operation.scatter_compatibility_slots ), "equationId": operation.equation_id, } for operation in stage.operations ], "evaluations": [ { "opcode": evaluation.opcode.value, "outputIndices": list(evaluation.output_indices), "equationIndices": list(evaluation.equation_indices), "equationIds": list(evaluation.equation_ids), "evaluatorSlot": evaluation.evaluator_slot, } for evaluation in stage.evaluations ], } for stage in self.effort_stages ], "flowStages": [ { "targetSlots": list(stage.target_slots), "scatterCompatibilitySlots": list( stage.scatter_compatibility_slots ), "equationIds": list(stage.equation_ids), "operations": [ { "opcode": operation.opcode.value, "outputIndices": list(operation.output_indices), "equationIndices": list(operation.equation_indices), "equationIds": list(operation.equation_ids), "evaluatorSlot": operation.evaluator_slot, } for operation in stage.operations ], } for stage in self.flow_stages ], } def calculate_structural_signature(self) -> str: payload = json.dumps( self.structural_dict(), ensure_ascii=True, separators=(",", ":"), sort_keys=True, ).encode("utf-8") return sha256(payload).hexdigest() @dataclass(frozen=True, slots=True) class CausalIRBindings: readers: tuple[Callable[[], float], ...] writers: tuple[Callable[[float], None], ...] evaluators: tuple[Callable[[], object], ...] @dataclass(slots=True) class CausalIRWorkspace: structural_signature: str canonical_values: "np.ndarray[Any, Any]" effort_residuals: "np.ndarray[Any, Any]" effort_written: "np.ndarray[Any, Any]" flow_values: "np.ndarray[Any, Any]" flow_written: "np.ndarray[Any, Any]" transaction_values: "np.ndarray[Any, Any]" @dataclass(frozen=True, slots=True) class CausalIRExecutionResult: success: bool fallback_reason: str | None structural_signature: str effort_assignment_count: int flow_assignment_count: int completed_effort_stage_count: int completed_flow_stage_count: int rolled_back: bool StageObserver = Callable[ [str, int, tuple[int, ...], tuple[float, ...]], None, ] @dataclass(frozen=True, slots=True) class CausalNumericIR: """Bound reference IR; its normal path performs no full snapshot.""" program: CausalIRProgram bindings: CausalIRBindings def create_workspace(self) -> CausalIRWorkspace: try: import numpy as np except ImportError as exc: # pragma: no cover raise RuntimeError("The causal numeric reference IR requires NumPy.") from exc return CausalIRWorkspace( structural_signature=self.program.structural_signature, canonical_values=np.empty( max(len(self.program.canonical_slots), 1), dtype=np.float64 ), effort_residuals=np.empty( max(self.program.maximum_effort_stage_width, 1), dtype=np.float64 ), effort_written=np.empty( max(self.program.maximum_effort_stage_width, 1), dtype=np.bool_ ), flow_values=np.empty( max(self.program.maximum_flow_stage_width, 1), dtype=np.float64 ), flow_written=np.empty( max(self.program.maximum_flow_stage_width, 1), dtype=np.bool_ ), transaction_values=np.empty( max(len(self.program.compatibility_slots), 1), dtype=np.float64 ), ) def execute( self, workspace: CausalIRWorkspace, *, effort_variables: tuple[str, ...] = ("p",), transactional: bool = False, stage_observer: StageObserver | None = None, ) -> CausalIRExecutionResult: """Interpret the IR; transactional snapshots are audit-only.""" program = self.program bindings = self.bindings signature = program.structural_signature if workspace.structural_signature != signature: raise ValueError("Causal IR workspace belongs to a different program.") if len(bindings.readers) != len(program.compatibility_slots) or len( bindings.writers ) != len(program.compatibility_slots): raise ValueError("Causal IR compatibility binding count is inconsistent.") if any(variable not in {"p", "x", "v"} for variable in effort_variables): return CausalIRExecutionResult( False, "unsupportedEffortVariable", signature, 0, 0, 0, 0, False ) snapshot_count = 0 if transactional: try: for slot, reader in enumerate(bindings.readers): workspace.transaction_values[slot] = float(reader()) snapshot_count += 1 except MemoryError: raise except (ArithmeticError, RuntimeError, TypeError, ValueError) as exc: return CausalIRExecutionResult( False, f"slotReadFailed:{type(exc).__name__}", signature, 0, 0, 0, 0, False, ) effort_count = 0 flow_count = 0 completed_effort_stages = 0 completed_flow_stages = 0 def failed(reason: str) -> CausalIRExecutionResult: rolled_back = False if transactional: for slot in range(snapshot_count): bindings.writers[slot](float(workspace.transaction_values[slot])) rolled_back = True return CausalIRExecutionResult( False, reason, signature, effort_count, flow_count, completed_effort_stages, completed_flow_stages, rolled_back, ) selected_efforts = frozenset(effort_variables) for stage_index, stage in enumerate(program.effort_stages): if stage.variable not in selected_efforts: continue width = len(stage.operations) workspace.effort_written[:width] = False for evaluation in stage.evaluations: try: evaluated = bindings.evaluators[evaluation.evaluator_slot]() if evaluation.opcode is CausalIROpcode.EFFORT_DIRECT_RESIDUAL: output = evaluation.output_indices[0] workspace.effort_residuals[output] = float(evaluated) workspace.effort_written[output] = True continue if not hasattr(evaluated, "__len__"): raise TypeError("component evaluator returned no sequence") for output, equation in zip( evaluation.output_indices, evaluation.equation_indices ): if equation >= len(evaluated): raise IndexError("component equation disappeared") workspace.effort_residuals[output] = float(evaluated[equation]) workspace.effort_written[output] = True except MemoryError: raise except ( ArithmeticError, IndexError, RuntimeError, TypeError, ValueError, ) as exc: return failed(f"effortEvaluationFailed:{type(exc).__name__}") if any(not bool(workspace.effort_written[index]) for index in range(width)): return failed("effortEvaluationCoverageMismatch") for output, operation in enumerate(stage.operations): try: anchor = float(bindings.readers[operation.anchor_compatibility_slot]()) target = anchor - float(workspace.effort_residuals[output]) except MemoryError: raise except ( ArithmeticError, IndexError, RuntimeError, TypeError, ValueError, ) as exc: return failed(f"effortAssignmentFailed:{type(exc).__name__}") if not isfinite(target) or ( stage.variable == "p" and target <= PRESSURE_LOWER_BOUND_PA ): return failed("nonFiniteOrInvalidEffortAssignment") workspace.canonical_values[operation.result_slot] = target for slot in operation.scatter_compatibility_slots: bindings.writers[slot](target) effort_count += 1 completed_effort_stages += 1 if stage_observer is not None: try: stage_observer( f"effort:{stage.variable}", stage_index, tuple(item.result_slot for item in stage.operations), tuple( float(workspace.canonical_values[item.result_slot]) for item in stage.operations ), ) except MemoryError: raise except Exception as exc: return failed(f"stageObserverFailed:{type(exc).__name__}") try: external_finite = all( isfinite(float(bindings.readers[slot]())) for slot in program.external_effort_compatibility_slots ) except MemoryError: raise except (ArithmeticError, RuntimeError, TypeError, ValueError) as exc: return failed(f"externalEffortReadFailed:{type(exc).__name__}") if not external_finite: return failed("nonFiniteExternalEffort") for slot in program.reset_compatibility_slots: bindings.writers[slot](0.0) for stage_index, stage in enumerate(program.flow_stages): width = len(stage.target_slots) workspace.flow_written[:width] = False for operation in stage.operations: try: evaluated = bindings.evaluators[operation.evaluator_slot]() if operation.opcode is CausalIROpcode.FLOW_DIRECT: output = operation.output_indices[0] workspace.flow_values[output] = float(evaluated) workspace.flow_written[output] = True continue if not hasattr(evaluated, "__len__"): raise TypeError("component evaluator returned no sequence") for output, equation in zip( operation.output_indices, operation.equation_indices ): if equation >= len(evaluated): raise IndexError("component equation disappeared") # Targets are zero before the stage; preserve -residual. workspace.flow_values[output] = -float(evaluated[equation]) workspace.flow_written[output] = True except MemoryError: raise except ( ArithmeticError, IndexError, RuntimeError, TypeError, ValueError, ) as exc: return failed(f"flowEvaluationFailed:{type(exc).__name__}") if any(not bool(workspace.flow_written[index]) for index in range(width)): return failed("flowAssignmentCoverageMismatch") for output, (canonical, compatibility) in enumerate( zip(stage.target_slots, stage.scatter_compatibility_slots) ): target = float(workspace.flow_values[output]) if not isfinite(target): return failed("nonFiniteFlowAssignment") workspace.canonical_values[canonical] = target bindings.writers[compatibility](target) flow_count += 1 completed_flow_stages += 1 if stage_observer is not None: try: stage_observer( "flow", stage_index, stage.target_slots, tuple(float(workspace.flow_values[i]) for i in range(width)), ) except MemoryError: raise except Exception as exc: return failed(f"stageObserverFailed:{type(exc).__name__}") return CausalIRExecutionResult( True, None, signature, effort_count, flow_count, completed_effort_stages, completed_flow_stages, False, ) @dataclass(frozen=True, slots=True) class CausalIRCompilation: ir: CausalNumericIR | None fallback_reason: str | None @property def supported(self) -> bool: return self.ir is not None and self.fallback_reason is None def _unsupported(reason: str) -> CausalIRCompilation: return CausalIRCompilation(ir=None, fallback_reason=reason) def _unique_slots(items: Iterable[int]) -> tuple[int, ...]: return tuple(dict.fromkeys(int(item) for item in items)) def _compile_effort_evaluations( operations: tuple[CausalIREffortOperation, ...], anchor_evaluators: tuple[Callable[[], float], ...], component_locations: dict[ str, tuple[object, Callable[[], tuple[float, ...]], int] ], evaluators: list[Callable[[], object]], ) -> tuple[CausalIREffortEvaluation, ...]: grouped: dict[int, list[tuple[int, int, str]]] = {} component_callbacks: dict[int, Callable[[], tuple[float, ...]]] = {} direct: list[tuple[int, Callable[[], float], str]] = [] for output, (operation, anchor_evaluate) in enumerate( zip(operations, anchor_evaluators) ): location = component_locations.get(operation.equation_id) if location is None: direct.append((output, anchor_evaluate, operation.equation_id)) continue owner, evaluate, equation = location key = id(owner) component_callbacks[key] = evaluate grouped.setdefault(key, []).append((output, equation, operation.equation_id)) compiled: list[CausalIREffortEvaluation] = [] for output, evaluate, equation_id in direct: evaluator = len(evaluators) evaluators.append(evaluate) compiled.append( CausalIREffortEvaluation( CausalIROpcode.EFFORT_DIRECT_RESIDUAL, (output,), (), (equation_id,), evaluator, ) ) for key, entries in grouped.items(): evaluator = len(evaluators) evaluators.append(component_callbacks[key]) compiled.append( CausalIREffortEvaluation( CausalIROpcode.EFFORT_COMPONENT_RESIDUAL, tuple(item[0] for item in entries), tuple(item[1] for item in entries), tuple(item[2] for item in entries), evaluator, ) ) return tuple(compiled) def compile_causal_numeric_ir(solver: object) -> CausalIRCompilation: """Lower a compile-proven global plan; unsupported plans fail closed.""" if not bool(getattr(solver, "_causal_fast_path_eligible", False)): return _unsupported( str( getattr(solver, "_causal_fast_path_fallback_reason", None) or "causalProofNotAvailable" ) ) try: unknowns = tuple(getattr(solver, "unknowns")) effort_plan = getattr(solver, "_causal_effort_plan_by_variable") flow_plan = tuple(getattr(solver, "_explicit_flow_plan")) component_plan = tuple(getattr(solver, "_component_equation_plan")) reset_unknowns = tuple( getattr(solver, "_explicit_flow_unknowns_by_variables")[ frozenset(("f", "m_flow")) ] ) external_unknowns = tuple( getattr(solver, "_causal_external_effort_unknowns") ) except (AttributeError, KeyError, TypeError): return _unsupported("unsupportedCausalSolverContract") unknown_ids = tuple(str(item.id) for item in unknowns) if len(set(unknown_ids)) != len(unknown_ids): return _unsupported("duplicateAlgebraicUnknown") compatibility_slot_by_id = { unknown_id: slot for slot, unknown_id in enumerate(unknown_ids) } compatibility_slots = tuple( CausalIRCompatibilitySlot(slot, unknown_id, str(unknown.variable)) for slot, (unknown_id, unknown) in enumerate(zip(unknown_ids, unknowns)) ) readers = tuple(item.read for item in unknowns) writers = tuple(item.write for item in unknowns) evaluators: list[Callable[[], object]] = [] canonical_slots: list[CausalIRCanonicalSlot] = [] component_locations: dict[ str, tuple[object, Callable[[], tuple[float, ...]], int] ] = {} try: for plan in component_plan: for equation, template in enumerate(plan.templates): component_locations[str(template.id)] = ( plan.component, plan.evaluate, equation, ) except (AttributeError, TypeError): return _unsupported("unsupportedComponentEvaluationContract") effort_stages: list[CausalIREffortStage] = [] try: for variable in ("p", "x", "v"): operations: list[CausalIREffortOperation] = [] anchors: list[Callable[[], float]] = [] for assignment in effort_plan[variable]: result = len(canonical_slots) equation_id = str(assignment.anchor.equation_id) scatter = tuple( compatibility_slot_by_id[item.id] for item in assignment.members ) if not scatter or len(set(scatter)) != len(scatter): return _unsupported("invalidEffortScatterSlots") canonical_slots.append( CausalIRCanonicalSlot( result, f"effort:{variable}:{equation_id}", variable, "effort_group", ) ) operations.append( CausalIREffortOperation( CausalIROpcode.EFFORT_BROADCAST, variable, result, compatibility_slot_by_id[assignment.anchor.unknown.id], scatter, equation_id, ) ) anchors.append(assignment.anchor.evaluate) operation_tuple = tuple(operations) effort_stages.append( CausalIREffortStage( variable, operation_tuple, _compile_effort_evaluations( operation_tuple, tuple(anchors), component_locations, evaluators, ), ) ) except (AttributeError, KeyError, TypeError): return _unsupported("unsupportedEffortPlanContract") flow_stages: list[CausalIRFlowStage] = [] try: for stage in flow_plan: scatter = tuple( compatibility_slot_by_id[item.unknown.id] for item in stage.assignments ) equation_ids = tuple(str(item.equation_id) for item in stage.assignments) if len(set(scatter)) != len(scatter): return _unsupported("duplicateFlowTargetInStage") targets: list[int] = [] for assignment in stage.assignments: target = len(canonical_slots) targets.append(target) canonical_slots.append( CausalIRCanonicalSlot( target, f"flow:{assignment.unknown.id}", str(assignment.unknown.variable), "flow_assignment", ) ) covered: list[int] = [] operations: list[CausalIRFlowOperation] = [] for output, evaluate in stage.direct_evaluations: output = int(output) evaluator = len(evaluators) evaluators.append(evaluate) operations.append( CausalIRFlowOperation( CausalIROpcode.FLOW_DIRECT, (output,), (), (equation_ids[output],), evaluator, ) ) covered.append(output) for evaluation in stage.component_evaluations: evaluator = len(evaluators) evaluators.append(evaluation.evaluate) outputs = tuple(int(item) for item in evaluation.assignment_indices) operations.append( CausalIRFlowOperation( CausalIROpcode.FLOW_COMPONENT_RESIDUAL, outputs, tuple(int(item) for item in evaluation.equation_indices), tuple(str(item) for item in evaluation.equation_ids), evaluator, ) ) covered.extend(outputs) if sorted(covered) != list(range(len(scatter))): return _unsupported("flowStageEvaluationCoverageMismatch") flow_stages.append( CausalIRFlowStage( tuple(targets), scatter, equation_ids, tuple(operations) ) ) except (AttributeError, IndexError, KeyError, TypeError): return _unsupported("unsupportedFlowPlanContract") try: reset_slots = _unique_slots( compatibility_slot_by_id[item.id] for item in reset_unknowns ) external_slots = _unique_slots( compatibility_slot_by_id[item.id] for item in external_unknowns ) except (AttributeError, KeyError): return _unsupported("unknownCausalBoundarySlot") flow_scatter = tuple( item for stage in flow_stages for item in stage.scatter_compatibility_slots ) if len(set(flow_scatter)) != len(flow_scatter): return _unsupported("duplicateExplicitFlowAssignment") if set(flow_scatter) != set(reset_slots): return _unsupported("incompleteExplicitFlowCoverage") program = CausalIRProgram( CAUSAL_NUMERIC_IR_SCHEMA_VERSION, tuple(canonical_slots), compatibility_slots, reset_slots, external_slots, tuple(effort_stages), tuple(flow_stages), "", ) program = replace( program, structural_signature=program.calculate_structural_signature() ) return CausalIRCompilation( CausalNumericIR( program, CausalIRBindings(readers, writers, tuple(evaluators)), ), None, )