scripts/c4_double_robustness.py

"""Claim 4 (Theorem 4.3, double robustness).

"The loss difference between imputed feature distributions decays at a compound
rate O_P(a_n b_n) even when both the predictive model and the sampler have
estimation error."

Theorem 4.3 (verbatim structure): for j in H0, t -> l(t, y) is C^2 with bounded
derivatives, m_hat differentiable in coordinate j with
a_n := d_{x^j} m_hat(X^{-j}, t), and nu_hat - nu = O_P(b_n).  Then

        l(m_hat(X~'), y) - l(m_hat(X~), y) = O_P(a_n b_n).

Here X~ uses the ORACLE nu and X~' the ESTIMATED nu_hat, with the SAME residual
index, so the difference isolates the imputer error.

Independent tests:

 1. FIGURE 3 reproduction.  y = 0.8X1 + 0.6X2 + 0.4X3 + 0.2X4 + sin(X1) + eps,
    eps ~ N(0, 0.5), n = 2000, X ~ N(0, Sigma), Sigma_ij = 0.5^|i-j|, nu_hat
    linear, for the null coordinate j = 0.  Blue: the semi-knockoff statistic
    l(m(X~'_1),y) - l(m(X~'_2),y).  Orange: l(m(X~'_1),y) - l(m(X~_1),y).
    Theorem 4.3 => orange is far more concentrated than blue.  Random Forest,
    Neural Network, Gradient Boosting.

 2. COMPOUND-RATE FIT.  Sweep n and measure, separately,
       a_n  = the model's empirical sensitivity to coordinate j
              (finite-difference d m_hat / d x^j),
       b_n  = ||nu_hat - nu||_2 / sqrt(n),
       D_n  = |mean_i l(m_hat(X~'_i), y_i) - l(m_hat(X~_i), y_i)|.
    Fit log D_n = c + s log n and compare s with the exponent of a_n b_n
    measured on the same runs.  A single-robustness (linear, i.e. O_P(b_n))
    decay would give s ~ -1/2; the compound rate gives s ~ -1.

 3. LINEAR-MODEL SPECIALISATION (Remark 4.4): m_hat(X) = beta_hat' X so
    a_n = beta_hat^j which is O_P(n^{-1/2}) under the null, and nu_hat is
    parametric so b_n = O_P(n^{-1/2}) -> D_n should be O_P(1/n).

 4. BOUNDARY PROBE.  Break the null: for an IMPORTANT coordinate j, a_n is
    O(1) rather than O_P(n^{-1/2}), so only the single rate b_n should survive
    (s ~ -1/2).  This is the discriminating prediction of Theorem 4.3.
"""

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, RandomForestRegressor
from sklearn.linear_model import LinearRegression, Ridge
from sklearn.neural_network import MLPRegressor

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

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


def _gen(rng, n, p=10, rho=0.5):
    """Figure 3 data-generating process."""
    Sigma = ar1_cov(p, rho)
    L = np.linalg.cholesky(Sigma)
    X = rng.standard_normal((n, p)) @ L.T
    y = (
        0.8 * X[:, 1]
        + 0.6 * X[:, 2]
        + 0.4 * X[:, 3]
        + 0.2 * X[:, 4]
        + np.sin(X[:, 1])
        + 0.5 * rng.standard_normal(n)
    )
    return X, y, Sigma  # coordinate 0 is null


def _model(kind, X, y, seed):
    if kind == "rf":
        return RandomForestRegressor(n_estimators=100, random_state=seed, n_jobs=1).fit(
            X, y
        )
    if kind == "nn":
        return MLPRegressor(
            hidden_layer_sizes=(64, 64), max_iter=800, random_state=seed
        ).fit(X, y)
    if kind == "gb":
        return GradientBoostingRegressor(random_state=seed).fit(X, y)
    if kind == "linear":
        return LinearRegression().fit(X, y)
    raise ValueError(kind)


# ------------------------------------------------------------------ Figure 3
def _fig3_rep(rep, kind, n=2000, p=10, j=0):
    rng = np.random.default_rng(SEED0 + 23_000_000 + 101 * rep)
    X, y, Sigma = _gen(rng, n, p)
    m = _model(kind, X, y, rep)
    cols = [k for k in range(p) if k != j]

    nu = gaussian_nu(X, j, Sigma)  # ORACLE nu
    nu_h = Ridge(alpha=1.0).fit(X[:, cols], X[:, j]).predict(X[:, cols])  # ESTIMATED
    e_or = X[:, j] - nu
    e_es = X[:, j] - nu_h

    pi1 = rng.permutation(n)
    pi2 = rng.permutation(n)

    def build(base, eps, perm):
        Z = X.copy()
        Z[:, j] = base + eps[perm]
        return Z

    L1p = sq_loss(m.predict(build(nu_h, e_es, pi1)), y)  # X~'_1 (estimated)
    L2p = sq_loss(m.predict(build(nu_h, e_es, pi2)), y)  # X~'_2 (estimated)
    L1o = sq_loss(m.predict(build(nu, e_or, pi1)), y)  # X~_1  (oracle)
    blue = L1p - L2p  # the semi-knockoff statistic itself
    orange = L1p - L1o  # theoretical vs estimated imputer
    return blue, orange


