cryosim/optimization/sensitivity.py

"""Parameter sensitivity analysis via central finite differences."""

from dataclasses import dataclass, field
from typing import List, Dict, Optional

import numpy as np

from cryosim import predict
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


METRICS = ["mdot_kgpm", "mass_eff", "Tc_peak_K", "kWh_extend"]

_NEAR_ZERO_THRESHOLD = 0.01


@dataclass
class SensitivityResult:
    """Result of sensitivity analysis.

    sensitivities: dict of {metric: {param: normalized_sensitivity}}
        where normalized_sensitivity = (Δmetric/metric) / (Δparam/param)
        i.e., % change in metric per % change in parameter.
    rankings: dict of {metric: [(param, sensitivity), ...]} sorted by |sensitivity|
    zero_sensitivity_notes: warnings about params with near-zero sensitivity
    """
    param_names: List[str]
    metrics: List[str]
    sensitivities: Dict[str, Dict[str, float]]
    rankings: Dict[str, List[tuple]]
    hardware: str
    Pexit: float
    speed: float
    zero_sensitivity_notes: List[str] = field(default_factory=list)


@dataclass
class MultiSensitivityResult:
    """Sensitivity analysis across multiple operating points."""
    results: Dict[float, SensitivityResult]  # Pexit -> result
    combined_rankings: Dict[str, List[tuple]]  # metric -> [(param, max_abs_sensitivity)]
    zero_sensitivity_notes: List[str]  # warnings about params with near-zero sensitivity


def run_sensitivity(
    hardware: str = "old_icv",
    speed: float = 0.65,
    Pexit: float = 350.0,
    delta_frac: float = 0.05,
    Ptank: float = 7.0,
    Psat: float = 2.0,
) -> SensitivityResult:
    """Run sensitivity analysis on all 7 calibratable parameters.

    For each parameter, perturbs by ±delta_frac (default ±5%) and measures
    the change in each metric. Returns normalized sensitivities:
    (% change in metric) / (% change in param).

    Args:
        hardware: Base config name.
        speed: Speed fraction.
        Pexit: Discharge pressure [barg].
        delta_frac: Fractional perturbation (0.05 = ±5%).
        Ptank: Inlet tank pressure [barg].
        Psat: Saturation pressure [barg].

    Returns:
        SensitivityResult with sensitivities and rankings.
    """
    cfg = load_config(hardware)
    nominal = get_nominal_values(hardware)

    # Baseline prediction
    baseline = predict(hardware=hardware, Pexit=Pexit, speed=speed, Ptank=Ptank, Psat=Psat)
    baseline_metrics = {
        "mdot_kgpm": baseline.mdot_kgpm,
        "mass_eff": baseline.mass_eff,
        "Tc_peak_K": baseline.Tc_peak_K,
        "kWh_extend": baseline.kWh_extend,
    }

    sensitivities = {m: {} for m in METRICS}

    for i, pname in enumerate(PARAM_NAMES):
        info = CALIBRATION_PARAMS[pname]
        val = nominal[i]

        # Perturb ±delta_frac, clamped to bounds
        delta = max(abs(val) * delta_frac, 1e-8)
        val_lo = max(val - delta, info["low"])
        val_hi = min(val + delta, info["high"])
        actual_delta = val_hi - val_lo

        if actual_delta < 1e-12:
            for m in METRICS:
                sensitivities[m][pname] = 0.0
            continue

        # Build overridden configs for low and high perturbations
        params_lo = list(nominal)
        params_lo[i] = val_lo
        params_hi = list(nominal)
        params_hi[i] = val_hi

        cfg_lo = apply_overrides(cfg, params_lo)
        cfg_hi = apply_overrides(cfg, params_hi)

        try:
            args_lo = cfg_lo.to_engine_args()
            args_hi = cfg_hi.to_engine_args()

            out_lo, hist_lo = ICV_open(
                Pexit_barg=Pexit, speed_f=speed,
                Ptank_barg=Ptank, Psat_barg=Psat, **args_lo,
            )
            out_hi, hist_hi = ICV_open(
                Pexit_barg=Pexit, speed_f=speed,
                Ptank_barg=Ptank, Psat_barg=Psat, **args_hi,
            )

            metrics_lo = {
                "mdot_kgpm": hist_lo["mdot_kgpm"],
                "mass_eff": float(out_lo[1, 0]),
                "Tc_peak_K": float(np.max(hist_lo["Tc_K"])),
                "kWh_extend": hist_lo["kWh_extend"],
            }
            metrics_hi = {
                "mdot_kgpm": hist_hi["mdot_kgpm"],
                "mass_eff": float(out_hi[1, 0]),
                "Tc_peak_K": float(np.max(hist_hi["Tc_K"])),
                "kWh_extend": hist_hi["kWh_extend"],
            }
        except Exception:
            for m in METRICS:
                sensitivities[m][pname] = 0.0
            continue

        # Normalized sensitivity: (Δmetric/metric_baseline) / (Δparam/param_nominal)
        for m in METRICS:
            dm = metrics_hi[m] - metrics_lo[m]
            base_val = baseline_metrics[m]
            if abs(base_val) > 1e-12 and abs(val) > 1e-12:
                sensitivities[m][pname] = (dm / base_val) / (actual_delta / val)
            else:
                sensitivities[m][pname] = 0.0

    # Build rankings (sorted by absolute sensitivity)
    rankings = {}
    for m in METRICS:
        ranked = sorted(sensitivities[m].items(), key=lambda x: abs(x[1]), reverse=True)
        rankings[m] = ranked

    return SensitivityResult(
        param_names=list(PARAM_NAMES),
        metrics=list(METRICS),
        sensitivities=sensitivities,
        rankings=rankings,
        hardware=hardware,
        Pexit=Pexit,
        speed=speed,
    )


