eval_hnf_ensemble.py

"""
eval_hnf_ensemble.py
--------------------
Evaluates hnf_ensemble_v1.pt across 8 gates and produces a report.

Usage:
    python eval_hnf_ensemble.py                          # expects .pt in same dir
    python eval_hnf_ensemble.py path/to/hnf_ensemble_v1.pt
    python eval_hnf_ensemble.py --teacher path/to/hnf_tight_150.pt

Gates:
    G0  Load + structural integrity
    G1  All models produce finite outputs
    G2  Ensemble always outputs negative delta (damping policy)
    G3  Fidelity R² vs teacher on held-out inputs
    G4  100% directional consensus across seeds
    G5  Uncertainty bounds (mean σ, max σ)
    G6  Weight identity (tight_150 in ensemble == direct load)
    G7  PRIME integration trace (standalone corrector mode)
"""

import sys
import argparse
import math
from pathlib import Path

# ── dependency check ──────────────────────────────────────────────────────────
try:
    import torch
    import torch.nn as nn
    import numpy as np
except ImportError as e:
    print(f"[ABORT] missing dependency: {e}")
    print("        pip install torch numpy")
    sys.exit(1)

# ── model definition (must match training) ────────────────────────────────────
class HNF(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(5, 32), nn.Tanh(),
            nn.Linear(32, 32), nn.Tanh(),
            nn.Linear(32, 1),
        )
    def forward(self, x):
        return self.net(x)


def load_net(state_dict):
    m = HNF()
    m.load_state_dict(state_dict)
    m.eval()
    return m


# ── ensemble predict ──────────────────────────────────────────────────────────
def ensemble_predict(nets, weights, H, t_norm, noise=0.0, gamma=0.3, eta=0.4, h_clip=100.0):
    Hc  = float(np.clip(H, -h_clip, h_clip))
    sig = 1.0 / (1.0 + math.exp(-gamma * Hc))
    inp = torch.tensor([[Hc, sig, noise, t_norm, eta]], dtype=torch.float32)
    with torch.no_grad():
        outs = np.array([net(inp).item() for net in nets])
    mean = float(np.dot(weights, outs))
    std  = float(np.sqrt(np.dot(weights, (outs - mean) ** 2)))
    return mean, std, outs


# ── gate helpers ──────────────────────────────────────────────────────────────
PASS = "PASS"
WARN = "WARN"
FAIL = "FAIL"

def result(status, msg):
    icon = {"PASS": "✓", "WARN": "⚠", "FAIL": "✗"}[status]
    return status, f"  {icon} {msg}"


def separator(title=""):
    w = 60
    if title:
        pad = (w - len(title) - 2) // 2
        print(f"\n{'─'*pad} {title} {'─'*(w - pad - len(title) - 2)}")
    else:
        print("─" * w)


# ── gates ─────────────────────────────────────────────────────────────────────
def gate0_load(ensemble_path):
    separator("G0  Load + structure")
    try:
        ckpt = torch.load(ensemble_path, map_location="cpu", weights_only=False)
    except Exception as e:
        s, m = result(FAIL, f"torch.load failed: {e}")
        print(m); return s, None

    required = ["ensemble_config", "weights", "model_state_dicts",
                "seed_metrics", "teacher_metrics", "ensemble_metrics"]
    missing  = [k for k in required if k not in ckpt]
    if missing:
        s, m = result(FAIL, f"missing keys: {missing}")
        print(m); return s, None
    print(result(PASS, f"all required keys present")[1])

    expected_models = ["tight_150", "seed_1", "seed_2", "seed_3"]
    missing_models  = [k for k in expected_models if k not in ckpt["model_state_dicts"]]
    if missing_models:
        s, m = result(FAIL, f"missing models: {missing_models}")
        print(m); return s, None
    print(result(PASS, f"models: {list(ckpt['model_state_dicts'].keys())}")[1])

    w_sum = sum(ckpt["weights"])
    if abs(w_sum - 1.0) > 1e-4:
        s, m = result(FAIL, f"weights sum = {w_sum:.6f} (expected 1.0)")
        print(m); return s, None
    print(result(PASS, f"weights sum = {w_sum:.6f}  values = {[round(w,4) for w in ckpt['weights']]}")[1])

    return PASS, ckpt


