hardware-counters

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

Implements the techniques the tabular-deep-learning literature identifies as
the ones that actually close the gap to gradient-boosted trees, and evaluates
them under the SAME corrected protocol as retrain2.py (merged configurations,
median-constant baseline, leave-one-application-out).

Techniques, with sources:

  PLR numerical embeddings -- Gorishniy et al., "On Embeddings for Numerical
      Features in Tabular Deep Learning" (NeurIPS 2022, arXiv:2203.05556).
      Each scalar x becomes concat[sin(v), cos(v)] with v = 2*pi*c*x and
      trainable frequencies c ~ N(0, sigma^2), followed by Linear+ReLU. This
      gives an MLP the high-frequency capacity it structurally lacks, which is
      the mechanism Grinsztajn et al. (NeurIPS 2022) identify as the reason
      trees beat MLPs on non-smooth tabular targets. sigma is the critical
      hyperparameter and is tuned here.

  Robust scaling + smooth clipping -- Holzmuller et al., "Better by Default"
      (NeurIPS 2024). Replaces standardisation/log1p. Scales by median and IQR
      so outliers cannot distort the inlier scale, then squashes tails with
      tanh instead of truncating. Relevant here because four features have
      |skew| > 2 and max/min up to 5.7e4.

  Deep ensembling -- Gorishniy et al., "TabM" (arXiv:2410.24210). Individual
      members are weak and overfitted; the average generalises. Already the
      best-performing tail model in earlier runs.

  RealMLP -- Holzmuller et al. (NeurIPS 2024), via pytabkit. A pre-tuned MLP
      combining robust scaling, PBLD embeddings, a diagonal weight layer, NTP
      parametrisation and Mish activation, with defaults meta-learned on a
      benchmark suite rather than tuned per dataset.

