run_definitive_tuned.py

#!/usr/bin/env python3
"""
DEFINITIVE TUNED PIPELINE
=========================
Loads best hyperparameters from Optuna studies (tuning/*_best.json) and
runs the full evaluation:
  1. LOCO CV on 9 train cohorts
  2. Refit on all train data
  3. Held-out test on GSE246337
  4. Full metrics + plots + comparison vs baseline (defaults)

Run AFTER tune_optuna.py has produced both elasticnet_loco_v1_best.json
and histgbr_loco_v1_best.json.

  python3 run_definitive_tuned.py
"""
import json, sys, gc, time, warnings
from pathlib import Path
import numpy as np
import pandas as pd
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.linear_model import ElasticNet
from sklearn.model_selection import LeaveOneGroupOut
from sklearn.metrics import mean_absolute_error
from scipy.stats import wilcoxon
import torch

warnings.filterwarnings("ignore")

PROJECT = Path(__file__).parent
ZENODO  = PROJECT / "data" / "zenodo" / "DeepStratv3"
PROC    = PROJECT / "data" / "processed"
RESULTS = PROJECT / "results"
TUNING  = PROJECT / "tuning"
SEED    = 42
SEED_MODEL = 1337


# ─── Load tuned params ───
def load_best(model_name):
    # Try v2 first (re-tuned on larger data), then v1
    for version in ["v2", "v1"]:
        fpath = TUNING / f"{model_name}_loco_{version}_best.json"
        if fpath.exists():
            with open(fpath) as f:
                d = json.load(f)
            print(f"  Loaded {model_name} ({version}): best CV MAE={d['best_value']:.4f} "
                  f"({d['n_trials']} trials, {d.get('elapsed_min','?')}min)", flush=True)
            return d["best_params"]
    print(f"  WARNING: no tuning results found for {model_name} — using defaults", flush=True)
    return None


# ─── Default params (baseline for comparison) ───
HGBR_DEFAULT = {
    "max_iter": 800, "max_depth": 4, "min_samples_leaf": 20,
    "learning_rate": 0.05, "max_leaf_nodes": 50, "l2_regularization": 0.1,
    "random_state": SEED, "early_stopping": True, "n_iter_no_change": 15,
    "validation_fraction": 0.15,
}
EN_DEFAULT = {
    "alpha": 0.01, "l1_ratio": 0.9, "max_iter": 5000, "random_state": SEED,
}


def build_hgbr_params(tuned):
    p = {
        "max_iter": 1000,
        "max_depth": None,
        "max_bins": 255,
        "early_stopping": True,
        "n_iter_no_change": 20,
        "validation_fraction": 0.15,
        "random_state": SEED_MODEL,
    }
    if tuned:
        p.update(tuned)
    else:
        p = HGBR_DEFAULT
    return p


def build_en_params(tuned):
    p = {
        "max_iter": 5000,
        "tol": 1e-3,
        "selection": "random",
        "random_state": SEED_MODEL,
    }
    if tuned:
        p.update(tuned)
    else:
        p = EN_DEFAULT
    return p


# ─── Load data ───
print("=== Loading datasets ===", flush=True)
DSA_CPGS = pd.read_csv(ZENODO / "resources" / "deep_strat_cpgs.csv", header=None)[0].tolist()
cpg_stats = pd.read_parquet(PROC / "epic_cpg_population_stats.parquet")
pop_means = cpg_stats["mean"].to_dict()

all_dfs = []
for f in sorted(PROC.glob("*.parquet")):
    if any(x in f.stem for x in ["combined", "epic_cpg", "part"]):
        continue
    df = pd.read_parquet(f)
    if "age" not in df.columns or len(df) == 0:
        continue
    all_dfs.append(df)
    geo = df["geo_id"].iloc[0]
    print(f"  {geo}: {len(df)}", flush=True)

combined = pd.concat(all_dfs, axis=0, join="outer")
del all_dfs; gc.collect()
meta = ["age", "sex", "geo_id"]
for cpg in DSA_CPGS:
    if cpg not in combined.columns:
        combined[cpg] = np.nan
combined = combined[DSA_CPGS + meta]

X_raw   = combined[DSA_CPGS].values.astype(np.float32)
y       = combined["age"].values.astype(np.float32)
sex     = combined["sex"].values.astype(np.float32)
cohorts = combined["geo_id"].values.astype(str)
print(f"\nTotal: {len(combined)} samples, {len(np.unique(cohorts))} cohorts", flush=True)