def gate1_finite(nets, weights, eta, gamma):
    separator("G1  Finite outputs")
    torch.manual_seed(0)
    test_inputs = torch.randn(200, 5)
    test_inputs[:, 1] = torch.sigmoid(0.3 * test_inputs[:, 0])
    test_inputs[:, 4] = eta

    all_ok = True
    for name, net in nets.items():
        with torch.no_grad():
            out = net(test_inputs)
        nan, inf = out.isnan().any().item(), out.isinf().any().item()
        lo, hi   = out.min().item(), out.max().item()
        if nan or inf:
            s, m = result(FAIL, f"{name}: nan={nan} inf={inf}")
            all_ok = False
        else:
            s, m = result(PASS, f"{name}: range=[{lo:.5f}, {hi:.5f}]")
        print(m)

    return PASS if all_ok else FAIL


def gate2_always_negative(nets_list, weights, eta, gamma):
    separator("G2  Ensemble always negative")
    H_sweep     = np.linspace(-10, 10, 2000)
    violations  = []
    for H in H_sweep:
        mean, _, _ = ensemble_predict(nets_list, weights, H, 0.5, gamma=gamma, eta=eta)
        if mean >= 0:
            violations.append((H, mean))

    if violations:
        s, m = result(WARN, f"{len(violations)} non-negative outputs found")
        print(m)
        for H, v in violations[:3]:
            print(f"       H={H:.3f} → {v:.6f}")
        return WARN
    else:
        _, m = result(PASS, f"0 violations across H∈[-10,10] (2000 points)")
        print(m)
        return PASS


def gate3_fidelity(nets_list, weights, teacher_net, eta, gamma, ref_loss, ref_lyap):
    separator("G3  Fidelity vs teacher (held-out)")
    np.random.seed(99)
    H_test = np.concatenate([np.random.uniform(-8, 8, 400), np.linspace(-10, 10, 200)])

    ref_outs, ens_outs = [], []
    for H in H_test:
        sig = 1.0 / (1.0 + math.exp(-gamma * H))
        inp = torch.tensor([[H, sig, 0.0, 0.5, eta]], dtype=torch.float32)
        with torch.no_grad():
            ref_outs.append(teacher_net(inp).item())
        mean, _, _ = ensemble_predict(nets_list, weights, H, 0.5, gamma=gamma, eta=eta)
        ens_outs.append(mean)

    ref_arr = np.array(ref_outs)
    ens_arr = np.array(ens_outs)
    mse = float(np.mean((ref_arr - ens_arr) ** 2))
    mae = float(np.mean(np.abs(ref_arr - ens_arr)))
    r2  = float(1 - mse / (np.var(ref_arr) + 1e-12))

    status = PASS if r2 >= 0.5 else WARN
    _, m = result(status, f"R²={r2:.6f}  MSE={mse:.8f}  MAE={mae:.6f}")
    print(m)
    print(f"       teacher range : [{ref_arr.min():.5f}, {ref_arr.max():.5f}]")
    print(f"       ensemble range: [{ens_arr.min():.5f}, {ens_arr.max():.5f}]")
    print(f"       teacher  loss={ref_loss:.6f}  lyapunov={ref_lyap:.6f}")
    return status, r2


def gate4_consensus(nets_list, weights, eta, gamma):
    separator("G4  Directional consensus")
    H_test = np.linspace(-10, 10, 1000)
    agree  = 0
    for H in H_test:
        _, _, outs = ensemble_predict(nets_list, weights, H, 0.5, gamma=gamma, eta=eta)
        if np.all(outs < 0):
            agree += 1

    pct    = agree / len(H_test) * 100
    status = PASS if pct >= 95 else WARN
    _, m   = result(status, f"all nets agree negative: {agree}/{len(H_test)} = {pct:.1f}%")
    print(m)
    return status


