SystemSimulationApp/app/simulation/solvers/stream.py

from __future__ import annotations

from dataclasses import dataclass

from app.simulation.core.base import Component, DynamicComponent
from app.simulation.core.ports import PortState
from app.simulation.performance import profile_phase
from app.simulation.systems.network import SimulationNetwork


class StreamSolveError(RuntimeError):
    def __init__(self, message: str, diagnostics: "StreamSolveDiagnostics") -> None:
        super().__init__(message)
        self.diagnostics = diagnostics


@dataclass(frozen=True)
class StreamSolveDiagnostics:
    converged: bool
    iterations: int
    max_delta: float

    def as_dict(self) -> dict[str, object]:
        return {
            "converged": self.converged,
            "iterations": self.iterations,
            "maxDelta": self.max_delta,
        }


@dataclass(frozen=True)
class _StreamConnectionBinding:
    component_name: str
    port_name: str
    connected_component: Component
    connected_port: PortState


class StreamResolver:
    """Resolve outflow enthalpy propagation after pressure and flow are known."""

    def __init__(
        self,
        network: SimulationNetwork,
        *,
        relative_tolerance: float = 1e-9,
        max_iterations: int = 100,
    ) -> None:
        self.network = network
        self.relative_tolerance = relative_tolerance
        self.max_iterations = max_iterations
        self._components = tuple(network.components.values())
        self._dynamic_components = tuple(
            component
            for component in self._components
            if isinstance(component, DynamicComponent)
        )
        self._non_dynamic_components = tuple(
            component
            for component in self._components
            if not isinstance(component, DynamicComponent)
        )
        # State ownership and pressure-flow stream sensitivity are independent
        # classifications. Compile this hook by behavior so algebraic
        # components such as PNL00R receive their upstream-temperature
        # references without dispatching a no-op to every component at runtime.
        self._flow_temperature_reference_components = tuple(
            component
            for component in self._components
            if type(component).update_flow_temperature_references
            is not Component.update_flow_temperature_references
        )
        self._ports = tuple(
            (component.name, port_name, port)
            for component in self._components
            for port_name, port in component.ports.items()
        )
        self._connection_bindings = self._build_connection_bindings()
        self.last_diagnostics: StreamSolveDiagnostics | None = None

    def _build_connection_bindings(self) -> tuple[_StreamConnectionBinding, ...]:
        result: list[_StreamConnectionBinding] = []
        for connection in self.network.connections:
            if connection.kind != "physical":
                continue
            first, second = connection.endpoints
            first_component = self.network.components[first.component]
            second_component = self.network.components[second.component]
            result.append(
                _StreamConnectionBinding(
                    component_name=first.component,
                    port_name=first.port,
                    connected_component=second_component,
                    connected_port=second_component.get_port(second.port),
                )
            )
            result.append(
                _StreamConnectionBinding(
                    component_name=second.component,
                    port_name=second.port,
                    connected_component=first_component,
                    connected_port=first_component.get_port(first.port),
                )
            )
        return tuple(result)

    def connected_enthalpies(self) -> dict[str, dict[str, float]]:
        values: dict[str, dict[str, float]] = {
            component.name: {} for component in self._components
        }
        for binding in self._connection_bindings:
            values[binding.component_name][binding.port_name] = (
                binding.connected_port.h_outflow
            )
        return values

    def connected_temperature_reference_enthalpies(
        self,
    ) -> dict[str, dict[str, float]]:
        """Return connector references used for upstream temperature only."""

        values: dict[str, dict[str, float]] = {
            component.name: {} for component in self._components
        }
        for binding in self._connection_bindings:
            values[binding.component_name][binding.port_name] = float(
                getattr(
                    binding.connected_component,
                    "temperature_reference_h",
                    binding.connected_port.h_outflow,
                )
            )
        return values

    @profile_phase("simulation.refresh", minimum_mode="audit")
    def refresh_flow_temperature_references(self) -> None:
        """Refresh pressure-flow property inputs without changing stream outflows."""

        connected = self.connected_temperature_reference_enthalpies()
        for component in self._flow_temperature_reference_components:
            component.update_flow_temperature_references(
                connected[component.name]
            )

    @profile_phase("simulation.refresh", minimum_mode="audit")
    def _refresh_dynamic_components(self) -> None:
        for component in self._dynamic_components:
            component.refresh_thermodynamic_ports()

    @profile_phase("simulation.refresh", minimum_mode="audit")
    def _refresh_stream_components(
        self,
        connected: dict[str, dict[str, float]],
    ) -> None:
        for component in self._non_dynamic_components:
            component.update_stream_outflows(connected[component.name])

    @profile_phase("simulation.stream", minimum_mode="audit")
    def solve(
        self,
        *,
        dynamic_ports_are_current: bool = False,
    ) -> tuple[StreamSolveDiagnostics, dict[str, dict[str, float]]]:
        if not dynamic_ports_are_current:
            self._refresh_dynamic_components()

        max_delta = 0.0
        for iteration in range(1, self.max_iterations + 1):
            previous = {
                (component_name, port_name): port.h_outflow
                for component_name, port_name, port in self._ports
            }
            connected = self.connected_enthalpies()
            self._refresh_stream_components(connected)

            deltas = [
                abs(port.h_outflow - previous[(component_name, port_name)])
                for component_name, port_name, port in self._ports
            ]
            magnitudes = [
                abs(port.h_outflow)
                for _component_name, _port_name, port in self._ports
            ]
            max_delta = max(deltas, default=0.0)
            scale = max(magnitudes + [1.0])
            if max_delta <= self.relative_tolerance * scale:
                diagnostics = StreamSolveDiagnostics(
                    converged=True,
                    iterations=iteration,
                    max_delta=max_delta,
                )
                self.last_diagnostics = diagnostics
                return diagnostics, self.connected_enthalpies()

        diagnostics = StreamSolveDiagnostics(
            converged=False,
            iterations=self.max_iterations,
            max_delta=max_delta,
        )
        self.last_diagnostics = diagnostics
        raise StreamSolveError(
            "Stream enthalpy propagation did not converge.",
            diagnostics,
        )