scripts/c1_pvalue_validity.py

"""Claim 1 (Theorem 3.3).

"Semi-knockoffs avoids the train-test data split required by prior CIT methods
such as HRT while still yielding valid p-values, via nonparametric paired tests
requiring only conditional expectations nu_j and rho_j rather than exact
knockoff construction."

Four independent tests:

 A. ORACLE validity + resampling-scheme audit.  Gaussian linear design, nu_j and
    rho_j in CLOSED FORM (the exact hypothesis of Theorem 3.3).  Null p-values of
    the Wilcoxon and sign-test versions collected over many replicates and tested
    for uniformity (one-sample KS).  Swept over n in {30, 100, 300}.
    Boundary probe: Algorithms 1-4 say "sample pi_{j,1}, pi_{j,2} PERMUTATIONS",
    while the proof in Appendix E.2 uses i.i.d. uniform indices U_1..U_n.  Under
    permutations the n paired differences are NOT independent, which the Wilcoxon
    signed-rank / sign test formally requires.  Both schemes are measured.

 B. NO-SPLIT vs SPLIT, everything ESTIMATED (Algorithm 1).  Same data, same
    gradient-boosting architecture:
       - SKO-Wcx, model trained on ALL n samples, tested on those same n samples;
       - HRT with NO split (model + sampler fit on all n, test on all n);
       - HRT done properly with a 50/50 split (df-corrected residual sampler);
       - HRT with a 50/50 split and the EXACT Gaussian conditional sampler.

 C. "ONLY CONDITIONAL EXPECTATIONS, NOT EXACT KNOCKOFFS".  Heavy-tailed
    multivariate-t(3) elliptical design.  There the conditional MEAN is still
    linear but the conditional distribution is heteroscedastic and non-Gaussian,
    so a Gaussian model-X knockoff / sampler is mis-specified.  SKO should stay
    valid (its exchangeability argument only needs rho = nu under the null);
    a Gaussian-sampler HRT should not.
"""

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
from sklearn.linear_model import Ridge

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

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


def _ks(pv):
    return stats.kstest(pv, "uniform")


def _pack(pv, secs=None):
    ks = _ks(pv)
    d = {
        "n_pvalues": int(len(pv)),
        "ks_stat": float(ks.statistic),
        "ks_pvalue": float(ks.pvalue),
        "mean_pvalue": float(np.mean(pv)),
        "type_I_at_0.05": float(np.mean(pv <= 0.05)),
        "type_I_at_0.10": float(np.mean(pv <= 0.10)),
        "type_I_at_0.20": float(np.mean(pv <= 0.20)),
    }
    if secs is not None:
        d["seconds"] = round(secs, 1)
    return d


# ---------------------------------------------------------------- experiment A
def _oracle_rep(rep, n, p, rho, test, scheme):
    tag = {"wilcoxon": 0, "sign": 1}[test] * 2 + {"perm": 0, "iid": 1}[scheme]
    rng = np.random.default_rng(SEED0 + 1000 * rep + 137 * tag + 31 * n)
    Sigma = ar1_cov(p, rho)
    L = np.linalg.cholesky(Sigma)
    X = rng.standard_normal((n, p)) @ L.T
    beta = np.zeros(p)
    beta[0] = 1.5
    beta[1] = 1.0  # supp = {0, 1}; every other j is null
    sigma = 1.0
    y = X @ beta + sigma * rng.standard_normal(n)

    # black-box pre-trained model, trained on ALL n samples (no split)
    model = GradientBoostingRegressor(random_state=rep).fit(X, y)

    out = []
    for j in range(2, min(p, 12)):  # null coordinates only
        nu = gaussian_nu(X, j, Sigma)
        rhoj = gaussian_rho(X, y, j, Sigma, beta, sigma)
        out.append(
            sko_pvalue(
                X,
                y,
                j,
                nu,
                rhoj,
                model.predict,
                sq_loss,
                rng,
                test=test,
                scheme=scheme,
            )
        )
    return out


def exp_A(p=50, reps=200, n_jobs=100):
    res = {}
    for n in (30, 100, 300):
        for test in ("wilcoxon", "sign"):
            for scheme in ("perm", "iid"):
                t0 = time.time()
                lists = Parallel(n_jobs=n_jobs)(
                    delayed(_oracle_rep)(r, n, p, 0.6, test, scheme)
                    for r in range(reps)
                )
                key = f"n{n}|{test}|{scheme}"
                res[key] = _pack(np.concatenate(lists), time.time() - t0)
                print(f"[A] {key:22s} {res[key]}", flush=True)
    return res


