tune_optuna.py

#!/usr/bin/env python3
"""
OPTUNA HYPERPARAMETER TUNING for HistGBR + ElasticNet
======================================================
Methodology:

CV strategy: single-level LOCO over 9 train cohorts (GSE246337 NEVER touched).
            GSE246337 is our outer test loop — nested CV is not needed.
Sampler:    TPESampler(multivariate=True), separate seed from CV/model.
Pruner:     Fold-level pruning (HistGBR has no native callback).
Objective:  mean LOCO MAE  (also tracks std + max-fold for diagnostics).
Storage:    SQLite local — resumable.
Stopping:   plateau callback (no improvement in 30 trials).

Critical guards:
- Test cohort GSE246337 is loaded but NEVER passed to any fit/predict during tuning.
- Assertion at the start fails if anyone tries to read test labels.
- Batch features computed ONCE, outside the CV loop (deterministic, not target-dependent).

Usage:
  python3 tune_optuna.py --model elasticnet --n-trials 100
  python3 tune_optuna.py --model histgbr    --n-trials 100
"""
import argparse, gc, sys, time, warnings
from pathlib import Path

import numpy as np
import pandas as pd
import optuna
from optuna.samplers import TPESampler
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.linear_model import ElasticNet
from sklearn.model_selection import LeaveOneGroupOut
from sklearn.metrics import mean_absolute_error

warnings.filterwarnings("ignore")
optuna.logging.set_verbosity(optuna.logging.WARNING)

# ─────────────────────────────────────────────────────────────────────
# Constants
# ─────────────────────────────────────────────────────────────────────
PROJECT = Path(__file__).parent
ZENODO  = PROJECT / "data" / "zenodo" / "DeepStratv3"
PROC    = PROJECT / "data" / "processed"
RESULTS = PROJECT / "results"
TUNING  = PROJECT / "tuning"
TUNING.mkdir(exist_ok=True)

TEST_COHORT = "GSE246337"   # held out — NEVER touched during tuning
SEED_SAMPLER = 42
SEED_MODEL   = 1337
PLATEAU_PATIENCE = 30


# ─────────────────────────────────────────────────────────────────────
# Data loading (mirrors run_definitive.py exactly, but excludes test)
# ─────────────────────────────────────────────────────────────────────
def load_train_data():
    print("=== Loading datasets (TEST cohort excluded) ===", 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)

    # ── HARD GUARD: drop test cohort BEFORE any modeling ──
    train_mask = cohorts != TEST_COHORT
    n_test_dropped = (~train_mask).sum()
    print(f"\n  Dropped {n_test_dropped} samples from test cohort {TEST_COHORT}",
          flush=True)
    print(f"  Train pool: {train_mask.sum()} samples from "
          f"{len(np.unique(cohorts[train_mask]))} cohorts", flush=True)

    X_raw_tr   = X_raw[train_mask]
    y_tr       = y[train_mask]
    sex_tr     = sex[train_mask]
    cohorts_tr = cohorts[train_mask]
    del X_raw, y, sex, cohorts, combined; gc.collect()

    # ── Imputation: pop-mean (deterministic, no target dependency) ──
    print("\n=== Imputation (pop-mean, target-independent) ===", flush=True)
    is_450k_tr = (np.isnan(X_raw_tr).mean(axis=1) > 0.3).astype(np.float32)
    print(f"  450K: {is_450k_tr.sum():.0f}, EPIC: {(1-is_450k_tr).sum():.0f}",
          flush=True)

    X_imp_tr = X_raw_tr.copy()
    for i, cpg in enumerate(DSA_CPGS):
        nm = np.isnan(X_imp_tr[:, i])
        if nm.any():
            X_imp_tr[nm, i] = pop_means.get(cpg, 0.5)
    X_imp_tr = np.nan_to_num(X_imp_tr, nan=0.5)
    print(f"  NaN remaining: {np.isnan(X_imp_tr).sum()}", flush=True)

    # ── Batch features (cohort-level, target-independent) ──
    print("\n=== Batch features (deterministic, computed once) ===", flush=True)
    shared_idx = np.where(np.isnan(X_raw_tr).mean(axis=0) < 0.01)[0]
    X_sh = np.nan_to_num(X_raw_tr[:, 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),
    ])
    print(f"  Shared CpGs: {len(shared_idx)}, Batch features: 5", flush=True)

    sex_clean = np.nan_to_num(sex_tr, nan=0.5)
    X_enhanced = np.column_stack([X_imp_tr, sex_clean, is_450k_tr, batch])  # ElasticNet
    X_hgbr     = np.column_stack([X_raw_tr, sex_tr, is_450k_tr, batch])     # HistGBR (NaN-native)

    assert np.isnan(X_enhanced).sum() == 0
    print(f"  X_enhanced shape: {X_enhanced.shape}, X_hgbr shape: {X_hgbr.shape}",
          flush=True)

    return X_enhanced, X_hgbr, y_tr, cohorts_tr