def gate5_uncertainty(nets_list, weights, eta, gamma):
    separator("G5  Uncertainty bounds")
    H_sweep = np.linspace(-10, 10, 1000)
    stds    = []
    for H in H_sweep:
        _, std, _ = ensemble_predict(nets_list, weights, H, 0.5, gamma=gamma, eta=eta)
        stds.append(std)
    stds     = np.array(stds)
    mean_std = float(stds.mean())
    max_std  = float(stds.max())
    max_H    = H_sweep[np.argmax(stds)]

    s1, m1 = result(PASS if mean_std < 0.05 else WARN, f"mean σ = {mean_std:.6f}")
    s2, m2 = result(PASS if max_std  < 0.10 else WARN, f"max  σ = {max_std:.6f}  at H≈{max_H:.2f}")
    print(m1); print(m2)
    return PASS if (mean_std < 0.05 and max_std < 0.10) else WARN


def gate6_weight_identity(ckpt, teacher_path):
    separator("G6  Weight identity")
    if teacher_path is None:
        _, m = result(WARN, "no --teacher path provided; skipping identity check")
        print(m); return WARN

    try:
        teacher_ckpt = torch.load(teacher_path, map_location="cpu", weights_only=False)
        sd_direct    = teacher_ckpt["model_state_dict"]
        sd_ensemble  = ckpt["model_state_dicts"]["tight_150"]
        all_match    = all(torch.allclose(sd_ensemble[k], sd_direct[k]) for k in sd_direct)
        status       = PASS if all_match else FAIL
        _, m = result(status, f"tight_150 in ensemble == direct load: {all_match}")
        print(m)
    except Exception as e:
        _, m = result(WARN, f"could not verify: {e}")
        print(m); return WARN

    s1 = ckpt["model_state_dicts"]["seed_1"]["net.0.weight"]
    s2 = ckpt["model_state_dicts"]["seed_2"]["net.0.weight"]
    s3 = ckpt["model_state_dicts"]["seed_3"]["net.0.weight"]
    s12 = torch.allclose(s1, s2)
    s13 = torch.allclose(s1, s3)
    _, m1 = result(FAIL if s12 else PASS, f"seed_1 != seed_2: {not s12}")
    _, m2 = result(FAIL if s13 else PASS, f"seed_1 != seed_3: {not s13}")
    print(m1); print(m2)
    return PASS if (all_match and not s12 and not s13) else FAIL


def gate7_integration(nets_list, weights, cfg):
    separator("G7  PRIME integration (standalone corrector)")
    eta, gamma  = cfg["eta"], cfg["gamma"]
    beta, eps   = cfg["beta"], cfg["epsilon"]
    h_clip      = cfg["h_clip"]
    rng         = np.random.RandomState(42)

    print("  Mode: ensemble IS the dynamics  (H_new = H + ensemble(H))")
    print()
    print(f"  {'H0':>6}  {'H_final':>10}  {'std':>8}  {'min':>8}  {'max':>8}  status")
    print(f"  {'─'*6}  {'─'*10}  {'─'*8}  {'─'*8}  {'─'*8}  {'─'*6}")

    results = {}
    for H0 in [0.5, 2.0, -2.0, 5.0, -5.0, 0.0, 10.0, -10.0]:
        H = float(H0)
        traj = [H]
        for t in range(100):
            Hc    = float(np.clip(H, -h_clip, h_clip))
            noise = float(rng.normal(0, 0.01))
            delta, _, _ = ensemble_predict(nets_list, weights, Hc, t/100,
                                           noise=noise, gamma=gamma, eta=eta, h_clip=h_clip)
            H = float(np.clip(Hc + delta, -h_clip, h_clip))
            traj.append(H)
        arr    = np.array(traj)
        stable = abs(arr[-1]) < h_clip * 0.8
        status = PASS if stable else WARN
        results[H0] = dict(final=arr[-1], std=arr.std(), status=status)
        _, m = result(status, f"{H0:>+6.1f}  {arr[-1]:>+10.4f}  {arr.std():>8.4f}"
                              f"  {arr.min():>+8.3f}  {arr.max():>+8.3f}")
        print(m)

    n_warn = sum(1 for r in results.values() if r["status"] == WARN)
    return WARN if n_warn > 2 else PASS


