loop1d.py

"""
loop1d.py — 1D thermal-hydraulic network model of a single pumped cooling loop.

Topology (series loop):

        heat in (chip load)          heat out (to secondary)
              |                              ^
              v                              |
    [pump] -> [cold plate] -> [heat exchanger] --+
      ^                                           |
      +-------------------------------------------+
                    one mass flow (series)

Two decoupled solves (constant properties => one-way coupling):
  1. Hydraulic : find the operating mass flow rate where pump head balances
                 the sum of component pressure losses.  Newton via fsolve.
                 (Scalar here; written so the same pattern extends to a
                  nodal network with parallel branches.)
  2. Thermal   : lumped-capacitance energy balances on a cold-plate solid
                 wall node + two fluid control volumes (cold plate, HX).
                 Transient: method-of-lines ODE system, stiff -> BDF.
                 Steady  : the same residual solved algebraically (fsolve),
                 cross-checked against a closed-form reference.

The heat exchanger uses an effectiveness-NTU model evaluated on the
instantaneous inlet (quasi-steady HX assumption: the exchanger responds
fast relative to loop transport time).

Extension hooks are marked with TODO:
  - add adiabatic/loss pipe control volumes between components
  - replace the scalar hydraulic solve with a nodal Newton-Raphson solve
    for parallel branches (rack manifolds, multiple cold plates)
  - wrap an outer Picard loop for temperature-dependent rho/mu/cp
  - distributed (multi-node) HX for fast transient fidelity

Author: Sahar Goudarzi
"""

from dataclasses import dataclass, field
import numpy as np
from scipy.optimize import fsolve
from scipy.integrate import solve_ivp
import pandas as pd


# --------------------------------------------------------------------------
# Working-fluid properties (constant for the skeleton; ~50/50 water-glycol)
# --------------------------------------------------------------------------
@dataclass
class Fluid:
    rho: float = 1050.0     # kg/m^3
    cp: float = 3400.0      # J/(kg.K)
    mu: float = 2.5e-3      # Pa.s   (kept for when pipe friction is computed)


# --------------------------------------------------------------------------
# Component parameters
# --------------------------------------------------------------------------
@dataclass
class Pump:
    """Quadratic pump curve: dp = dp0 * speed^2 - k * mdot^2  (affinity law on speed)."""
    dp0: float = 2.2e5      # Pa, shutoff head at speed = 1.0
    k: float = 1.0e6        # Pa / (kg/s)^2, curve steepness
    speed: float = 1.0      # 0..1 fractional speed (sweep variable)

    def dp(self, mdot: float) -> float:
        return self.dp0 * self.speed**2 - self.k * mdot**2


@dataclass
class FlowResistance:
    """Generic turbulent loss element: dp = K * mdot^2 (lumps f.L/D + minor losses)."""
    K: float

    def dp(self, mdot: float) -> float:
        return self.K * mdot**2


@dataclass
class ColdPlate:
    K: float = 3.0e6        # Pa/(kg/s)^2, hydraulic loss
    hA: float = 1800.0      # W/K, wall<->fluid conductance
    m_fluid: float = 0.05   # kg, fluid hold-up (thermal mass)
    m_wall: float = 0.40    # kg, base/wall mass
    cw: float = 900.0       # J/(kg.K), wall material (Al-ish)

    def dp(self, mdot: float) -> float:
        return self.K * mdot**2


@dataclass
class HeatExchanger:
    K: float = 4.0e6        # Pa/(kg/s)^2, hydraulic loss
    UA: float = 1500.0      # W/K, overall conductance
    m_fluid: float = 0.08   # kg, primary-side hold-up
    # Secondary (facility) side, treated as a fixed-inlet stream:
    mdot_sec: float = 0.30  # kg/s
    cp_sec: float = 3400.0  # J/(kg.K)

    def dp(self, mdot: float) -> float:
        return self.K * mdot**2

    def effectiveness(self, mdot: float, cp: float) -> tuple[float, float]:
        """Counterflow effectiveness-NTU. Returns (eps, Cmin)."""
        Ch = mdot * cp                 # hot (primary coolant) stream
        Cc = self.mdot_sec * self.cp_sec
        Cmin, Cmax = min(Ch, Cc), max(Ch, Cc)
        Cr = Cmin / Cmax
        NTU = self.UA / Cmin
        if abs(1.0 - Cr) < 1e-9:       # Cr -> 1 limit
            eps = NTU / (1.0 + NTU)
        else:
            e = np.exp(-NTU * (1.0 - Cr))
            eps = (1.0 - e) / (1.0 - Cr * e)
        return eps, Cmin


