hardware-counters

#!/usr/bin/env python3
"""
State-of-the-art tabular NN techniques for HPC runtime prediction.

Same protocol as retrain2.py: merged configurations, median-constant baseline,
leave-one-application-out, error = max(pred/act, act/pred).

Techniques and sources:
  PLR numerical embeddings  Gorishniy et al., NeurIPS 2022 (arXiv:2203.05556)
  robust scale + smooth clip  Holzmuller et al., NeurIPS 2024 ("Better by Default")
  deep ensembling           Gorishniy et al., TabM (arXiv:2410.24210)
  RealMLP                   Holzmuller et al., NeurIPS 2024, via pytabkit

COST NOTE: a full nested grid (8 outer x 7 inner x 12 configs) was measured at
69 s per inner evaluation, i.e. ~13 hours. That is disproportionate for 191
rows, so the inner search uses:
  * a 4-config random subsample of the grid (fixed seed, so it is reproducible
    and identical across folds);
  * 3-fold grouped inner CV instead of full inner LOAO;
  * cheaper inner fits (1 ensemble member, 150 epochs) with the full ensemble
    only for the final refit.
This keeps hyperparameter selection honest -- no configuration is chosen using
the evaluation fold -- while costing ~30 min instead of ~13 h.
"""
import os, warnings, sys, time
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.impute import SimpleImputer
from scipy.stats import wilcoxon

warnings.filterwarnings("ignore")
torch.set_num_threads(32)
# Data and output roots. Overridable so the pipeline can be relocated
# (a different account, a scratch copy, another site) without editing code;
# the default keeps existing invocations working unchanged.
D      = os.environ.get("DISS_ROOT", "/work/project/project/user")
OUT    = Path(f"{D}/analysis/out"); OUT.mkdir(parents=True, exist_ok=True)
F_PEAK = 2.25e9

sys.path.insert(0, f"{D}/analysis")
from retrain2 import load_merged, build_features, fac
from advanced_nn import RobustSmoothScaler, PLRNet, train_plr


def prep():
    df = load_merged()
    X  = build_features(df)
    y  = df.runtime_s.values
    t_an = (df.PAPI_TOT_CYC / F_PEAK).values
    return df, X, y, t_an, t_an / y


def scaled(Xtr, Xte):
    imp = SimpleImputer(strategy="median").fit(Xtr)
    sc  = RobustSmoothScaler().fit(imp.transform(Xtr))
    return sc.transform(imp.transform(Xtr)), sc.transform(imp.transform(Xte))


FULL_GRID = [{"sigma": s, "k": k, "hidden": h}
             for s in (0.01, 0.05, 0.2)
             for k in (16, 32)
             for h in ((128, 64), (256, 128))]
rng = np.random.RandomState(0)
GRID = [FULL_GRID[i] for i in rng.choice(len(FULL_GRID), 4, replace=False)]