def fig3(reps=8, n_jobs=24):
    res = {"n": 2000, "p": 10, "replicates": reps, "null_coordinate": 0}
    for kind in ("rf", "nn", "gb"):
        got = Parallel(n_jobs=n_jobs)(delayed(_fig3_rep)(r, kind) for r in range(reps))
        blue = np.concatenate([g[0] for g in got])
        orange = np.concatenate([g[1] for g in got])
        res[kind] = {
            "blue_std": float(np.std(blue)),
            "orange_std": float(np.std(orange)),
            "blue_iqr": float(np.subtract(*np.percentile(blue, [75, 25]))),
            "orange_iqr": float(np.subtract(*np.percentile(orange, [75, 25]))),
            "concentration_ratio_std": float(np.std(blue) / np.std(orange)),
            "blue_mean": float(np.mean(blue)),
            "orange_mean": float(np.mean(orange)),
            "blue_sample": blue[:4000].tolist(),
            "orange_sample": orange[:4000].tolist(),
        }
        print(
            f"[fig3] {kind}: blue_std={np.std(blue):.5f} "
            f"orange_std={np.std(orange):.5f} "
            f"ratio={np.std(blue)/np.std(orange):.2f}",
            flush=True,
        )
    return res


# ------------------------------------------------------------- compound rate
def _rate_rep(rep, n, kind, j, p=10):
    rng = np.random.default_rng(SEED0 + 29_000_000 + 313 * rep + 11 * n)
    X, y, Sigma = _gen(rng, n, p)
    m = _model(kind, X, y, rep)
    cols = [k for k in range(p) if k != j]

    nu = gaussian_nu(X, j, Sigma)
    nu_h = Ridge(alpha=1.0).fit(X[:, cols], X[:, j]).predict(X[:, cols])
    b_n = float(np.linalg.norm(nu_h - nu) / np.sqrt(n))  # ||nu_hat - nu||

    # a_n : empirical sensitivity of m_hat to coordinate j.  Tree ensembles are
    # piecewise constant, so an infinitesimal derivative is identically 0; we
    # use a SECANT sensitivity at scale h = 0.25 sd(X^j), which reduces to the
    # derivative for smooth models and stays informative for trees.
    h = 0.25 * float(np.std(X[:, j]))
    Xp, Xm = X.copy(), X.copy()
    Xp[:, j] += h
    Xm[:, j] -= h
    a_n = float(np.mean(np.abs(m.predict(Xp) - m.predict(Xm)) / (2 * h)))

    perm = rng.permutation(n)
    e_or = X[:, j] - nu
    e_es = X[:, j] - nu_h

    def build(base, eps):
        Z = X.copy()
        Z[:, j] = base + eps[perm]
        return Z

    d = sq_loss(m.predict(build(nu_h, e_es)), y) - sq_loss(
        m.predict(build(nu, e_or)), y
    )
    # Theorem 4.3 is a POINTWISE O_P statement, so use pointwise scales of |d|
    # rather than the average over i (which would add its own 1/sqrt(n)).
    return (a_n, b_n, float(np.mean(np.abs(d))), float(np.quantile(np.abs(d), 0.9)))


def _slope(ns, v):
    v = np.asarray(v, float)
    ok = v > 0
    if ok.sum() < 3:
        return float("nan")
    return float(np.polyfit(np.log(np.asarray(ns)[ok]), np.log(v[ok]), 1)[0])


def rate(ns, kind, j, reps=40, n_jobs=100, label=""):
    A, B, Dm, Dq = [], [], [], []
    for n in ns:
        got = Parallel(n_jobs=n_jobs)(
            delayed(_rate_rep)(r, n, kind, j) for r in range(reps)
        )
        arr = np.array(got)
        A.append(float(np.median(arr[:, 0])))
        B.append(float(np.median(arr[:, 1])))
        Dm.append(float(np.median(arr[:, 2])))
        Dq.append(float(np.median(arr[:, 3])))
    sa, sb = _slope(ns, A), _slope(ns, B)
    sdm, sdq = _slope(ns, Dm), _slope(ns, Dq)
    res = {
        "ns": list(ns),
        "model": kind,
        "coordinate": j,
        "null": bool(j == 0),
        "replicates": reps,
        "a_n": A,
        "b_n": B,
        "D_n_mean_abs": Dm,
        "D_n_q90_abs": Dq,
        "exponent_a_n": sa,
        "exponent_b_n": sb,
        "exponent_D_n_mean_abs": sdm,
        "exponent_D_n_q90_abs": sdq,
        "predicted_compound_exponent": sa + sb,
        "single_robust_exponent_b_only": sb,
    }
    print(
        f"[rate {label}] a_n exp={sa:+.3f}  b_n exp={sb:+.3f}  "
        f"=> compound pred={sa+sb:+.3f} | measured D_n exp: "
        f"mean|d|={sdm:+.3f} q90|d|={sdq:+.3f}",
        flush=True,
    )
    return res


if __name__ == "__main__":
    os.makedirs(OUT, exist_ok=True)
    t0 = time.time()
    ns = [100, 200, 400, 800, 1600, 3200, 6400]
    res = {"seed0": SEED0}
    res["figure3"] = fig3()
    res["rate_linear_null"] = rate(ns, "linear", 0, label="linear j=0 (null)")
    res["rate_linear_important"] = rate(ns, "linear", 1, label="linear j=1 (important)")
    res["rate_gb_null"] = rate(ns, "gb", 0, reps=24, label="GB j=0 (null)")
    res["rate_gb_important"] = rate(ns, "gb", 1, reps=24, label="GB j=1 (important)")
    res["rate_rf_null"] = rate(ns[:-1], "rf", 0, reps=16, label="RF j=0 (null)")
    res["rate_nn_null"] = rate(ns[:-1], "nn", 0, reps=16, label="NN j=0 (null)")
    res["rate_nn_important"] = rate(
        ns[:-1], "nn", 1, reps=16, label="NN j=1 (important)"
    )
    res["seconds"] = round(time.time() - t0, 1)
    with open(os.path.join(OUT, "claim4_double_robustness.json"), "w") as f:
        json.dump(res, f, indent=2)
    print("wrote", os.path.join(OUT, "claim4_double_robustness.json"))