scripts/c3_stability.py

"""Claim 3 (Theorem 4.1, optimization stability).

"For null (non-relevant) features, regularized empirical risk minimizers trained
with and without the feature remain close, with a ||theta_tilde^j - theta_hat||_2
<= O_P(sqrt(log(1/delta)/n)) bound."

Setup exactly as in Section 4.2 / Eq. (2):
   R_n(theta) = (1/n) sum_i l(theta' chi_i, z_i) + lambda ||theta||^2
   theta_hat        = argmin over R^p
   theta_hat^{-j}   = argmin over R^{p-1} (coordinate j dropped)
   theta_tilde^j    = (0, theta_hat^{-j}) re-embedded in R^p
(sklearn Ridge minimises ||Xw-y||^2 + alpha||w||^2, so alpha = n * lambda.)

Five independent tests:

 1. FIGURE 1 reproduction (n=300, p=50, Sigma_ij = 0.6^|i-j|, beta 0.25-sparse in
    blocks of 5, noise ||chi beta||/2), Ridge and Lasso, null vs important j.
 2. RATE IN n: fit  log Q_{1-delta}(||theta_tilde^j - theta_hat||_2) = a log n + c
    over n in [50, 3200] for null j.  Theorem 4.1 predicts a = -1/2.
 3. RATE IN delta at fixed n.  Theorem 4.1 predicts Q ~ sqrt(log(1/delta)),
    i.e. exponent +1/2 on log(1/delta).
 4. BOUNDARY PROBE on the assumptions.  The proof bounds
    ||theta_hat - theta_tilde|| <= (2/lambda) |grad_j R_n(theta_tilde)| and then
    uses Assumption E.4 with theta* the UNREGULARISED population minimiser.  For
    a fixed lambda > 0 the regularised population minimiser is NOT that theta*,
    so a bias floor of order lambda should appear and the n^{-1/2} decay should
    stall.  We sweep lambda in {1e-1, 1e-2, 1e-3, 1e-6} to see whether it does.
    We also break Assumption E.2 (subgaussian design) with a t_2 design.
 5. THE QUANTITY THE METHOD ACTUALLY USES: coefficients of nu_hat (X^j on X^{-j})
    versus rho_hat (X^j on X^{-j}, y).  Here the "extra uninformative feature" is
    y itself, which is exactly what Section 4.2 invokes.  Null vs important j and
    the rate in n.
"""

from __future__ import annotations

import json
import os
import sys
import time

import numpy as np
from joblib import Parallel, delayed
from sklearn.linear_model import Lasso, Ridge

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

OUT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "outputs"
)
FIGS = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "figs")
SEED0 = 20260725
EPS = 1e-12


def _design(rng, n, p, rho=0.6, heavy=False):
    Sigma = ar1_cov(p, rho)
    L = np.linalg.cholesky(Sigma)
    g = rng.standard_normal((n, p)) @ L.T
    if heavy:  # t_2: NOT subgaussian (Assumption E.2 fails)
        g = g * np.sqrt(2.0 / rng.chisquare(2.0, size=n))[:, None]
    return g