# ---------------------------------------------------------------- experiment B
def _split_rep(rep, n, p):
    rng = np.random.default_rng(SEED0 + 7_000_000 + rep)
    X, y, beta, Sigma, supp = gen_adjacent(n, p, rng)
    nulls = np.array([j for j in range(p) if beta[j] == 0])[:12]

    m_full = GradientBoostingRegressor(random_state=rep).fit(X, y)
    idx = rng.permutation(n)
    tr, te = idx[: n // 2], idx[n // 2 :]
    m_split = GradientBoostingRegressor(random_state=rep).fit(X[tr], y[tr])
    r2_split = float(1 - np.mean((m_split.predict(X[te]) - y[te]) ** 2) / np.var(y[te]))

    rows = []
    for j in nulls:
        cols = [k for k in range(p) if k != j]
        nu, rhoj = fit_nu_rho(X, y, j, alpha=1.0)
        p_sko = sko_pvalue(X, y, j, nu, rhoj, m_full.predict, sq_loss, rng)

        # HRT with NO split at all
        res_full = X[:, j] - nu
        p_hrt_nosplit = hrt_pvalue(
            X, y, j, nu, res_full, m_full.predict, sq_loss, rng, K=200
        )

        # HRT with a proper 50/50 split; residuals df-corrected so the resampled
        # conditional variance is unbiased (n_tr - p degrees of freedom).
        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))
        nu_te = rg.predict(X[np.ix_(te, cols)])
        p_hrt_split = hrt_pvalue(
            X[te], y[te], j, nu_te, res_tr * dfc, m_split.predict, sq_loss, rng, K=200
        )

        # HRT with a split AND the exact Gaussian conditional sampler
        nu_or = gaussian_nu(X, j, Sigma)[te]
        S_mj = Sigma[np.ix_(cols, cols)]
        cvar = float(
            Sigma[j, j] - Sigma[j, cols] @ np.linalg.solve(S_mj, Sigma[cols, j])
        )
        pool = np.sqrt(cvar) * rng.standard_normal(4000)
        p_hrt_split_or = hrt_pvalue(
            X[te], y[te], j, nu_or, pool, m_split.predict, sq_loss, rng, K=200
        )
        rows.append((p_sko, p_hrt_nosplit, p_hrt_split, p_hrt_split_or))
    return rows, r2_split


def exp_B(n=300, p=50, reps=100, n_jobs=100):
    t0 = time.time()
    got = Parallel(n_jobs=n_jobs)(delayed(_split_rep)(r, n, p) for r in range(reps))
    rows = np.array([r for g in got for r in g[0]])
    res = {
        "R2_with_split_heldout": float(np.mean([g[1] for g in got])),
        "n_null_pvalues": int(rows.shape[0]),
        "replicates": reps,
        "seconds": round(time.time() - t0, 1),
    }
    for k, nm in enumerate(
        ["SKO_Wcx_nosplit", "HRT_nosplit", "HRT_split", "HRT_split_oracle_sampler"]
    ):
        res[nm] = _pack(rows[:, k])
        print(f"[B] {nm:26s} {res[nm]}", flush=True)
    return res


# ---------------------------------------------------------------- experiment C
def _t_rep(rep, n, p, df=3.0):
    """Elliptical multivariate-t(df) design: conditional mean still linear,
    conditional law heteroscedastic and non-Gaussian."""
    rng = np.random.default_rng(SEED0 + 5_000_000 + rep)
    Sigma = ar1_cov(p, 0.6)
    L = np.linalg.cholesky(Sigma)
    g = rng.standard_normal((n, p)) @ L.T
    w = np.sqrt(df / rng.chisquare(df, size=n))[:, None]
    X = g * w  # multivariate t_df(0, Sigma)
    beta = np.zeros(p)
    beta[0], beta[1] = 1.5, 1.0
    y = X @ beta + rng.standard_normal(n)
    m_full = GradientBoostingRegressor(random_state=rep).fit(X, y)
    idx = rng.permutation(n)
    tr, te = idx[: n // 2], idx[n // 2 :]
    m_split = GradientBoostingRegressor(random_state=rep).fit(X[tr], y[tr])

    out = []
    for j in range(2, 12):
        cols = [k for k in range(p) if k != j]
        # ORACLE nu for an elliptical law: same linear map as the Gaussian case.
        nu = gaussian_nu(X, j, Sigma)
        rhoj = nu.copy()  # under the null rho == nu exactly
        p_sko = sko_pvalue(X, y, j, nu, rhoj, m_full.predict, sq_loss, rng)

        # HRT with a split and a mis-specified GAUSSIAN conditional sampler
        S_mj = Sigma[np.ix_(cols, cols)]
        cvar = float(
            Sigma[j, j] - Sigma[j, cols] @ np.linalg.solve(S_mj, Sigma[cols, j])
        )
        pool = np.sqrt(cvar) * rng.standard_normal(4000)
        p_hrt = hrt_pvalue(
            X[te], y[te], j, nu[te], pool, m_split.predict, sq_loss, rng, K=200
        )
        out.append((p_sko, p_hrt))
    return out


def exp_C(n=300, p=50, reps=200, n_jobs=100):
    t0 = time.time()
    got = Parallel(n_jobs=n_jobs)(delayed(_t_rep)(r, n, p) for r in range(reps))
    rows = np.array([r for g in got for r in g])
    res = {
        "design": "multivariate t(3), AR(1) rho=0.6",
        "replicates": reps,
        "seconds": round(time.time() - t0, 1),
    }
    for k, nm in enumerate(["SKO_Wcx_oracle_nu", "HRT_split_gaussian_sampler"]):
        res[nm] = _pack(rows[:, k])
        print(f"[C] {nm:28s} {res[nm]}", flush=True)
    return res


if __name__ == "__main__":
    os.makedirs(OUT, exist_ok=True)
    result = {
        "seed0": SEED0,
        "A_oracle_uniformity": exp_A(),
        "B_split_vs_nosplit": exp_B(),
        "C_heavy_tailed_elliptical": exp_C(),
    }
    with open(os.path.join(OUT, "claim1_pvalue_validity.json"), "w") as f:
        json.dump(result, f, indent=2)
    print("wrote", os.path.join(OUT, "claim1_pvalue_validity.json"))