# ── main ──────────────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(description="Evaluate hnf_ensemble_v1.pt")
    parser.add_argument("ensemble",  nargs="?",
                        default=str(Path(__file__).parent / "hnf_ensemble_v1.pt"),
                        help="path to hnf_ensemble_v1.pt")
    parser.add_argument("--teacher", default=None,
                        help="path to hnf_tight_150.pt for G6 identity check")
    parser.add_argument("--verbose", action="store_true")
    args = parser.parse_args()

    print()
    print("╔══════════════════════════════════════════════════╗")
    print("║    HNF Ensemble Evaluation — hnf_ensemble_v1    ║")
    print("╚══════════════════════════════════════════════════╝")
    print(f"  checkpoint : {args.ensemble}")
    print(f"  teacher    : {args.teacher or '(none — G6 will be skipped)'}")
    print()

    # G0
    g0_status, ckpt = gate0_load(args.ensemble)
    if g0_status == FAIL:
        print("\n[ABORT] G0 failed — cannot continue"); sys.exit(1)

    cfg     = ckpt["ensemble_config"]["prime_config"]
    eta     = cfg["eta"]
    gamma   = cfg["gamma"]
    weights = np.array(ckpt["weights"])

    nets_dict = {name: load_net(sd) for name, sd in ckpt["model_state_dicts"].items()}
    nets_list = list(nets_dict.values())

    # Load teacher for G3 + G6
    teacher_net = nets_dict["tight_150"]  # already in ensemble
    ref_loss    = ckpt["teacher_metrics"]["loss"]
    ref_lyap    = ckpt["teacher_metrics"]["lyapunov"]

    # G1–G7
    g1 = gate1_finite(nets_dict, weights, eta, gamma)
    g2 = gate2_always_negative(nets_list, weights, eta, gamma)
    g3, r2 = gate3_fidelity(nets_list, weights, teacher_net, eta, gamma, ref_loss, ref_lyap)
    g4 = gate4_consensus(nets_list, weights, eta, gamma)
    g5 = gate5_uncertainty(nets_list, weights, eta, gamma)
    g6 = gate6_weight_identity(ckpt, args.teacher)
    g7 = gate7_integration(nets_list, weights, cfg)

    # Summary
    gates   = {"G0 Load+structure":     g0_status,
               "G1 Finite outputs":     g1,
               "G2 Always negative":    g2,
               f"G3 Fidelity R²={r2:.3f}": g3,
               "G4 Directional consensus": g4,
               "G5 Uncertainty bounds": g5,
               "G6 Weight identity":    g6,
               "G7 PRIME integration":  g7}

    separator("SUMMARY")
    for name, status in gates.items():
        icon = {"PASS": "✓", "WARN": "⚠", "FAIL": "✗"}[status]
        print(f"  {icon}  {name}: {status}")

    n_pass = sum(1 for v in gates.values() if v == PASS)
    n_warn = sum(1 for v in gates.values() if v == WARN)
    n_fail = sum(1 for v in gates.values() if v == FAIL)

    print()
    print(f"  {n_pass}/8 PASS   {n_warn}/8 WARN   {n_fail}/8 FAIL")

    if n_fail > 0:
        verdict = "FAIL — do not use"
    elif n_warn > 2:
        verdict = "CONDITIONAL — review warnings before use"
    elif n_warn > 0:
        verdict = "CONDITIONAL SHIP — known limitations documented"
    else:
        verdict = "SHIP"

    print(f"  Verdict: {verdict}")

    # Known limitation note
    separator()
    print("  NOTE: G7 standalone mode is the correct integration pattern.")
    print("  Do NOT use as 0.05*net(H) addon inside PRIME loop — mixing")
    print("  coefficient is insufficient to overcome η·H·σ(γH) for H>0.")
    print("  Use as: H_new = H + ensemble(H, t_norm)")
    print()

    sys.exit(0 if n_fail == 0 else 1)


if __name__ == "__main__":
    main()