scripts/c5c_derand_validity.py

"""Claim 5 addendum: is the 5-permutation derandomised Wilcoxon still VALID?

The n-sweep in c5b_derandomization.py reproduces the paper's power gain from
derandomisation, but the type-I error rises with the number of permutations.
Theorem 3.3 covers Algorithm 1/2, which draws ONE pair of permutations; the
Rao-Blackwellised variant of Section 5 has no accompanying theory.

This script isolates the question under a clean GLOBAL NULL (y drawn completely
independently of X, so EVERY coordinate is null) with the ORACLE nu_j = rho_j,
i.e. the exact hypothesis of Theorem 3.3.  If the type-I error still inflates
with the number of permutations there, the inflation is a property of the
Rao-Blackwellisation itself, not of the masked-correlation design.
"""

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.ensemble import GradientBoostingRegressor

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

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


def _rep(rep, n, p=30, oracle=True):
    rng = np.random.default_rng(
        SEED0 + 71_000_000 + 733 * rep + 3 * n + (0 if oracle else 5)
    )
    Sigma = ar1_cov(p, 0.6)
    X = rng.standard_normal((n, p)) @ np.linalg.cholesky(Sigma).T
    y = rng.standard_normal(n)  # GLOBAL NULL: y independent of X
    m = GradientBoostingRegressor(random_state=rep).fit(X, y)
    out = {k: [] for k in PERMS}
    for j in range(min(p, 10)):
        if oracle:
            nu = gaussian_nu(X, j, Sigma)  # under the global null rho = nu
            rho = nu.copy()
        else:
            nu, rho = fit_nu_rho(X, y, j, alpha=1.0)
        for k in PERMS:
            out[k].append(
                sko_pvalue(X, y, j, nu, rho, m.predict, sq_loss, rng, n_perm=k)
            )
    return out


def run(n, reps=200, n_jobs=100, oracle=True):
    t0 = time.time()
    got = Parallel(n_jobs=n_jobs)(
        delayed(_rep)(r, n, oracle=oracle) for r in range(reps)
    )
    res = {"n": n, "p": 30, "replicates": reps, "oracle": oracle,
           "seconds": round(time.time() - t0, 1)}
    for k in PERMS:
        pv = np.concatenate([g[k] for g in got])
        ks = stats.kstest(pv, "uniform")
        res[f"perm{k}"] = {
            "n_pvalues": int(len(pv)),
            "type_I_at_0.05": float(np.mean(pv <= 0.05)),
            "type_I_at_0.10": float(np.mean(pv <= 0.10)),
            "mean_pvalue": float(np.mean(pv)),
            "ks_pvalue": float(ks.pvalue),
        }
        print(
            f"[global-null {'oracle' if oracle else 'estim'} n={n}] perms={k:<3} "
            f"typeI@0.05={res[f'perm{k}']['type_I_at_0.05']:.4f} "
            f"typeI@0.10={res[f'perm{k}']['type_I_at_0.10']:.4f} "
            f"KS p={ks.pvalue:.2e}",
            flush=True,
        )
    return res


if __name__ == "__main__":
    os.makedirs(OUT, exist_ok=True)
    res = {
        "seed0": SEED0,
        "permutation_grid": PERMS,
        "design": "global null, oracle nu=rho, GB black box",
    }
    for n in (100, 300):
        res[f"oracle_n{n}"] = run(n, oracle=True)
    for n in (100, 300, 1000):
        res[f"estimated_n{n}"] = run(n, oracle=False)
    with open(os.path.join(OUT, "claim5c_derand_validity.json"), "w") as f:
        json.dump(res, f, indent=2)
    print("wrote", os.path.join(OUT, "claim5c_derand_validity.json"))