# ─────────────────────────────────────────────────────────────────────
# Plateau-stopping callback
# ─────────────────────────────────────────────────────────────────────
class PlateauCallback:
    def __init__(self, patience=30):
        self.patience = patience
        self.best = float("inf")
        self.no_improve = 0

    def __call__(self, study, trial):
        if trial.state != optuna.trial.TrialState.COMPLETE:
            return
        v = trial.value
        if v < self.best - 1e-4:
            self.best = v
            self.no_improve = 0
        else:
            self.no_improve += 1
        if self.no_improve >= self.patience:
            print(f"\n  >>> Plateau: no improvement in {self.patience} trials. "
                  f"Stopping. <<<", flush=True)
            study.stop()


# ─────────────────────────────────────────────────────────────────────
# Objective: ElasticNet
# ─────────────────────────────────────────────────────────────────────
def make_objective_en(X, y, cohorts, logo, subsample_frac=None):
    rng_sub = np.random.RandomState(SEED_MODEL)

    def objective(trial):
        params = {
            "alpha":     trial.suggest_float("alpha",     1e-4, 10.0, log=True),
            "l1_ratio":  trial.suggest_float("l1_ratio",  0.05, 0.95),
            "max_iter":  5000,
            "tol":       1e-3,
            "selection": "random",
            "random_state": SEED_MODEL,
        }
        fold_maes = []
        for fold_i, (tri, tei) in enumerate(logo.split(X, y, cohorts)):
            if subsample_frac is not None and subsample_frac < 1.0:
                n_sub = int(len(tri) * subsample_frac)
                tri = rng_sub.choice(tri, size=n_sub, replace=False)
            en = ElasticNet(**params)
            en.fit(X[tri], y[tri])
            yp = en.predict(X[tei])
            mae = mean_absolute_error(y[tei], yp)
            fold_maes.append(mae)
            trial.report(np.mean(fold_maes), step=fold_i)
            if trial.should_prune():
                raise optuna.TrialPruned()
        mean_mae = float(np.mean(fold_maes))
        std_mae  = float(np.std(fold_maes))
        max_mae  = float(np.max(fold_maes))
        trial.set_user_attr("std_mae", std_mae)
        trial.set_user_attr("max_mae", max_mae)
        return mean_mae
    return objective


# ─────────────────────────────────────────────────────────────────────
# Objective: HistGBR
# ─────────────────────────────────────────────────────────────────────
def make_objective_hgbr(X, y, cohorts, logo, subsample_frac=None):
    """
    When subsample_frac < 1.0, each fold's training set is randomly subsampled
    (without replacement). Validation fold remains full-size. Tune on subsample
    for speed, then refit on full data with best params.
    """
    rng_sub = np.random.RandomState(SEED_MODEL)

    def objective(trial):
        # Budget-constrained ranges: lr >=0.03 forces faster convergence,
        # max_iter=300 + aggressive early stopping (n_iter_no_change=8).
        params = {
            "learning_rate":      trial.suggest_float("learning_rate", 0.03, 0.2, log=True),
            "max_iter":           300,
            "max_leaf_nodes":     trial.suggest_int("max_leaf_nodes", 15, 127, log=True),
            "min_samples_leaf":   trial.suggest_int("min_samples_leaf", 20, 200, log=True),
            "l2_regularization":  trial.suggest_float("l2_regularization", 1e-6, 10.0, log=True),
            "max_features":       trial.suggest_float("max_features", 0.3, 1.0),
            "max_bins":           255,
            "max_depth":          None,
            "early_stopping":     True,
            "n_iter_no_change":   8,
            "validation_fraction": 0.15,
            "random_state":       SEED_MODEL,
        }
        fold_maes = []
        for fold_i, (tri, tei) in enumerate(logo.split(X, y, cohorts)):
            # Subsample train (preserves val fold integrity)
            if subsample_frac is not None and subsample_frac < 1.0:
                n_sub = int(len(tri) * subsample_frac)
                sub_idx = rng_sub.choice(tri, size=n_sub, replace=False)
                tri = sub_idx
            m = HistGradientBoostingRegressor(**params)
            m.fit(X[tri], y[tri])
            yp = m.predict(X[tei])
            mae = mean_absolute_error(y[tei], yp)
            fold_maes.append(mae)
            trial.report(np.mean(fold_maes), step=fold_i)
            if trial.should_prune():
                raise optuna.TrialPruned()
        mean_mae = float(np.mean(fold_maes))
        std_mae  = float(np.std(fold_maes))
        max_mae  = float(np.max(fold_maes))
        trial.set_user_attr("std_mae", std_mae)
        trial.set_user_attr("max_mae", max_mae)
        return mean_mae
    return objective