def _beta_blocks(rng, p, sparsity=0.25, block=5):
    k = int(round(sparsity * p))
    nb = max(1, k // block)
    starts = rng.choice(np.arange(0, p - block + 1, block), size=nb, replace=False)
    beta = np.zeros(p)
    supp = []
    for s in starts:
        beta[s : s + block] = rng.uniform(1.0, 2.0, size=block)
        supp += list(range(s, s + block))
    return beta, np.array(sorted(supp))


def _fit(chi, z, learner, lam):
    n = chi.shape[0]
    if learner == "ridge":
        est = Ridge(alpha=n * lam, fit_intercept=False)
    else:
        est = Lasso(alpha=lam, fit_intercept=False, max_iter=50000, tol=1e-8)
    est.fit(chi, z)
    return est.coef_.ravel()


def _drops(chi, z, learner, lam):
    """||theta_tilde^j - theta_hat||_2 for every coordinate j."""
    p = chi.shape[1]
    th = _fit(chi, z, learner, lam)
    out = np.empty(p)
    for j in range(p):
        cols = [k for k in range(p) if k != j]
        tt = np.zeros(p)
        tt[cols] = _fit(chi[:, cols], z, learner, lam)
        out[j] = float(np.linalg.norm(tt - th))
    return out


# ---------------------------------------------------------------- 1. Figure 1
def _fig1_rep(rep, n=300, p=50, lam=1e-2):
    rng = np.random.default_rng(SEED0 + 11_000_000 + rep)
    chi = _design(rng, n, p)
    beta, supp = _beta_blocks(rng, p)
    sig = np.linalg.norm(chi @ beta) / np.sqrt(n) / 2.0  # noise level ||chi beta||/2
    z = chi @ beta + sig * rng.standard_normal(n)
    mask = np.zeros(p, bool)
    mask[supp] = True
    return _drops(chi, z, "ridge", lam), _drops(chi, z, "lasso", lam), mask


def fig1(reps=50, n_jobs=50):
    got = Parallel(n_jobs=n_jobs)(delayed(_fig1_rep)(r) for r in range(reps))
    ridge = np.array([g[0] for g in got])
    lasso = np.array([g[1] for g in got])
    mask = np.array([g[2] for g in got])
    res = {"replicates": reps, "n": 300, "p": 50, "lambda": 1e-2}
    for nm, arr in (("ridge", ridge), ("lasso", lasso)):
        res[nm] = {
            "null_mean": float(arr[~mask].mean()),
            "null_q95": float(np.quantile(arr[~mask], 0.95)),
            "important_mean": float(arr[mask].mean()),
            "important_q05": float(np.quantile(arr[mask], 0.05)),
            "separation_ratio": float(arr[mask].mean() / arr[~mask].mean()),
            "auc_important_over_null": float(
                np.mean(arr[mask].ravel()[:, None] > arr[~mask].ravel()[None, :])
            ),
        }
        print(f"[fig1] {nm}: {res[nm]}", flush=True)
    r, la, mk = _fig1_rep(0)
    res["example"] = {
        "ridge": r.tolist(),
        "lasso": la.tolist(),
        "important": mk.tolist(),
    }
    return res


# ---------------------------------------------------------------- 2/3/4. rates
def _rate_rep(rep, n, p, learner, lam, heavy):
    rng = np.random.default_rng(SEED0 + 13_000_000 + 977 * rep + 7 * n)
    chi = _design(rng, n, p, heavy=heavy)
    beta, supp = _beta_blocks(rng, p)
    sig = np.linalg.norm(chi @ beta) / np.sqrt(n) / 2.0
    z = chi @ beta + sig * rng.standard_normal(n)
    d = _drops(chi, z, learner, lam)
    mask = np.zeros(p, bool)
    mask[supp] = True
    return d[~mask]


def rate_in_n(
    ns,
    learner="ridge",
    lam=1e-3,
    p=20,
    reps=60,
    heavy=False,
    deltas=(0.5, 0.1, 0.05, 0.01),
    n_jobs=100,
    label="",
):
    rows = {}
    for n in ns:
        got = Parallel(n_jobs=n_jobs)(
            delayed(_rate_rep)(r, n, p, learner, lam, heavy) for r in range(reps)
        )
        rows[n] = np.concatenate(got)
    res = {
        "ns": list(ns),
        "learner": learner,
        "lambda": lam,
        "p": p,
        "replicates": reps,
        "heavy_tailed_design": heavy,
        "quantiles": {},
    }
    for d in deltas:
        q = np.array([max(np.quantile(rows[n], 1 - d), EPS) for n in ns])
        a, c = np.polyfit(np.log(ns), np.log(q), 1)
        pred = a * np.log(ns) + c
        r2 = 1 - np.sum((np.log(q) - pred) ** 2) / np.sum(
            (np.log(q) - np.log(q).mean()) ** 2
        )
        # local exponent over the last three sample sizes (detects a bias floor)
        a_tail, _ = np.polyfit(np.log(ns[-3:]), np.log(q[-3:]), 1)
        res["quantiles"][str(d)] = {
            "q_values": q.tolist(),
            "fitted_exponent": float(a),
            "fitted_exponent_last3": float(a_tail),
            "r2": float(r2),
        }
        print(
            f"[rate-n {label}] delta={d}: exp={a:+.4f} (pred -0.5) "
            f"tail_exp={a_tail:+.4f} R2={r2:.4f}",
            flush=True,
        )
    big = rows[max(ns)]
    ds = np.array([0.5, 0.2, 0.1, 0.05, 0.02, 0.01, 0.005])
    qs = np.array([max(np.quantile(big, 1 - d), EPS) for d in ds])
    b, _ = np.polyfit(np.log(np.log(1 / ds)), np.log(qs), 1)
    lin = np.polyfit(np.sqrt(np.log(1 / ds)), qs, 1)
    fit = np.polyval(lin, np.sqrt(np.log(1 / ds)))
    r2lin = 1 - np.sum((qs - fit) ** 2) / np.sum((qs - qs.mean()) ** 2)
    res["delta_dependence"] = {
        "n": int(max(ns)),
        "deltas": ds.tolist(),
        "quantiles": qs.tolist(),
        "fitted_exponent_on_log(1/delta)": float(b),
        "r2_linear_in_sqrt_log(1/delta)": float(r2lin),
    }
    print(
        f"[rate-delta {label}] exp on log(1/delta)={b:+.4f} (pred +0.5), "
        f"R2(linear in sqrt(log(1/delta)))={r2lin:.4f}",
        flush=True,
    )
    return res


# ------------------------------------------- 5. the nu_hat / rho_hat coefficients
def _nurho_rep(rep, n, p, lam=1e-3):
    """The extra 'uninformative feature' is y itself: nu_hat regresses X^j on
    X^{-j}; rho_hat regresses X^j on (X^{-j}, y)."""
    rng = np.random.default_rng(SEED0 + 17_000_000 + 613 * rep + 5 * n)
    X = _design(rng, n, p)
    beta, supp = _beta_blocks(rng, p)
    y = X @ beta + rng.standard_normal(n)
    mask = np.zeros(p, bool)
    mask[supp] = True
    out = np.empty(p)
    for j in range(p):
        cols = [k for k in range(p) if k != j]
        Xm = X[:, cols]
        Xy = np.column_stack([Xm, y])
        th_full = _fit(Xy, X[:, j], "ridge", lam)  # theta_hat  (with y)
        th_red = _fit(Xm, X[:, j], "ridge", lam)  # theta_hat^{-y}
        tt = np.concatenate([th_red, [0.0]])  # theta_tilde
        out[j] = float(np.linalg.norm(tt - th_full))
    return out, mask


def nurho(ns=(50, 100, 200, 400, 800, 1600, 3200), p=20, reps=40, n_jobs=100):
    res = {"ns": list(ns), "p": p, "replicates": reps}
    qn, qi = [], []
    for n in ns:
        got = Parallel(n_jobs=n_jobs)(delayed(_nurho_rep)(r, n, p) for r in range(reps))
        d = np.array([g[0] for g in got])
        m = np.array([g[1] for g in got])
        qn.append(float(np.quantile(d[~m], 0.95)))
        qi.append(float(np.quantile(d[m], 0.95)))
    a, _ = np.polyfit(np.log(ns), np.log(qn), 1)
    ai, _ = np.polyfit(np.log(ns), np.log(qi), 1)
    res["null_q95"] = qn
    res["important_q95"] = qi
    res["null_fitted_exponent"] = float(a)
    res["important_fitted_exponent"] = float(ai)
    print(
        f"[nu/rho] null q95 exponent={a:+.4f} (pred -0.5); "
        f"important exponent={ai:+.4f}",
        flush=True,
    )
    print(f"[nu/rho] null q95={np.round(qn,5).tolist()}", flush=True)
    print(f"[nu/rho] imp  q95={np.round(qi,5).tolist()}", flush=True)
    return res


if __name__ == "__main__":
    os.makedirs(OUT, exist_ok=True)
    os.makedirs(FIGS, exist_ok=True)
    t0 = time.time()
    ns = [50, 100, 200, 400, 800, 1600, 3200]
    res = {"seed0": SEED0}
    res["figure1"] = fig1()
    res["rate_ridge_lam1e-3"] = rate_in_n(ns, "ridge", 1e-3, label="ridge lam=1e-3")
    res["rate_ridge_lam1e-2"] = rate_in_n(ns, "ridge", 1e-2, label="ridge lam=1e-2")
    res["rate_ridge_lam1e-1"] = rate_in_n(ns, "ridge", 1e-1, label="ridge lam=1e-1")
    res["rate_ridge_lam1e-6"] = rate_in_n(ns, "ridge", 1e-6, label="ridge lam=1e-6")
    res["rate_ridge_heavy_t2"] = rate_in_n(
        ns, "ridge", 1e-3, heavy=True, label="ridge t2 design"
    )
    res["rate_lasso_lam1e-2"] = rate_in_n(ns, "lasso", 1e-2, label="lasso lam=1e-2")
    res["nu_rho_coefficients"] = nurho()
    res["seconds"] = round(time.time() - t0, 1)
    with open(os.path.join(OUT, "claim3_stability.json"), "w") as f:
        json.dump(res, f, indent=2)
    print("wrote", os.path.join(OUT, "claim3_stability.json"))