def main():
    df, X, y, t_an, eta = prep()
    apps = sorted(df.app.unique())
    tgt, Xv = np.log10(eta), X.values

    L, rows = [], {}
    say = lambda s="": (print(s, flush=True), L.append(s))
    say("Advanced tabular NN techniques")
    say("=" * 78)
    say(f"{len(df)} merged configurations, {X.shape[1]} features, {len(apps)} applications")
    say(f"inner search grid (random subsample of {len(FULL_GRID)}): "
        + ", ".join(f"s{g['sigma']}/k{g['k']}/w{g['hidden'][0]}" for g in GRID))
    say("")

    def record(name, fn):
        t0 = time.time()
        r = np.full(len(y), np.nan)
        for a in apps:
            te = (df.app == a).values
            r[te] = fac(fn(Xv[~te], tgt[~te], Xv[te], t_an[te], df.app[~te].values), y[te])
        rows[name] = r
        say(f"  {name:36s} median {np.nanmedian(r):.4f}  p90 {np.nanpercentile(r,90):.4f}"
            f"   [{time.time()-t0:.0f}s]")
        return r

    record("Constant (median log-eta)",
           lambda a, b, c, ta, ap: ta / (10 ** np.full(len(c), np.median(b))))

    def rf(Xtr, ytr, Xte, ta, ap):
        s1, s2 = scaled(Xtr, Xte)
        return ta / (10 ** RandomForestRegressor(n_estimators=400, min_samples_leaf=2,
                     random_state=0, n_jobs=-1).fit(s1, ytr).predict(s2))
    record("Random Forest", rf)

    def gbq(Xtr, ytr, Xte, ta, ap):
        s1, s2 = scaled(Xtr, Xte)
        return ta / (10 ** GradientBoostingRegressor(loss="quantile", alpha=0.5,
                     n_estimators=300, max_depth=3, learning_rate=0.05,
                     random_state=0).fit(s1, ytr).predict(s2))
    record("GBM quantile(0.5)", gbq)

    # ---- PLR at fixed literature defaults (no search) --------------------
    def plr_fixed(Xtr, ytr, Xte, ta, ap):
        s1, s2 = scaled(Xtr, Xte)
        return ta / (10 ** train_plr(s1, ytr, s2, sigma=0.05, k=24,
                                     hidden=(128, 64), n_ens=5, epochs=400))
    record("PLR ensemble (fixed defaults)", plr_fixed)

    # ---- PLR with cheap but honest nested selection ----------------------
    picks = []
    def plr_nested(Xtr, ytr, Xte, ta, ap):
        uapps = np.unique(ap)
        # 3 grouped inner folds instead of 7, to bound cost
        folds = np.array_split(uapps, 3)
        best, bg = None, None
        for g in GRID:
            errs = []
            for f in folds:
                ite = np.isin(ap, f)
                if ite.sum() == 0 or (~ite).sum() < 20:
                    continue
                s1, s2 = scaled(Xtr[~ite], Xtr[ite])
                p = train_plr(s1, ytr[~ite], s2, n_ens=1, epochs=150, **g)
                errs.append(np.median(np.abs(p - ytr[ite])))
            e = float(np.mean(errs)) if errs else 1e9
            if best is None or e < best:
                best, bg = e, g
        picks.append(f"s{bg['sigma']}/k{bg['k']}/w{bg['hidden'][0]}")
        s1, s2 = scaled(Xtr, Xte)
        return ta / (10 ** train_plr(s1, ytr, s2, n_ens=5, epochs=400, **bg))
    record("PLR ensemble (nested select)", plr_nested)
    say(f"      inner picks: {', '.join(picks)}")

    # ---- RealMLP ---------------------------------------------------------
    try:
        from pytabkit import RealMLP_TD_Regressor
        def realmlp(Xtr, ytr, Xte, ta, ap):
            s1, s2 = scaled(Xtr, Xte)
            m = RealMLP_TD_Regressor(random_state=0, device="cpu", verbosity=0)
            m.fit(s1, ytr)
            return ta / (10 ** np.asarray(m.predict(s2)).ravel())
        record("RealMLP-TD (pre-tuned)", realmlp)
    except Exception as e:
        say(f"  RealMLP failed: {type(e).__name__}: {e}")

    # ---- hybrid ----------------------------------------------------------
    def hybrid(Xtr, ytr, Xte, ta, ap):
        s1, s2 = scaled(Xtr, Xte)
        p1 = train_plr(s1, ytr, s2, sigma=0.05, k=24, hidden=(128, 64),
                       n_ens=5, epochs=400)
        p2 = RandomForestRegressor(n_estimators=400, min_samples_leaf=2,
                                   random_state=0, n_jobs=-1).fit(s1, ytr).predict(s2)
        return ta / (10 ** (0.5 * p1 + 0.5 * p2))
    record("PLR + RF hybrid", hybrid)

    say("")
    say(f"--- paired Wilcoxon vs constant baseline (n={len(y)}) ---")
    base = rows["Constant (median log-eta)"]
    for k_, v in rows.items():
        if k_.startswith("Constant"):
            continue
        m = ~np.isnan(v) & ~np.isnan(base)
        try:
            _, p = wilcoxon(v[m], base[m])
        except ValueError:
            p = float("nan")
        tag = "better" if np.nanmedian(v) < np.nanmedian(base) else "WORSE"
        say(f"  {k_:36s} wins {int((v[m]<base[m]).sum()):3d}/{int(m.sum())}  "
            f"p={p:.4g}  ({tag})")

    say("")
    say("--- per-application median ---")
    tb = pd.DataFrame(rows); tb["app"] = df.app.values
    say(tb.groupby("app").median().round(3).to_string())

    tb.to_csv(OUT / "advanced_nn_perrow.csv", index=False)
    (OUT / "advanced_nn_summary.txt").write_text("\n".join(L) + "\n")


if __name__ == "__main__":
    main()