# ─── Imputation + features (mirrors tuning script exactly) ───
is_450k = (np.isnan(X_raw).mean(axis=1) > 0.3).astype(np.float32)
X_imp = X_raw.copy()
for i, cpg in enumerate(DSA_CPGS):
    nm = np.isnan(X_imp[:, i])
    if nm.any():
        X_imp[nm, i] = pop_means.get(cpg, 0.5)
X_imp = np.nan_to_num(X_imp, nan=0.5)

shared_idx = np.where(np.isnan(X_raw).mean(axis=0) < 0.01)[0]
X_sh = np.nan_to_num(X_raw[:, shared_idx], nan=0.5)
batch = np.column_stack([
    np.mean(X_sh, axis=1), np.std(X_sh, axis=1),
    (X_sh < 0.1).mean(axis=1), (X_sh > 0.9).mean(axis=1),
    np.median(X_sh, axis=1),
])
sex_clean = np.nan_to_num(sex, nan=0.5)
X_enhanced = np.column_stack([X_imp, sex_clean, is_450k, batch])
X_hgbr     = np.column_stack([X_raw, sex,       is_450k, batch])
X_dsa      = np.hstack([sex_clean.reshape(-1,1), X_imp]).astype(np.float32)

# ─── Test split ───
TEST = "GSE246337"
te = cohorts == TEST
tr = ~te
y_tr, y_te = y[tr], y[te]
cohorts_tr = cohorts[tr]
logo = LeaveOneGroupOut()
print(f"\n=== Test: {TEST} (n={te.sum()}) | Train: {tr.sum()} ({len(np.unique(cohorts_tr))} cohorts) ===", flush=True)

# ─── Load tuned params ───
print("\n=== Loading tuned hyperparameters ===", flush=True)
hgbr_tuned = load_best("histgbr")
en_tuned   = load_best("elasticnet")
hgbr_params = build_hgbr_params(hgbr_tuned)
en_params   = build_en_params(en_tuned)
print(f"\n  HistGBR: {hgbr_params}")
print(f"  ElasticNet: {en_params}", flush=True)


def run_loco(model_factory, X, y_tr, cohorts_tr, name):
    cv_preds = np.full(tr.sum(), np.nan)
    fold_maes = []
    for tri, tei in logo.split(X[tr], y_tr, cohorts_tr):
        m = model_factory()
        m.fit(X[tr][tri], y_tr[tri])
        cv_preds[tei] = m.predict(X[tr][tei])
        coh = cohorts_tr[tei][0]
        mae = mean_absolute_error(y_tr[tei], cv_preds[tei])
        fold_maes.append(mae)
        print(f"  CV {coh}: MAE={mae:.2f}", flush=True)
    return cv_preds, fold_maes


# ─── HistGBR (tuned) ───
print("\n--- HistGBR (TUNED) ---", flush=True)
t0 = time.time()
cv_preds_h, _ = run_loco(lambda: HistGradientBoostingRegressor(**hgbr_params),
                          X_hgbr, y_tr, cohorts_tr, "HistGBR")
print(f"  LOCO CV MAE: {mean_absolute_error(y_tr, cv_preds_h):.3f}", flush=True)
hgbr = HistGradientBoostingRegressor(**hgbr_params)
hgbr.fit(X_hgbr[tr], y_tr)
hp_test = hgbr.predict(X_hgbr[te])
print(f"  TEST MAE: {mean_absolute_error(y_te, hp_test):.3f} ({time.time()-t0:.0f}s)", flush=True)

# ─── ElasticNet (tuned) ───
print("\n--- ElasticNet (TUNED) ---", flush=True)
t0 = time.time()
cv_preds_e, _ = run_loco(lambda: ElasticNet(**en_params),
                          X_enhanced, y_tr, cohorts_tr, "EN")
print(f"  LOCO CV MAE: {mean_absolute_error(y_tr, cv_preds_e):.3f}", flush=True)
en = ElasticNet(**en_params)
en.fit(X_enhanced[tr], y_tr)
ep_test = en.predict(X_enhanced[te])
print(f"  TEST MAE: {mean_absolute_error(y_te, ep_test):.3f} ({time.time()-t0:.0f}s)", flush=True)