@dataclass
class Loop:
    fluid: Fluid = field(default_factory=Fluid)
    pump: Pump = field(default_factory=Pump)
    cold_plate: ColdPlate = field(default_factory=ColdPlate)
    hx: HeatExchanger = field(default_factory=HeatExchanger)
    pipe: FlowResistance = field(default_factory=lambda: FlowResistance(K=1.5e6))
    # Boundary conditions
    Q_chip: float = 1000.0  # W, heat load
    T_cold_in: float = 30.0 # degC, facility coolant inlet to HX


# ==========================================================================
# 1. HYDRAULIC SOLVE
# ==========================================================================
def hydraulic_residual(mdot: float, loop: Loop) -> float:
    """Series loop: pump head must equal the sum of all component losses.

    For a nodal network this scalar residual becomes a vector of node mass
    balances; swap fsolve(scalar) for fsolve(vector of nodal pressures).
    """
    mdot = float(np.asarray(mdot).flat[0])
    losses = (loop.cold_plate.dp(mdot)
              + loop.hx.dp(mdot)
              + loop.pipe.dp(mdot))
    return loop.pump.dp(mdot) - losses


def solve_hydraulic(loop: Loop, mdot_guess: float = 1.0) -> float:
    """Return the operating mass flow [kg/s]."""
    mdot = fsolve(hydraulic_residual, mdot_guess, args=(loop,), full_output=False)[0]
    return float(mdot)


def hydraulic_check(loop: Loop, mdot: float) -> float:
    """Closed-form reference for the single series loop (validation)."""
    Ksum = loop.cold_plate.K + loop.hx.K + loop.pipe.K + loop.pump.k
    return np.sqrt(loop.pump.dp0 * loop.pump.speed**2 / Ksum)


# ==========================================================================
# 2. THERMAL SOLVE
# ==========================================================================
# State vector y = [T_wall, T_fluid_coldplate, T_fluid_hx]   (all degC)
#   loop closure:  cold-plate inlet = HX outlet (T_fluid_hx)
#                  HX inlet         = cold-plate outlet (T_fluid_coldplate)

def thermal_rhs(t, y, loop: Loop, mdot: float):
    Tw, Tcp, Thx = y
    cp = loop.fluid.cp
    mc = mdot * cp

    # Cold-plate solid wall: chip heat in, convects to fluid
    dTw = (loop.Q_chip - loop.cold_plate.hA * (Tw - Tcp)) \
        / (loop.cold_plate.m_wall * loop.cold_plate.cw)

    # Cold-plate fluid CV: advect from HX outlet (upstream) + wall heat
    dTcp = (mc * (Thx - Tcp) + loop.cold_plate.hA * (Tw - Tcp)) \
        / (loop.cold_plate.m_fluid * cp)

    # HX fluid CV: advect from cold-plate outlet (upstream), reject Q_hx
    eps, Cmin = loop.hx.effectiveness(mdot, cp)
    Q_hx = eps * Cmin * (Tcp - loop.T_cold_in)   # quasi-steady eff-NTU
    dThx = (mc * (Tcp - Thx) - Q_hx) \
        / (loop.hx.m_fluid * cp)

    return [dTw, dTcp, dThx]


def solve_transient(loop: Loop, mdot: float, t_end: float = 600.0,
                    T0: float | None = None):
    """Integrate to steady state with a stiff solver (BDF)."""
    if T0 is None:
        T0 = [loop.T_cold_in, loop.T_cold_in, loop.T_cold_in]
    sol = solve_ivp(thermal_rhs, (0.0, t_end), T0, args=(loop, mdot),
                    method="BDF", rtol=1e-7, atol=1e-9, dense_output=True)
    return sol


def solve_steady(loop: Loop, mdot: float):
    """Solve the steady thermal system directly (residual = 0)."""
    def F(y):
        return thermal_rhs(0.0, y, loop, mdot)
    y0 = [loop.T_cold_in + 20, loop.T_cold_in + 10, loop.T_cold_in + 5]
    Tw, Tcp, Thx = fsolve(F, y0, args=())
    return Tw, Tcp, Thx


