scripts/c5_power_hrt.py

"""Claim 5 (Figures 4 and 5, Section 5.1).

"On simulated data with adjacent-feature support, Semi-knockoffs maintains
type-I error control while achieving higher power than HRT, and derandomization
with 5 permutations under masked correlation further increases power."

Two settings, both n = 300, p = 50, 50 replicates (the paper's scale):

 Figure 4 - ADJACENT SUPPORT.  X ~ N(0, Sigma), Sigma_ij = 0.6^|i-j|,
   y = beta'X + eps, the first 0.25p coordinates of beta in [1, 2], rest zero,
   eps ~ N(0, 1).  Pre-trained black box: gradient boosting.

 Figure 5 - MASKED CORRELATION.  X ~ N(0, Sigma), Sigma_ij = 0.6^|i-j|.  One
   relevant coordinate l, y = X_l + 0.5 eps1.  A correlated NULL variable is
   built as X_{l-1} = X_l + 0.5 eps2.  Pre-trained black box: neural network.

Methods compared (all at alpha = 0.05 on the p-value scale):
   HRT          - Tansey et al. (2022), 50/50 train/test split, K = 200 draws
   SKO_Wcx      - Algorithm 1, single permutation, NO split
   SKO_Wcx_p5   - Algorithm 1 Rao-Blackwellised over 5 permutations, NO split
   LOCO_Wcx     - refit-without-feature baseline with a Wilcoxon test on a
                  held-out split (a "VIM"-style comparator)
Reported: power, type-I error, AUC - the three panels of Figures 4 and 5.
"""

from __future__ import annotations

import json
import os
import sys
import time

import numpy as np
from joblib import Parallel, delayed
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.linear_model import Ridge
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import cross_val_score
from sklearn.neural_network import MLPRegressor
from scipy import stats

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__))))
from semiknockoffs import (  # noqa: E402
    fit_nu_rho,
    gen_adjacent,
    gen_masked,
    hrt_pvalue,
    sko_pvalue,
    sq_loss,
)

OUT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "outputs"
)
SEED0 = 20260725
ALPHA = 0.05


def _blank(kind, seed):
    if kind == "gb":
        return GradientBoostingRegressor(random_state=seed)
    if kind == "nn":
        return MLPRegressor(random_state=seed)
    raise ValueError(kind)


def _fit(kind, X, y, seed):
    if kind == "gb":
        return GradientBoostingRegressor(random_state=seed).fit(X, y)
    if kind == "nn":
        # scikit-learn defaults, which reproduce the paper's reported
        # R2 = 0.567 (split) / 0.581 (no split) for the masked setting.
        return MLPRegressor(random_state=seed).fit(X, y)
    raise ValueError(kind)


