scripts/c2_fdr.py

"""Claim 2 (Theorem 3.4): FDR(S_SKO) <= q.

Decisive design: a known ground-truth support, MANY replicates, and the realised
FDR (mean FDP over replicates) compared with the nominal q at several q values.

Two regimes:
  * ORACLE  (Algorithm 3, the exact hypothesis "Given nu_j and rho_j"):
    Gaussian design with KNOWN Sigma and beta, so nu_j and rho_j are the exact
    closed-form conditional expectations.  If Theorem 3.4 is true, FDR <= q must
    hold here for EVERY q, at finite n, for ANY model m.
  * ESTIMATED (Algorithm 4): nu_j, rho_j are ridge regressions fitted on the
    same n samples.  Theorem 3.4 does not formally cover this case (Sections
    4.2-4.5 argue it approximately holds).

Boundary probes:
  * a deliberately mis-specified / nonlinear design where the ridge imputer is
    biased (rho_hat != nu_hat under the null) - does estimated SKO break?
  * the "BH on Wilcoxon p-values" alternative the paper says should be worse.
"""

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 Lasso

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

OUT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "outputs"
)
SEED0 = 20260725
QS = [0.05, 0.1, 0.2, 0.3, 0.5]


def _make(rep, n, p, rho, k, tag):
    rng = np.random.default_rng(SEED0 + 3_000_000 + 1009 * rep + 7 * tag)
    Sigma = ar1_cov(p, rho)
    L = np.linalg.cholesky(Sigma)
    X = rng.standard_normal((n, p)) @ L.T
    beta = np.zeros(p)
    supp = np.arange(k)
    beta[supp] = rng.uniform(1.0, 2.0, size=k)
    sigma = 1.0
    y = X @ beta + sigma * rng.standard_normal(n)
    return rng, X, y, beta, Sigma, sigma, supp


def _model(kind, X, y, rep):
    if kind == "gb":
        return GradientBoostingRegressor(random_state=rep).fit(X, y).predict
    if kind == "lasso":
        return Lasso(alpha=0.05).fit(X, y).predict
    raise ValueError(kind)


def _rep_oracle(rep, n, p, k, kind):
    rng, X, y, beta, Sigma, sigma, supp = _make(rep, n, p, 0.6, k, 0)
    pred = _model(kind, X, y, rep)
    W = np.empty(p)
    for j in range(p):
        nu = gaussian_nu(X, j, Sigma)
        rh = gaussian_rho(X, y, j, Sigma, beta, sigma)
        W[j] = sko_statistic(X, y, j, nu, rh, pred, sq_loss, rng)
    return [
        (fdp(knockoff_select(W, q), supp), power(knockoff_select(W, q), supp))
        for q in QS
    ]


def _rep_estimated(rep, n, p, k, kind, alpha=1.0, nonlinear=False):
    rng, X, y, beta, Sigma, sigma, supp = _make(rep, n, p, 0.6, k, 1)
    if nonlinear:
        # break the linear-imputer specification: X^j depends nonlinearly on X^-j
        Z = X.copy()
        X = np.sign(Z) * np.abs(Z) ** 1.5
        y = Z @ beta + sigma * rng.standard_normal(n)
    pred = _model(kind, X, y, rep)
    W = np.empty(p)
    for j in range(p):
        nu, rh = fit_nu_rho(X, y, j, alpha=alpha)
        W[j] = sko_statistic(X, y, j, nu, rh, pred, sq_loss, rng)
    return [
        (fdp(knockoff_select(W, q), supp), power(knockoff_select(W, q), supp))
        for q in QS
    ]


def _rep_bh(rep, n, p, k, kind):
    """Wilcoxon p-values + Benjamini-Hochberg, the alternative route."""
    rng, X, y, beta, Sigma, sigma, supp = _make(rep, n, p, 0.6, k, 2)
    pred = _model(kind, X, y, rep)
    pv = np.empty(p)
    for j in range(p):
        nu = gaussian_nu(X, j, Sigma)
        rh = gaussian_rho(X, y, j, Sigma, beta, sigma)
        pv[j] = sko_pvalue(X, y, j, nu, rh, pred, sq_loss, rng)
    out = []
    order = np.argsort(pv)
    ranked = pv[order]
    for q in QS:
        thr = ranked <= q * np.arange(1, p + 1) / p
        nsel = np.max(np.where(thr)[0]) + 1 if thr.any() else 0
        sel = order[:nsel]
        out.append((fdp(sel, supp), power(sel, supp)))
    return out


def summarise(rows, label):
    arr = np.array(rows)  # (reps, len(QS), 2)
    res = {}
    for i, q in enumerate(QS):
        f = arr[:, i, 0]
        pw = arr[:, i, 1]
        se = float(np.std(f, ddof=1) / np.sqrt(len(f)))
        res[str(q)] = {
            "FDR": float(np.mean(f)),
            "FDR_se": se,
            "FDR_upper95": float(np.mean(f) + 1.96 * se),
            "power": float(np.nanmean(pw)),
            "controlled": bool(np.mean(f) <= q),
        }
        print(
            f"[{label}] q={q:<5} FDR={np.mean(f):.4f} (+-{1.96*se:.4f}) "
            f"power={np.nanmean(pw):.3f} ctrl={np.mean(f) <= q}",
            flush=True,
        )
    res["replicates"] = int(arr.shape[0])
    return res


def run(fn, reps, n_jobs, label, **kw):
    t0 = time.time()
    rows = Parallel(n_jobs=n_jobs)(delayed(fn)(r, **kw) for r in range(reps))
    out = summarise(rows, label)
    out["seconds"] = round(time.time() - t0, 1)
    return out


if __name__ == "__main__":
    os.makedirs(OUT, exist_ok=True)
    res = {"seed0": SEED0, "q_grid": QS, "n": 300, "p": 50, "k_support": 12}
    res["oracle_gb"] = run(
        _rep_oracle, 500, 100, "oracle-GB", n=300, p=50, k=12, kind="gb"
    )
    res["oracle_lasso"] = run(
        _rep_oracle, 500, 100, "oracle-Lasso", n=300, p=50, k=12, kind="lasso"
    )
    res["estimated_gb"] = run(
        _rep_estimated, 500, 100, "estim-GB", n=300, p=50, k=12, kind="gb"
    )
    res["estimated_gb_nonlinear"] = run(
        _rep_estimated,
        300,
        100,
        "estim-GB-nonlin",
        n=300,
        p=50,
        k=12,
        kind="gb",
        nonlinear=True,
    )
    res["oracle_gb_BH"] = run(
        _rep_bh, 300, 100, "oracle-GB-BH", n=300, p=50, k=12, kind="gb"
    )
    res["oracle_gb_global_null"] = run(
        _rep_oracle, 300, 100, "oracle-GB-globalnull", n=300, p=50, k=0, kind="gb"
    )
    with open(os.path.join(OUT, "claim2_fdr.json"), "w") as f:
        json.dump(res, f, indent=2)
    print("wrote", os.path.join(OUT, "claim2_fdr.json"))