# ─── DeepStrataAge (unchanged) ───
print("\n--- DeepStrataAge ---", flush=True)
sys.path.insert(0, str(ZENODO))
from DeepStrataAge.model import EnsembleModel
dsa = EnsembleModel(input_dim=12235, model_paths=[
    str(ZENODO / "resources" / f"{n}_holdout.pth") for n in ["MGH", "BOA", "CIBMTR"]])
with torch.no_grad():
    dp_test = dsa(torch.tensor(X_dsa[te])).numpy().flatten()
print(f"  TEST MAE: {mean_absolute_error(y_te, dp_test):.3f}", flush=True)


# ─── Metrics ───
def metrics(name, yt, yp):
    err = yp - yt; ae = np.abs(err)
    mae = np.mean(ae); medae = np.median(ae)
    r2 = 1 - np.sum(err**2) / np.sum((yt - yt.mean())**2)
    pear = np.corrcoef(yt, yp)[0, 1]
    rng = np.random.RandomState(SEED)
    boot = [np.mean(np.abs(yt[j := rng.randint(0, len(yt), len(yt))] - yp[j])) for _ in range(1000)]
    ci = f"[{np.percentile(boot, 2.5):.2f},{np.percentile(boot, 97.5):.2f}]"
    return {"Model": name, "MAE": f"{mae:.2f}", "95%CI": ci, "MedAE": f"{medae:.2f}",
            "R2": f"{r2:.3f}", "Pearson": f"{pear:.3f}",
            "<=5y": f"{np.mean(ae <= 5) * 100:.1f}%", ">10y": f"{np.mean(ae > 10) * 100:.1f}%"}

print(f"\n{'='*80}\n  TUNED RESULTS\n{'='*80}", flush=True)
rows = [
    {"Split": "CV",   **metrics("HistGBR (tuned)",    y_tr, cv_preds_h)},
    {"Split": "CV",   **metrics("ElasticNet (tuned)", y_tr, cv_preds_e)},
    {"Split": "TEST", **metrics("HistGBR (tuned)",    y_te, hp_test)},
    {"Split": "TEST", **metrics("ElasticNet (tuned)", y_te, ep_test)},
    {"Split": "TEST", **metrics("DeepStrataAge",      y_te, dp_test)},
]
table = pd.DataFrame(rows)
print(table[["Split","Model","MAE","95%CI","MedAE","R2","Pearson","<=5y",">10y"]].to_string(index=False), flush=True)
table.to_csv(RESULTS / "definitive_tuned.csv", index=False)

# Paired tests
print(f"\n--- Paired Tests (test set) ---", flush=True)
for n1, p1, n2, p2 in [("HGBR_tuned", hp_test, "EN_tuned", ep_test),
                        ("HGBR_tuned", hp_test, "DSA",      dp_test),
                        ("EN_tuned",   ep_test, "DSA",      dp_test)]:
    _, p = wilcoxon(np.abs(p1 - y_te), np.abs(p2 - y_te))
    d = mean_absolute_error(y_te, p1) - mean_absolute_error(y_te, p2)
    sig = "***" if p<.001 else ("**" if p<.01 else ("*" if p<.05 else "n.s."))
    w = n1 if d<0 else n2
    print(f"  {n1} vs {n2}: D={d:+.2f} p={p:.6f} {sig} -> {w}", flush=True)

# Plots
import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
for i, (name, yp, c) in enumerate([("HistGBR (tuned)", hp_test, "#2E86C1"),
                                     ("ElasticNet (tuned)", ep_test, "#E74C3C"),
                                     ("DeepStrataAge", dp_test, "#8E44AD")]):
    ax = axes[i]
    ax.scatter(y_te, yp, alpha=.4, s=10, c=c)
    ax.plot([15, 95], [15, 95], "k--", lw=.8)
    ax.set_xlabel("Actual Age"); ax.set_ylabel("Predicted")
    ax.set_title(f"{name}\nMAE={mean_absolute_error(y_te, yp):.2f}")
plt.suptitle(f"TUNED | Test: {TEST}", y=1.02)
plt.tight_layout()
plt.savefig(RESULTS / "definitive_tuned_test.png", dpi=200)
plt.close()

print(f"\nSaved: {RESULTS}/definitive_tuned.csv + .png", flush=True)
print("DONE!", flush=True)