def _rep(rep, setting, kind, n=300, p=50, K=200):
    rng = np.random.default_rng(
        SEED0 + 41_000_000 + 1231 * rep + (0 if setting == "adjacent" else 7)
    )
    if setting == "adjacent":
        X, y, beta, Sigma, supp = gen_adjacent(n, p, rng)
    else:
        X, y, l, supp = gen_masked(n, p, rng)
    truth = np.zeros(p, bool)
    truth[supp] = True

    m_full = _fit(kind, X, y, rep)
    r2_full = float(
        np.mean(cross_val_score(_blank(kind, rep), X, y, cv=5, scoring="r2"))
    )  # generalisation R2,
    #   the closest comparable to the paper's "R2 without split"
    idx = rng.permutation(n)
    tr, te = idx[: n // 2], idx[n // 2 :]
    m_split = _fit(kind, X[tr], y[tr], rep)
    r2_split = float(1 - np.mean((m_split.predict(X[te]) - y[te]) ** 2) / np.var(y[te]))

    P = {k: np.ones(p) for k in ("HRT", "SKO_Wcx", "SKO_Wcx_p5", "LOCO_Wcx")}
    for j in range(p):
        cols = [k for k in range(p) if k != j]
        nu, rho = fit_nu_rho(X, y, j, alpha=1.0)
        P["SKO_Wcx"][j] = sko_pvalue(X, y, j, nu, rho, m_full.predict, sq_loss, rng)
        P["SKO_Wcx_p5"][j] = sko_pvalue(
            X, y, j, nu, rho, m_full.predict, sq_loss, rng, n_perm=5
        )

        # HRT: needs the split
        rg = Ridge(alpha=1.0).fit(X[np.ix_(tr, cols)], X[tr, j])
        res_tr = X[tr, j] - rg.predict(X[np.ix_(tr, cols)])
        dfc = np.sqrt(len(tr) / max(1.0, len(tr) - p))
        P["HRT"][j] = hrt_pvalue(
            X[te],
            y[te],
            j,
            rg.predict(X[np.ix_(te, cols)]),
            res_tr * dfc,
            m_split.predict,
            sq_loss,
            rng,
            K=K,
        )

        # LOCO with a Wilcoxon test: refit the model without feature j on the
        # train half, compare per-observation losses on the test half.
        m_red = _fit(kind, X[np.ix_(tr, cols)], y[tr], rep)
        d = sq_loss(m_red.predict(X[np.ix_(te, cols)]), y[te]) - sq_loss(
            m_split.predict(X[te]), y[te]
        )
        nz = d[d != 0]
        P["LOCO_Wcx"][j] = (
            float(stats.wilcoxon(nz, alternative="greater").pvalue) if len(nz) else 1.0
        )

    out = {"R2_no_split_cv5": r2_full, "R2_split_heldout": r2_split}
    for k, pv in P.items():
        out[k] = {
            "power": float(np.mean(pv[truth] <= ALPHA)),
            "type_I": float(np.mean(pv[~truth] <= ALPHA)),
            "auc": float(roc_auc_score(truth.astype(int), -pv)),
        }
    return out


def run(setting, kind, reps=50, n_jobs=50):
    t0 = time.time()
    got = Parallel(n_jobs=n_jobs)(delayed(_rep)(r, setting, kind) for r in range(reps))
    res = {
        "setting": setting,
        "model": kind,
        "replicates": reps,
        "n": 300,
        "p": 50,
        "alpha": ALPHA,
        "R2_no_split_cv5": float(
            np.mean([g["R2_no_split_cv5"] for g in got])
        ),
        "R2_split_heldout": float(np.mean([g["R2_split_heldout"] for g in got])),
        "seconds": round(time.time() - t0, 1),
    }
    for k in ("HRT", "SKO_Wcx", "SKO_Wcx_p5", "LOCO_Wcx"):
        for met in ("power", "type_I", "auc"):
            v = np.array([g[k][met] for g in got])
            res[f"{k}|{met}"] = float(np.mean(v))
            res[f"{k}|{met}_se"] = float(np.std(v, ddof=1) / np.sqrt(len(v)))
        print(
            f"[{setting}/{kind}] {k:11s} power={res[k+'|power']:.3f} "
            f"typeI={res[k+'|type_I']:.3f} auc={res[k+'|auc']:.3f}",
            flush=True,
        )
    # paired comparisons on the same replicates
    for a, b in (("SKO_Wcx", "HRT"), ("SKO_Wcx_p5", "SKO_Wcx"), ("SKO_Wcx_p5", "HRT")):
        da = np.array([g[a]["power"] for g in got])
        db = np.array([g[b]["power"] for g in got])
        d = da - db
        t = stats.ttest_rel(da, db)
        res[f"paired_power_{a}_minus_{b}"] = {
            "mean_diff": float(d.mean()),
            "se": float(np.std(d, ddof=1) / np.sqrt(len(d))),
            "t_pvalue": float(t.pvalue),
        }
        print(
            f"    paired power {a} - {b} = {d.mean():+.4f} " f"(p={t.pvalue:.2e})",
            flush=True,
        )
    return res


if __name__ == "__main__":
    os.makedirs(OUT, exist_ok=True)
    res = {"seed0": SEED0}
    res["figure4_adjacent_gb"] = run("adjacent", "gb")
    res["figure5_masked_nn"] = run("masked", "nn")
    res["figure4_adjacent_nn"] = run("adjacent", "nn")
    res["figure5_masked_gb"] = run("masked", "gb")
    with open(os.path.join(OUT, "claim5_power.json"), "w") as f:
        json.dump(res, f, indent=2)
    print("wrote", os.path.join(OUT, "claim5_power.json"))