SystemSimulationApp/app/simulation/components/experimental/junctions/tee.py

from __future__ import annotations

from collections.abc import Mapping

from app.simulation.core.base import AlgebraicComponent
from app.simulation.core.catalog import ComponentDisplaySpec, PortDisplaySpec
from app.simulation.core.equations import EquationResidual
from app.simulation.core.medium import IdealGasMedium
from app.simulation.core.ports import PortDefinition


class Tee(AlgebraicComponent):
    """Python port of ModelicaModels.Mytee."""

    MODEL_TYPE = "tee"
    MODEL_VERSION = "1.0.0"
    PRESSURE_FLOW_DEPENDS_ON_STREAM = False
    PRESSURE_FLOW_EXACT_SUM_TO_ZERO_EQUATION_SUFFIXES = frozenset(
        ("mass_flow_balance",)
    )
    PORTS = (
        PortDefinition.pneumatic("port_in", nominal_role="bidirectional"),
        PortDefinition.pneumatic("port_out1", nominal_role="bidirectional"),
        PortDefinition.pneumatic("port_out2", nominal_role="bidirectional"),
    )
    PARAMETERS = ()
    RESULT_VARIABLES = ()
    DISPLAY = ComponentDisplaySpec(
        label="三通",
        library_id="experimental",
        category_id="junctions",
        symbol="tee",
        ports=(
            PortDisplaySpec("port_in", "left", order=10),
            PortDisplaySpec("port_out1", "right", order=20),
            PortDisplaySpec("port_out2", "right", order=30),
        ),
        order=50,
    )

    def __init__(self, name: str) -> None:
        super().__init__(name=name)
        self.set_parameter_values({})
        self.port_in = self.register_declared_port("port_in")
        self.port_out1 = self.register_declared_port("port_out1")
        self.port_out2 = self.register_declared_port("port_out2")

    @classmethod
    def create(
        cls,
        *,
        name: str,
        medium: IdealGasMedium,
        parameters: Mapping[str, float],
    ) -> Tee:
        return cls(name=name)

    def pressure_flow_equation_residuals(self) -> tuple[EquationResidual, ...]:
        return (
            EquationResidual(
                id=f"{self.name}:common_pressure_out1",
                owner="component",
                owner_id=self.name,
                relation="equal",
                variables=(f"{self.name}.port_in.p", f"{self.name}.port_out1.p"),
                role="effort",
                value=self.port_in.p - self.port_out1.p,
            ),
            EquationResidual(
                id=f"{self.name}:common_pressure_out2",
                owner="component",
                owner_id=self.name,
                relation="equal",
                variables=(f"{self.name}.port_in.p", f"{self.name}.port_out2.p"),
                role="effort",
                value=self.port_in.p - self.port_out2.p,
            ),
            EquationResidual(
                id=f"{self.name}:mass_flow_balance",
                owner="component",
                owner_id=self.name,
                relation="sumToZero",
                variables=(
                    f"{self.name}.port_in.m_flow",
                    f"{self.name}.port_out1.m_flow",
                    f"{self.name}.port_out2.m_flow",
                ),
                role="flow",
                value=(
                    self.port_in.m_flow
                    + self.port_out1.m_flow
                    + self.port_out2.m_flow
                ),
            ),
        )

    def update_stream_outflows(self, connected_h: Mapping[str, float]) -> None:
        incoming = [
            (port.m_flow, connected_h[name])
            for name, port in self.ports.items()
            if port.m_flow > 1e-12
        ]
        total_flow = sum(m_flow for m_flow, _ in incoming)
        if total_flow > 1e-12:
            mixed_h = sum(
                m_flow * enthalpy for m_flow, enthalpy in incoming
            ) / total_flow
        else:
            values = list(connected_h.values())
            mixed_h = sum(values) / len(values) if values else 0.0
        for port in self.ports.values():
            port.h_outflow = mixed_h

    def mixed_inlet_enthalpy(
        self,
        branch1_m_flow: float,
        branch1_h: float,
        branch2_m_flow: float,
        branch2_h: float,
        fallback_h: float = 0.0,
    ) -> float:
        positive_1 = max(branch1_m_flow, 0.0)
        positive_2 = max(branch2_m_flow, 0.0)
        total = positive_1 + positive_2
        if total <= 1e-9:
            return fallback_h
        return (positive_1 * branch1_h + positive_2 * branch2_h) / total

    def inlet_stream_enthalpy(
        self,
        branch1_m_flow: float,
        branch1_h: float,
        branch2_m_flow: float,
        branch2_h: float,
        fallback_h: float,
    ) -> float:
        """Approximate `inStream(port_in.h_outflow)` for the current tee topology."""

        return self.mixed_inlet_enthalpy(
            branch1_m_flow,
            branch1_h,
            branch2_m_flow,
            branch2_h,
            fallback_h=fallback_h,
        )

    def branch_actual_stream_enthalpy(
        self,
        branch_m_flow: float,
        branch_h: float,
        inlet_h: float,
    ) -> float:
        """Approximate `actualStream(branch.h_outflow)` for a tee branch port."""

        return inlet_h if branch_m_flow > 0.0 else branch_h

    @staticmethod
    def _solve_linear_2x2(
        a11: float,
        a12: float,
        a21: float,
        a22: float,
        b1: float,
        b2: float,
    ) -> tuple[float, float] | None:
        determinant = a11 * a22 - a12 * a21
        if abs(determinant) <= 1e-12:
            return None
        x1 = (b1 * a22 - b2 * a12) / determinant
        x2 = (a11 * b2 - a21 * b1) / determinant
        return x1, x2

    def solve_branch_outlet_flows_from_energy_balance(
        self,
        *,
        ratio_branch1: float,
        ratio_branch2: float,
        inlet_h_branch1: float,
        inlet_h_branch2: float,
        branch1_h: float,
        branch2_h: float,
        inlet_h: float,
        q_in_branch1: float,
        q_in_branch2: float,
        tolerance: float = 1e-12,
    ) -> tuple[float, float]:
        """Solve branch outlet flows for the current three-port downstream tee use-case."""

        rhs_branch1 = q_in_branch1 * inlet_h_branch1
        rhs_branch2 = q_in_branch2 * inlet_h_branch2

        def solve_both_forward() -> tuple[float, float] | None:
            return self._solve_linear_2x2(
                (1.0 + ratio_branch1) * branch1_h,
                ratio_branch1 * branch2_h,
                ratio_branch2 * branch1_h,
                (1.0 + ratio_branch2) * branch2_h,
                rhs_branch1,
                rhs_branch2,
            )

        def solve_one_reverse(
            *,
            branch1_reverse: bool,
        ) -> tuple[float, float] | None:
            if branch1_reverse:
                return self._solve_linear_2x2(
                    inlet_h * (1.0 + ratio_branch1),
                    ratio_branch1 * inlet_h,
                    ratio_branch2 * inlet_h,
                    branch2_h + ratio_branch2 * inlet_h,
                    rhs_branch1,
                    rhs_branch2,
                )

            return self._solve_linear_2x2(
                branch1_h + ratio_branch1 * inlet_h,
                ratio_branch1 * inlet_h,
                ratio_branch2 * inlet_h,
                inlet_h * (1.0 + ratio_branch2),
                rhs_branch1,
                rhs_branch2,
            )

        def solve_both_reverse() -> tuple[float, float] | None:
            return self._solve_linear_2x2(
                inlet_h * (1.0 + ratio_branch1),
                ratio_branch1 * inlet_h,
                ratio_branch2 * inlet_h,
                inlet_h * (1.0 + ratio_branch2),
                rhs_branch1,
                rhs_branch2,
            )

        candidate_solvers = (
            (
                solve_both_forward,
                lambda q1, q2: q1 >= -tolerance and q2 >= -tolerance,
            ),
            (
                lambda: solve_one_reverse(branch1_reverse=True),
                lambda q1, q2: q1 < -tolerance and q2 >= -tolerance and q1 + q2 > tolerance,
            ),
            (
                lambda: solve_one_reverse(branch1_reverse=True),
                lambda q1, q2: q1 < -tolerance and q2 >= -tolerance and q1 + q2 <= tolerance,
            ),
            (
                lambda: solve_one_reverse(branch1_reverse=False),
                lambda q1, q2: q2 < -tolerance and q1 >= -tolerance and q1 + q2 > tolerance,
            ),
            (
                lambda: solve_one_reverse(branch1_reverse=False),
                lambda q1, q2: q2 < -tolerance and q1 >= -tolerance and q1 + q2 <= tolerance,
            ),
            (
                solve_both_reverse,
                lambda q1, q2: q1 < -tolerance and q2 < -tolerance,
            ),
        )

        for solver, predicate in candidate_solvers:
            candidate = solver()
            if candidate is None:
                continue
            q_out_branch1, q_out_branch2 = candidate
            if predicate(q_out_branch1, q_out_branch2):
                return q_out_branch1, q_out_branch2

        return solve_both_forward() or (0.0, 0.0)