def steady_check(loop: Loop, mdot: float):
    """Closed-form steady reference (single loop) for validation.

    T_cp  = T_cold_in + Q/(eps*Cmin)
    T_hx  = T_cp - Q/(mdot*cp)
    T_w   = T_cp + Q/hA
    R_th  = (T_w - T_in_cp)/Q = 1/hA + 1/(mdot*cp)   [conv + caloric]
    """
    cp = loop.fluid.cp
    eps, Cmin = loop.hx.effectiveness(mdot, cp)
    Q = loop.Q_chip
    Tcp = loop.T_cold_in + Q / (eps * Cmin)
    Thx = Tcp - Q / (mdot * cp)
    Tw = Tcp + Q / loop.cold_plate.hA
    R_th = 1.0 / loop.cold_plate.hA + 1.0 / (mdot * cp)
    return Tw, Tcp, Thx, R_th


# ==========================================================================
# 3. PARAMETER SWEEP  ->  ROM / surrogate training data
# ==========================================================================
def sweep(loop_template: Loop,
          Q_chip_grid, speed_grid, T_cold_grid,
          out_csv: str = "rom_dataset.csv") -> pd.DataFrame:
    """Sweep loads, pump speeds, and coolant inlet temps; record steady metrics.

    Features the surrogate learns from : Q_chip, pump_speed, mdot, T_cold_in
    Targets it predicts                : R_th, T_wall (junction proxy), T_cp, eps
    """
    import copy
    rows = []
    for speed in speed_grid:
        for Q in Q_chip_grid:
            for Tc in T_cold_grid:
                loop = copy.deepcopy(loop_template)
                loop.pump.speed = speed
                loop.Q_chip = Q
                loop.T_cold_in = Tc

                mdot = solve_hydraulic(loop)
                Tw, Tcp, Thx = solve_steady(loop, mdot)
                eps, _ = loop.hx.effectiveness(mdot, loop.fluid.cp)
                R_th = (Tw - Thx) / Q          # (T_junction - T_inlet)/Q
                dp_pump = loop.pump.dp(mdot)

                rows.append(dict(
                    Q_chip=Q, pump_speed=speed, T_cold_in=Tc,
                    mdot=mdot, dp_pump=dp_pump, eps=eps,
                    T_wall=Tw, T_cp=Tcp, T_hx=Thx,
                    dT_fluid=Tcp - Thx, R_th=R_th,
                ))
    df = pd.DataFrame(rows)
    df.to_csv(out_csv, index=False)
    return df


# ==========================================================================
# DEMO / SELF-TEST
# ==========================================================================
if __name__ == "__main__":
    loop = Loop()

    # --- Hydraulic ---
    mdot = solve_hydraulic(loop)
    mdot_ref = hydraulic_check(loop, mdot)
    print(f"Hydraulic: mdot = {mdot:.4f} kg/s   (closed-form {mdot_ref:.4f}, "
          f"resid {hydraulic_residual(mdot, loop):.2e} Pa)")

    # --- Thermal (steady via fsolve, validated against closed form) ---
    Tw, Tcp, Thx = solve_steady(loop, mdot)
    Tw_r, Tcp_r, Thx_r, R_r = steady_check(loop, mdot)
    print(f"Thermal steady : T_wall={Tw:6.2f}  T_cp={Tcp:6.2f}  T_hx={Thx:6.2f} degC")
    print(f"Closed-form ref: T_wall={Tw_r:6.2f}  T_cp={Tcp_r:6.2f}  T_hx={Thx_r:6.2f} degC")
    print(f"Cold-plate thermal resistance R_th = {R_r*1e3:.3f} mK/W "
          f"(= 1/hA + 1/(mdot*cp))")

    # --- Energy-balance validation: rejected heat must equal chip load ---
    eps, Cmin = loop.hx.effectiveness(mdot, loop.fluid.cp)
    Q_hx = eps * Cmin * (Tcp - loop.T_cold_in)
    print(f"Energy balance : Q_chip={loop.Q_chip:.1f} W  Q_rejected={Q_hx:.1f} W "
          f"(closes to {100*abs(Q_hx-loop.Q_chip)/loop.Q_chip:.3f}% error)")

    # --- Transient should land on the steady solution ---
    sol = solve_transient(loop, mdot)
    Tw_t, Tcp_t, Thx_t = sol.y[:, -1]
    print(f"Transient t->ss: T_wall={Tw_t:6.2f}  T_cp={Tcp_t:6.2f}  T_hx={Thx_t:6.2f} degC")

    # --- Sweep -> ROM dataset ---
    df = sweep(
        loop,
        Q_chip_grid=np.linspace(250, 1000, 6),
        speed_grid=np.linspace(0.6, 1.0, 5),
        T_cold_grid=np.linspace(20, 40, 3),
        out_csv="rom_dataset.csv",
    )
    print(f"\nSweep complete: {len(df)} cases -> rom_dataset.csv")
    print(df.head().to_string(index=False))