def run_multi_sensitivity(
    hardware: str = "old_icv",
    speed: float = 0.65,
    pressures: Optional[List[float]] = None,
    delta_frac: float = 0.05,
    Ptank: float = 7.0,
    Psat: float = 2.0,
) -> MultiSensitivityResult:
    """Run sensitivity analysis at multiple operating points.

    Reveals pressure-dependent behaviour that single-point analysis misses.
    For example, dcv_leakKv may show zero sensitivity at 350 bar (below the
    flow cliff) but significant sensitivity at 500 bar.

    Args:
        hardware: Base config name.
        speed: Speed fraction.
        pressures: List of Pexit values to test. Default [200, 380, 500].
        delta_frac: Fractional perturbation (0.05 = ±5%).
        Ptank: Inlet tank pressure [barg].
        Psat: Saturation pressure [barg].

    Returns:
        MultiSensitivityResult with per-pressure results, combined rankings,
        and diagnostic notes about near-zero sensitivities.
    """
    if pressures is None:
        pressures = [200.0, 380.0, 500.0]

    results: Dict[float, SensitivityResult] = {}
    for P in pressures:
        results[P] = run_sensitivity(
            hardware=hardware, speed=speed, Pexit=P,
            delta_frac=delta_frac, Ptank=Ptank, Psat=Psat,
        )

    # Combined rankings: for each param, take max |sensitivity| across pressures
    combined_rankings: Dict[str, List[tuple]] = {}
    for metric in METRICS:
        param_max: Dict[str, float] = {}
        for pname in PARAM_NAMES:
            max_abs = 0.0
            for P in pressures:
                s = abs(results[P].sensitivities[metric].get(pname, 0.0))
                if s > max_abs:
                    max_abs = s
            param_max[pname] = max_abs
        ranked = sorted(param_max.items(), key=lambda x: x[1], reverse=True)
        combined_rankings[metric] = ranked

    # Zero-sensitivity notes
    notes: List[str] = []
    for pname in PARAM_NAMES:
        # Collect max |sensitivity| across all metrics and all pressures
        per_pressure_max: Dict[float, float] = {}
        for P in pressures:
            max_across_metrics = max(
                abs(results[P].sensitivities[m].get(pname, 0.0))
                for m in METRICS
            )
            per_pressure_max[P] = max_across_metrics

        all_near_zero = all(v < _NEAR_ZERO_THRESHOLD for v in per_pressure_max.values())
        if all_near_zero:
            notes.append(
                f"{pname} shows near-zero sensitivity across all tested pressures "
                f"— may not be identifiable from the data"
            )
        else:
            # Check for pressure-dependent behavior: zero at some, significant at others
            zero_pressures = [P for P, v in per_pressure_max.items() if v < _NEAR_ZERO_THRESHOLD]
            sig_pressures = [P for P, v in per_pressure_max.items() if v >= _NEAR_ZERO_THRESHOLD]
            if zero_pressures and sig_pressures:
                zero_str = ", ".join(f"{p:.0f}" for p in zero_pressures)
                sig_str = ", ".join(f"{p:.0f}" for p in sig_pressures)
                notes.append(
                    f"{pname} has near-zero sensitivity at {zero_str} bar "
                    f"but significant sensitivity at {sig_str} bar "
                    f"— pressure-dependent"
                )

    return MultiSensitivityResult(
        results=results,
        combined_rankings=combined_rankings,
        zero_sensitivity_notes=notes,
    )