scripts/c5b_derandomization.py

"""Claim 5, second half: does 5-permutation derandomisation raise power?

In the paper's own masked-correlation configuration (n = 300, p = 50) our
reproduction saturates: SKO_Wcx, SKO_Wcx_p5 and HRT all reach power 1.000, so
the derandomisation gain cannot be measured there.  This script sweeps the
sample size down into the non-saturated regime and measures the SKO_Wcx ->
SKO_Wcx_p5 power gain (paired over replicates) as a function of n, and also
sweeps the number of permutations at a fixed n.
"""

from __future__ import annotations

import json
import os
import sys
import time

import numpy as np
from joblib import Parallel, delayed
from scipy import stats
from sklearn.neural_network import MLPRegressor

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

OUT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "outputs"
)
SEED0 = 20260725
ALPHA = 0.05
PERMS = [1, 2, 5, 10, 20]


def _rep(rep, n, p=50):
    rng = np.random.default_rng(SEED0 + 61_000_000 + 1231 * rep + 13 * n)
    X, y, l, supp = gen_masked(n, p, rng)
    truth = np.zeros(p, bool)
    truth[supp] = True
    m = MLPRegressor(random_state=rep).fit(X, y)
    P = {k: np.ones(p) for k in PERMS}
    for j in range(p):
        nu, rho = fit_nu_rho(X, y, j, alpha=1.0)
        for k in PERMS:
            P[k][j] = sko_pvalue(X, y, j, nu, rho, m.predict, sq_loss, rng, n_perm=k)
    return {
        k: (float(np.mean(P[k][truth] <= ALPHA)), float(np.mean(P[k][~truth] <= ALPHA)))
        for k in PERMS
    }


def run(n, reps=100, n_jobs=100):
    t0 = time.time()
    got = Parallel(n_jobs=n_jobs)(delayed(_rep)(r, n) for r in range(reps))
    res = {
        "n": n,
        "p": 50,
        "replicates": reps,
        "alpha": ALPHA,
        "seconds": round(time.time() - t0, 1),
    }
    base = np.array([g[1][0] for g in got])
    for k in PERMS:
        pw = np.array([g[k][0] for g in got])
        t1 = np.array([g[k][1] for g in got])
        res[f"perm{k}"] = {
            "power": float(pw.mean()),
            "power_se": float(pw.std(ddof=1) / np.sqrt(reps)),
            "type_I": float(t1.mean()),
            "type_I_se": float(t1.std(ddof=1) / np.sqrt(reps)),
        }
        if k > 1:
            d = pw - base
            res[f"perm{k}"]["paired_gain_over_perm1"] = float(d.mean())
            res[f"perm{k}"]["paired_gain_pvalue"] = float(
                stats.ttest_rel(pw, base).pvalue
            )
        print(
            f"[n={n}] perms={k:<3} power={pw.mean():.3f} typeI={t1.mean():.3f}"
            + (
                f"  gain vs 1 perm = {(pw-base).mean():+.4f} "
                f"(p={stats.ttest_rel(pw, base).pvalue:.2e})"
                if k > 1
                else ""
            ),
            flush=True,
        )
    return res


if __name__ == "__main__":
    os.makedirs(OUT, exist_ok=True)
    res = {
        "seed0": SEED0,
        "setting": "masked correlation (Figure 5), NN black box",
        "permutation_grid": PERMS,
    }
    for n in (50, 75, 100, 150, 300):
        res[f"n{n}"] = run(n)
    with open(os.path.join(OUT, "claim5b_derandomization.json"), "w") as f:
        json.dump(res, f, indent=2)
    print("wrote", os.path.join(OUT, "claim5b_derandomization.json"))