Every learned model is selected by an INNER leave-one-application-out loop over
the training applications only, so no hyperparameter sees the evaluation fold.
"""
import warnings, sys, json
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(16)
D      = "/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


# ----------------------------------------------------------------- scaling
class RobustSmoothScaler:
    """
    Robust scaling + smooth clipping (Holzmuller et al., NeurIPS 2024).

    x' = (x - median) / (IQR/1.349)      -- IQR/1.349 approximates sigma for a
                                            normal, so the scale matches
                                            StandardScaler on clean data but is
                                            immune to outliers
    x'' = 3 * tanh(x' / 3)               -- smooth clip: linear near 0, tails
                                            squashed to +/-3 without the hard
                                            discontinuity of clipping
    """
    def fit(self, X):
        X = np.asarray(X, float)
        self.med_ = np.nanmedian(X, axis=0)
        q1, q3 = np.nanpercentile(X, [25, 75], axis=0)
        iqr = (q3 - q1) / 1.349
        # fall back to std, then 1.0, for constant / near-constant columns
        sd = np.nanstd(X, axis=0)
        iqr = np.where(iqr > 1e-12, iqr, np.where(sd > 1e-12, sd, 1.0))
        self.scale_ = iqr
        return self

    def transform(self, X):
        z = (np.asarray(X, float) - self.med_) / self.scale_
        return 3.0 * np.tanh(z / 3.0)

    def fit_transform(self, X):
        return self.fit(X).transform(X)


# ----------------------------------------------------------------- PLR
class PLREmbedding(nn.Module):
    """
    Periodic-Linear-ReLU embedding (Gorishniy et al., NeurIPS 2022, eq. 2).

    For each of n_features scalars x_i:
        v_i   = 2*pi * c_i * x_i          c_i in R^k, trainable, init N(0, s^2)
        p_i   = [sin(v_i), cos(v_i)]      in R^{2k}
        e_i   = ReLU(Linear(p_i))         in R^{d}
    The per-feature embeddings are concatenated into R^{n*d}.
    """
    def __init__(self, n_features, k=24, sigma=0.05, d=16):
        super().__init__()
        self.coeffs = nn.Parameter(torch.randn(n_features, k) * sigma)
        self.lin = nn.Linear(2 * k, d)
        self.n, self.d = n_features, d

    def forward(self, x):                        # x: (B, n)
        v = 2 * np.pi * x.unsqueeze(-1) * self.coeffs     # (B, n, k)
        p = torch.cat([torch.sin(v), torch.cos(v)], dim=-1)
        return torch.relu(self.lin(p)).flatten(1)         # (B, n*d)


class PLRNet(nn.Module):
    def __init__(self, n_features, k=24, sigma=0.05, d=16, hidden=(128, 64),
                 dropout=0.1):
        super().__init__()
        self.emb = PLREmbedding(n_features, k, sigma, d)
        layers, prev = [], n_features * d
        for h in hidden:
            layers += [nn.Linear(prev, h), nn.Mish(), nn.Dropout(dropout)]
            prev = h
        layers.append(nn.Linear(prev, 1))
        self.mlp = nn.Sequential(*layers)

    def forward(self, x):
        return self.mlp(self.emb(x)).squeeze(-1)


def train_plr(Xtr, ytr, Xte, k=24, sigma=0.05, d=16, hidden=(128, 64),
              wd=1e-2, epochs=400, lr=3e-3, seed=0, n_ens=5, dropout=0.1):
    """Deep ensemble of PLR nets; returns the mean prediction (TabM rationale)."""
    Xtr_t = torch.tensor(Xtr, dtype=torch.float32)
    ytr_t = torch.tensor(ytr, dtype=torch.float32)
    Xte_t = torch.tensor(Xte, dtype=torch.float32)
    preds = []
    for m in range(n_ens):
        torch.manual_seed(seed * 100 + m)
        net = PLRNet(Xtr.shape[1], k, sigma, d, hidden, dropout)
        opt = torch.optim.AdamW(net.parameters(), lr=lr, weight_decay=wd)
        sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=lr,
                                                    total_steps=epochs)
        net.train()
        for _ in range(epochs):
            opt.zero_grad()
            loss = nn.functional.mse_loss(net(Xtr_t), ytr_t)
            loss.backward()
            opt.step(); sched.step()
        net.eval()
        with torch.no_grad():
            preds.append(net(Xte_t).numpy())
    return np.mean(preds, axis=0)


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


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))


# PLR hyperparameter grid; sigma is the critical one per the paper
PLR_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))]


def main():
    df, X, y, t_an, eta = prep()
    apps = sorted(df.app.unique())
    tgt  = np.log10(eta)
    Xv   = 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("Leave-one-application-out; error = max(pred/act, act/pred)")
    say("")

    def record(name, fn):
        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:38s} median {np.nanmedian(r):.4f}   "
            f"p90 {np.nanpercentile(r,90):.4f}")
        return r

    # ---- references -----------------------------------------------------
    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)
        m = RandomForestRegressor(n_estimators=400, min_samples_leaf=2,
                                  random_state=0, n_jobs=-1).fit(s1, ytr)
        return ta / (10 ** m.predict(s2))
    record("Random Forest", rf)

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

    # ---- PLR with NESTED hyperparameter selection ------------------------
    picks = []
    def plr_nested(Xtr, ytr, Xte, ta, ap):
        best, bg = None, None
        for g in PLR_GRID:
            errs = []
            for ia in np.unique(ap):                 # inner LOAO, train apps only
                ite = (ap == ia)
                if ite.sum() == 0 or (~ite).sum() < 10:
                    continue
                s1, s2 = scaled(Xtr[~ite], Xtr[ite])
                p = train_plr(s1, ytr[~ite], s2, n_ens=2, epochs=250, **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"{bg['sigma']}/{bg['k']}/{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 (sigma/k/width): {', '.join(picks)}")

    # ---- PLR at a fixed sensible default (no per-fold 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)

    # ---- RealMLP (pre-tuned defaults, no per-dataset tuning) -------------
    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",
                                     n_threads=16, verbosity=0)
            m.fit(s1, ytr)
            return ta / (10 ** np.asarray(m.predict(s2)).ravel())
        record("RealMLP-TD (pre-tuned defaults)", realmlp)
    except Exception as e:
        say(f"  RealMLP unavailable: {type(e).__name__}: {e}")

    # ---- hybrid: PLR + trees --------------------------------------------
    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)

    # ---- significance ----------------------------------------------------
    say("")
    say(f"--- paired Wilcoxon vs constant baseline (n={len(y)} configs) ---")
    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_:38s} wins {int((v[m]<base[m]).sum()):3d}/{int(m.sum())}  "
            f"p={p:.4g}  ({tag})")

    say("")
    say("--- per-application median error factor ---")
    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()