# ─────────────────────────────────────────────────────────────────────
# Trial logger callback
# ─────────────────────────────────────────────────────────────────────
def log_trial(study, trial):
    import sys
    if trial.state == optuna.trial.TrialState.COMPLETE:
        std = trial.user_attrs.get("std_mae", 0.0)
        mx  = trial.user_attrs.get("max_mae", 0.0)
        flag = " ★ NEW BEST" if trial.value <= study.best_value + 1e-9 else ""
        sys.stdout.write(f"  Trial {trial.number:3d}: MAE={trial.value:.4f} "
                         f"std={std:.3f} max={mx:.3f}{flag}\n")
        sys.stdout.flush()
    elif trial.state == optuna.trial.TrialState.PRUNED:
        sys.stdout.write(f"  Trial {trial.number:3d}: PRUNED\n")
        sys.stdout.flush()


# ─────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", choices=["elasticnet", "histgbr"], required=True)
    ap.add_argument("--n-trials", type=int, default=100)
    ap.add_argument("--study-name", type=str, default=None)
    ap.add_argument("--subsample", type=float, default=None,
                    help="HistGBR only: subsample fraction (e.g. 0.6 = 60%% of train per fold)")
    args = ap.parse_args()

    X_enh, X_hgbr, y, cohorts = load_train_data()
    logo = LeaveOneGroupOut()
    n_folds = logo.get_n_splits(X_enh, y, cohorts)
    print(f"\n  LOCO folds: {n_folds}", flush=True)

    study_name = args.study_name or f"{args.model}_loco_v1"
    storage = f"sqlite:///{TUNING}/{study_name}.db"

    sampler = TPESampler(
        multivariate=True,
        seed=SEED_SAMPLER,
        n_startup_trials=20,
        constant_liar=True,
    )
    pruner = optuna.pruners.MedianPruner(
        n_startup_trials=10, n_warmup_steps=3, interval_steps=1)

    study = optuna.create_study(
        study_name=study_name,
        storage=storage,
        load_if_exists=True,
        direction="minimize",
        sampler=sampler,
        pruner=pruner,
    )

    if args.model == "elasticnet":
        objective = make_objective_en(X_enh, y, cohorts, logo,
                                      subsample_frac=args.subsample)
        sub_note = f" (subsample={args.subsample})" if args.subsample else ""
        eta_per_trial = f"~3 min{sub_note}"
    else:
        objective = make_objective_hgbr(X_hgbr, y, cohorts, logo,
                                         subsample_frac=args.subsample)
        sub_note = f" (subsample={args.subsample})" if args.subsample else ""
        eta_per_trial = f"~5 min{sub_note}"

    print(f"\n=== Tuning {args.model} ({args.n_trials} trials, ETA {eta_per_trial}/trial) ===")
    print(f"  Storage: {storage}")
    print(f"  Already completed: {len(study.trials)} trials\n", flush=True)

    plateau = PlateauCallback(patience=PLATEAU_PATIENCE)
    t0 = time.time()
    study.optimize(
        objective,
        n_trials=args.n_trials,
        callbacks=[log_trial, plateau],
        gc_after_trial=True,
        show_progress_bar=False,
    )
    elapsed = time.time() - t0

    print(f"\n=== Done in {elapsed/60:.1f} min ===")
    print(f"  Best MAE: {study.best_value:.4f}")
    print(f"  Best params:")
    for k, v in study.best_params.items():
        print(f"    {k}: {v}")

    # Save best params
    out = TUNING / f"{study_name}_best.json"
    import json
    with open(out, "w") as f:
        json.dump({
            "model": args.model,
            "best_value": study.best_value,
            "best_params": study.best_params,
            "n_trials": len(study.trials),
            "elapsed_min": round(elapsed/60, 1),
        }, f, indent=2)
    print(f"\n  Saved: {out}")


if __name__ == "__main__":
    main()