"""Semi-knockoffs claims 4 (strengthened), 5 and 6.""" import json import os import warnings import numpy as np from sklearn.datasets import load_breast_cancer from sklearn.ensemble import GradientBoostingRegressor, RandomForestRegressor from sklearn.linear_model import LinearRegression from sklearn.neural_network import MLPRegressor from sklearn.preprocessing import StandardScaler warnings.filterwarnings("ignore") from skocore import ar1_design, sko_pvalue, sko_statistic, fdp_power os.makedirs("outputs", exist_ok=True) OUT = json.load(open("outputs/results.json")) if os.path.exists("outputs/results.json") else {} def model_for(kind, X, y, seed=0): return {"gb": GradientBoostingRegressor(random_state=seed), "rf": RandomForestRegressor(n_estimators=100, random_state=seed, n_jobs=-1), "nn": MLPRegressor(hidden_layer_sizes=(64, 32), max_iter=600, random_state=seed), }[kind].fit(X, y) # ---------------------------------------------------------------- HRT baseline def hrt_pvalue(X, y, j, model, rng, n_perm=200, split=0.5, seed=0): """Holdout Randomization Test: fit on a TRAIN split, test on a HELD-OUT split by resampling feature j from its conditional distribution there. This is the method Semi-knockoffs claims to improve on by removing the split. """ n = len(y) idx = rng.permutation(n) ntr = int(split * n) tr, te = idx[:ntr], idx[ntr:] m = model.__class__(**model.get_params()).fit(X[tr], y[tr]) Xte, yte = X[te], y[te] # conditional model for X^j | X^{-j} fitted on the TRAIN half only cond = LinearRegression().fit(np.delete(X[tr], j, axis=1), X[tr, j]) mu = cond.predict(np.delete(Xte, j, axis=1)) resid = X[tr, j] - cond.predict(np.delete(X[tr], j, axis=1)) obs = float(np.mean((m.predict(Xte) - yte) ** 2)) cnt = 0 for _ in range(n_perm): Xp = Xte.copy() Xp[:, j] = mu + rng.choice(resid, size=len(te), replace=True) if float(np.mean((m.predict(Xp) - yte) ** 2)) <= obs: cnt += 1 return (1.0 + cnt) / (1.0 + n_perm) # ------------------------------- claim 4 (strengthened): double robustness arms def claim4b(ns=(150, 300, 600, 1200, 2400), p=20, reps=30, rho=0.5): """Decay of |W_j| for a null feature under two nuisance-quality regimes.""" res = {} for arm, nuis in (("well-specified", "ridge"), ("degraded", "degraded")): rows = [] for n in ns: w = [] for r in range(reps): rng = np.random.default_rng(4242 + r) X = ar1_design(n, p, rho, rng) beta = np.zeros(p); beta[:5] = 1.0 y = X @ beta + rng.standard_normal(n) m = model_for("gb", X, y, seed=r) if nuis == "degraded": # deliberately weak nuisances: use only 3 of the p-1 covariates Xs = X.copy() keep = list(range(3)) + [p - 1] Xs = Xs[:, keep] w.append(abs(sko_statistic(Xs, y, len(keep) - 1, model_for("gb", Xs, y, seed=r), rng, seed=r))) else: w.append(abs(sko_statistic(X, y, p - 1, m, rng, seed=r))) rows.append({"n": n, "mean_absW": float(np.mean(w)), "se": float(np.std(w, ddof=1) / np.sqrt(reps))}) print(f" claim4b[{arm}] n={n} |W|={np.mean(w):.6f}", flush=True) lx = np.log([r["n"] for r in rows]); ly = np.log([r["mean_absW"] for r in rows]) A = np.vstack([lx, np.ones_like(lx)]).T sl, ic = np.linalg.lstsq(A, ly, rcond=None)[0] pred = A @ np.array([sl, ic]) r2 = 1 - float(((ly - pred) ** 2).sum()) / float(((ly - ly.mean()) ** 2).sum()) res[arm] = {"rows": rows, "slope": float(sl), "r2": float(r2)} print(f" claim4b[{arm}] slope {sl:.4f} R2 {r2:.4f}", flush=True) res["both_faster_than_root_n"] = bool(all(v["slope"] < -0.5 for v in res.values() if isinstance(v, dict) and "slope" in v)) OUT["claim4b"] = res print("claim4b done", flush=True) # ------------------------- claim 5: adjacent support, Semi-KO vs HRT + derandom def claim5(reps=40, n=200, p=30, rho=0.8, alpha=0.05, n_perm=5, betas=(0.15, 0.25, 0.4, 0.8)): """Adjacent-feature support at several signal strengths. A single operating point is uninformative: at a strong signal both methods saturate at power 1.0 and the comparison says nothing. HRT's cost is that it must TRAIN on half the data and test on the other half, so its disadvantage should appear when the signal is weak relative to n. We therefore sweep the signal strength and report the whole curve. """ res = {"reps": reps, "n": n, "p": p, "rho": rho, "alpha": alpha, "n_permutations": n_perm, "support": "adjacent (features 10-14)", "betas": list(betas), "curves": {}} for mk in ("gb", "rf"): curve = [] for b in betas: sko_t1, sko_pw, hrt_t1, hrt_pw, der_pw = [], [], [], [], [] for r in range(reps): rng = np.random.default_rng(2100 + r) X = ar1_design(n, p, rho, rng) beta = np.zeros(p); beta[10:15] = b y = X @ beta + rng.standard_normal(n) m = model_for(mk, X, y, seed=r) ja, jn = 12, 25 sko_pw.append(sko_pvalue(X, y, ja, m, rng, seed=r) <= alpha) sko_t1.append(sko_pvalue(X, y, jn, m, rng, seed=r) <= alpha) hrt_pw.append(hrt_pvalue(X, y, ja, m, rng, seed=r) <= alpha) hrt_t1.append(hrt_pvalue(X, y, jn, m, rng, seed=r) <= alpha) der = [sko_pvalue(X, y, ja, m, rng, seed=r) for _ in range(n_perm)] der_pw.append(float(np.median(der)) <= alpha) row = {"beta": b, "sko_power": float(np.mean(sko_pw)), "sko_type_I": float(np.mean(sko_t1)), "hrt_power": float(np.mean(hrt_pw)), "hrt_type_I": float(np.mean(hrt_t1)), "sko_derandomised_power": float(np.mean(der_pw))} row["power_gap"] = row["sko_power"] - row["hrt_power"] row["derand_gain"] = row["sko_derandomised_power"] - row["sko_power"] curve.append(row) print(f" claim5[{mk}] beta={b:<5} SKO {row['sko_power']:.3f} (t1 {row['sko_type_I']:.3f}) | " f"HRT {row['hrt_power']:.3f} (t1 {row['hrt_type_I']:.3f}) | " f"derand {row['sko_derandomised_power']:.3f}", flush=True) res["curves"][mk] = curve res[f"{mk}_max_power_gap"] = max(r["power_gap"] for r in curve) res[f"{mk}_sko_ge_hrt_everywhere"] = bool(all(r["power_gap"] >= 0 for r in curve)) res[f"{mk}_max_sko_type_I"] = max(r["sko_type_I"] for r in curve) OUT["claim5"] = res # ------------------------------ claim 6: Wisconsin Breast Cancer, model-agnostic def claim6(reps=30, alpha=0.05): data = load_breast_cancer() Xr, yr = data.data, data.target.astype(float) Xs = StandardScaler().fit_transform(Xr) res = {"dataset": "Wisconsin Breast Cancer", "n": int(Xs.shape[0]), "p_original": int(Xs.shape[1]), "reps": reps, "alpha": alpha} for mk in ("rf", "nn", "gb"): t1, pw = [], [] for r in range(reps): rng = np.random.default_rng(6100 + r) # inject a conditionally-null feature: a function of the others plus # independent noise, so it carries no information about y given X noise = rng.standard_normal(len(yr)) null_feat = Xs[:, :5].mean(axis=1) + noise X = np.column_stack([Xs, null_feat]) m = model_for(mk, X, yr, seed=r) jnull = X.shape[1] - 1 t1.append(sko_pvalue(X, yr, jnull, m, rng, seed=r) <= alpha) # a genuinely predictive feature for reference (worst mean radius) pw.append(sko_pvalue(X, yr, 0, m, rng, seed=r) <= alpha) res[mk] = {"type_I_injected_null": float(np.mean(t1)), "rejects_real_feature": float(np.mean(pw))} print(f" claim6[{mk}] type-I on injected null {np.mean(t1):.3f} | " f"rejects real feature {np.mean(pw):.3f}", flush=True) res["model_agnostic"] = bool(all(res[k]["type_I_injected_null"] <= 0.10 for k in ("rf", "nn", "gb"))) OUT["claim6"] = res if __name__ == "__main__": import sys fns = {"4b": claim4b, "5": claim5, "6": claim6} for w in (sys.argv[1:] or ["6", "5", "4b"]): print("=== claim", w, flush=True) fns[w]() json.dump(OUT, open("outputs/results.json", "w"), indent=2) print("saved")