cryosim/optimization/sweep.py

"""1D and 2D parameter sweeps for hardware design exploration."""

from dataclasses import dataclass
from typing import List, Optional

import numpy as np
import pandas as pd

from cryosim.calibration.params import (
    CALIBRATION_PARAMS, PARAM_NAMES, get_nominal_values, apply_overrides,
)
from cryosim.hardware.config import load_config
from cryosim.engine.fast import ICV_open


@dataclass
class SweepResult:
    """Result of a 1D parameter sweep."""
    param_name: str
    param_values: np.ndarray
    metric_name: str
    metric_values: np.ndarray
    hardware: str
    Pexit: float
    speed: float


def _extract_metric(out: np.ndarray, hist: dict, metric: str) -> float:
    """Extract a named metric from engine output."""
    if metric == "mdot_kgpm":
        return hist["mdot_kgpm"]
    elif metric == "mass_eff":
        return float(out[1, 0])
    elif metric == "Tc_peak_K":
        return float(np.max(hist["Tc_K"]))
    elif metric == "kWh_extend":
        return hist["kWh_extend"]
    else:
        return hist.get(metric, np.nan)


def sweep_1d(
    hardware: str,
    param_name: str,
    values: List[float],
    speed: float = 0.65,
    Pexit: float = 350.0,
    metric: str = "mdot_kgpm",
    Ptank: float = 7.0,
    Psat: float = 2.0,
) -> SweepResult:
    """Sweep one parameter across a range of values.

    Args:
        hardware: Base config name.
        param_name: One of the 7 calibratable param names.
        values: List of values to sweep.
        speed, Pexit: Operating conditions.
        metric: Which metric to track ("mdot_kgpm", "mass_eff", "Tc_peak_K", "kWh_extend").
        Ptank: Inlet tank pressure [barg].
        Psat: Saturation pressure [barg].

    Returns:
        SweepResult with param_values and metric_values arrays.
    """
    if param_name not in CALIBRATION_PARAMS:
        raise ValueError(f"Unknown param '{param_name}'. Available: {list(CALIBRATION_PARAMS.keys())}")

    cfg = load_config(hardware)
    nominal = get_nominal_values(hardware)
    param_idx = PARAM_NAMES.index(param_name)

    metric_vals = np.zeros(len(values))

    for i, val in enumerate(values):
        params = list(nominal)
        params[param_idx] = val
        override_cfg = apply_overrides(cfg, params)
        args = override_cfg.to_engine_args()

        try:
            out, hist = ICV_open(
                Pexit_barg=Pexit, speed_f=speed,
                Ptank_barg=Ptank, Psat_barg=Psat, **args,
            )
            metric_vals[i] = _extract_metric(out, hist, metric)
        except Exception:
            metric_vals[i] = np.nan

    return SweepResult(
        param_name=param_name,
        param_values=np.array(values),
        metric_name=metric,
        metric_values=metric_vals,
        hardware=hardware,
        Pexit=Pexit,
        speed=speed,
    )


def sweep_2d(
    hardware: str,
    param1: str,
    values1: List[float],
    param2: str,
    values2: List[float],
    speed: float = 0.65,
    Pexit: float = 350.0,
    metric: str = "mdot_kgpm",
    Ptank: float = 7.0,
    Psat: float = 2.0,
) -> pd.DataFrame:
    """Sweep two parameters on a 2D grid.

    Args:
        hardware: Base config name.
        param1: First parameter name.
        values1: Values for first parameter.
        param2: Second parameter name.
        values2: Values for second parameter.
        speed, Pexit: Operating conditions.
        metric: Which metric to track.
        Ptank: Inlet tank pressure [barg].
        Psat: Saturation pressure [barg].

    Returns:
        DataFrame with columns: param1, param2, metric.
    """
    if param1 not in CALIBRATION_PARAMS or param2 not in CALIBRATION_PARAMS:
        raise ValueError(f"Unknown params. Available: {list(CALIBRATION_PARAMS.keys())}")

    cfg = load_config(hardware)
    nominal = get_nominal_values(hardware)
    idx1 = PARAM_NAMES.index(param1)
    idx2 = PARAM_NAMES.index(param2)

    rows = []
    for v1 in values1:
        for v2 in values2:
            params = list(nominal)
            params[idx1] = v1
            params[idx2] = v2
            override_cfg = apply_overrides(cfg, params)
            args = override_cfg.to_engine_args()

            try:
                out, hist = ICV_open(
                    Pexit_barg=Pexit, speed_f=speed,
                    Ptank_barg=Ptank, Psat_barg=Psat, **args,
                )
                val = _extract_metric(out, hist, metric)
            except Exception:
                val = np.nan

            rows.append({param1: v1, param2: v2, metric: val})

    return pd.DataFrame(rows)