scripts/c6_wdbc.py

"""Claim 6 (Figure 6, Section 5.2).

"On the Wisconsin Breast Cancer real dataset, Semi-knockoffs is applied across
Random Forest, Neural Network, and Gradient Boosting models to demonstrate
model-agnostic feature selection."

Data: sklearn.datasets.load_breast_cancer -> the Wisconsin Diagnostic Breast
Cancer (WDBC) set, n = 569, p = 30 (matching Appendix F.6).

Exactly as in Appendix F.6, the true support is unknown, so an ARTIFICIAL NULL
feature correlated at 0.6 with the original inputs is appended; the rejection
rate of that feature over repetitions estimates the type-I error.

Reported per model (Random Forest / Neural Network / Gradient Boosting), over
repetitions with distinct seeds:
  * number of discoveries at alpha = 0.05 (Semi-knockoffs Wilcoxon, no split)
  * the same with 5-permutation derandomisation
  * HRT with a 50/50 split, as the split-based comparator
  * the artificial-null rejection rate = estimated type-I error
  * stability of the discovery set across seeds (the paper's claim that
    Semi-knockoffs "exhibits stable power across models and consistent numbers
    of discoveries, unlike other methods whose results vary significantly with
    the random seed")
  * the FDR-controlling knockoff-threshold variant (Algorithm 4) at q = 0.1
"""

from __future__ import annotations

import json
import os
import sys
import time

import numpy as np
from joblib import Parallel, delayed
from sklearn.datasets import load_breast_cancer
from sklearn.ensemble import GradientBoostingClassifier, RandomForestClassifier
from sklearn.linear_model import Ridge
from sklearn.neural_network import MLPClassifier
from sklearn.preprocessing import StandardScaler

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__))))
from semiknockoffs import (  # noqa: E402
    fit_nu_rho,
    hrt_pvalue,
    knockoff_select,
    logloss,
    sko_pvalue,
    sko_statistic,
)

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


def load_wdbc(rng, corr=0.6):
    d = load_breast_cancer()
    X = StandardScaler().fit_transform(d.data)
    y = d.target.astype(float)
    n, p = X.shape
    # artificial null feature correlated at ~0.6 with the original inputs:
    # a normalised random combination of the standardised columns plus noise.
    w = rng.standard_normal(p)
    w /= np.linalg.norm(w)
    base = X @ w
    base /= np.std(base)
    z = corr * base + np.sqrt(1 - corr**2) * rng.standard_normal(n)
    Xa = np.column_stack([X, z])
    names = list(d.feature_names) + ["ARTIFICIAL_NULL"]
    return Xa, y, names


def _model(kind, X, y, seed):
    if kind == "rf":
        return RandomForestClassifier(
            n_estimators=100, random_state=seed, n_jobs=1
        ).fit(X, y)
    if kind == "nn":
        return MLPClassifier(
            hidden_layer_sizes=(64, 64), max_iter=800, random_state=seed
        ).fit(X, y)
    if kind == "gb":
        return GradientBoostingClassifier(random_state=seed).fit(X, y)
    raise ValueError(kind)


def _rep(rep, kind, K=200):
    rng = np.random.default_rng(SEED0 + 53_000_000 + 977 * rep)
    X, y, names = load_wdbc(rng)
    n, p = X.shape

    m_full = _model(kind, X, y, rep)

    def pred_full(Z):
        return m_full.predict_proba(Z)[:, 1]

    idx = rng.permutation(n)
    tr, te = idx[: n // 2], idx[n // 2 :]
    m_split = _model(kind, X[tr], y[tr], rep)

    def pred_split(Z):
        return m_split.predict_proba(Z)[:, 1]

    acc_full = float(np.mean((pred_full(X) > 0.5) == (y > 0.5)))
    acc_split = float(np.mean((pred_split(X[te]) > 0.5) == (y[te] > 0.5)))

    p_sko = np.ones(p)
    p_sko5 = np.ones(p)
    p_hrt = np.ones(p)
    W = np.zeros(p)
    for j in range(p):
        cols = [k for k in range(p) if k != j]
        nu, rho = fit_nu_rho(X, y, j, alpha=1.0)
        p_sko[j] = sko_pvalue(X, y, j, nu, rho, pred_full, logloss, rng)
        p_sko5[j] = sko_pvalue(X, y, j, nu, rho, pred_full, logloss, rng, n_perm=5)
        W[j] = sko_statistic(X, y, j, nu, rho, pred_full, logloss, rng, n_perm=5)

        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))
        p_hrt[j] = hrt_pvalue(
            X[te],
            y[te],
            j,
            rg.predict(X[np.ix_(te, cols)]),
            res_tr * dfc,
            pred_split,
            logloss,
            rng,
            K=K,
        )

    return {
        "acc_full": acc_full,
        "acc_split": acc_split,
        "p_sko": p_sko.tolist(),
        "p_sko5": p_sko5.tolist(),
        "p_hrt": p_hrt.tolist(),
        "W": W.tolist(),
        "sel_knockoff": {str(q): knockoff_select(W, q).tolist() for q in QS},
        "names": names,
    }


def run(kind, reps=50, n_jobs=50):
    t0 = time.time()
    got = Parallel(n_jobs=n_jobs)(delayed(_rep)(r, kind) for r in range(reps))
    names = got[0]["names"]
    p = len(names)
    null_idx = p - 1
    res = {
        "model": kind,
        "replicates": reps,
        "n": 569,
        "p_with_artificial": p,
        "accuracy_no_split_insample": float(np.mean([g["acc_full"] for g in got])),
        "accuracy_split_heldout": float(np.mean([g["acc_split"] for g in got])),
        "seconds": round(time.time() - t0, 1),
    }
    for key, lab in (
        ("p_sko", "SKO_Wcx"),
        ("p_sko5", "SKO_Wcx_p5"),
        ("p_hrt", "HRT_split"),
    ):
        P = np.array([g[key] for g in got])  # (reps, p)
        disc = P <= ALPHA
        n_disc = disc[:, :null_idx].sum(1)  # real features only
        res[lab] = {
            "mean_discoveries": float(n_disc.mean()),
            "sd_discoveries": float(n_disc.std(ddof=1)),
            "min_discoveries": int(n_disc.min()),
            "max_discoveries": int(n_disc.max()),
            "artificial_null_rejection_rate": float(disc[:, null_idx].mean()),
            "jaccard_stability_across_seeds": _jaccard(disc[:, :null_idx]),
        }
        print(
            f"[{kind}] {lab:11s} disc={n_disc.mean():.2f}+-{n_disc.std(ddof=1):.2f} "
            f"typeI(artificial null)={disc[:, null_idx].mean():.3f} "
            f"jaccard={res[lab]['jaccard_stability_across_seeds']:.3f}",
            flush=True,
        )
    res["SKO_knockoff_threshold"] = {}
    for q in QS:
        sel = [set(g["sel_knockoff"][str(q)]) for g in got]
        nsel = np.array([len([s for s in x if s != null_idx]) for x in sel])
        res["SKO_knockoff_threshold"][str(q)] = {
            "mean_selected": float(nsel.mean()),
            "sd_selected": float(nsel.std(ddof=1)),
            "artificial_null_selection_rate": float(
                np.mean([null_idx in s for s in sel])
            ),
        }
        print(
            f"[{kind}] knockoff q={q}: sel={nsel.mean():.2f}"
            f"+-{nsel.std(ddof=1):.2f} artificial-null selected in "
            f"{np.mean([null_idx in s for s in sel]):.3f}",
            flush=True,
        )
    # per-feature rejection frequency for SKO_Wcx_p5
    P5 = np.array([g["p_sko5"] for g in got])
    res["per_feature_rejection_rate_sko_p5"] = {
        names[j]: float(np.mean(P5[:, j] <= ALPHA)) for j in range(p)
    }
    return res


def _jaccard(disc):
    """Mean pairwise Jaccard similarity of the selected sets across seeds."""
    r = disc.shape[0]
    vals = []
    for a in range(r):
        for b in range(a + 1, r):
            u = np.logical_or(disc[a], disc[b]).sum()
            i = np.logical_and(disc[a], disc[b]).sum()
            vals.append(1.0 if u == 0 else i / u)
    return float(np.mean(vals))


if __name__ == "__main__":
    os.makedirs(OUT, exist_ok=True)
    res = {
        "seed0": SEED0,
        "alpha": ALPHA,
        "q": Q,
        "dataset": "sklearn.datasets.load_breast_cancer (WDBC, n=569, p=30)",
    }
    for kind in ("rf", "nn", "gb"):
        res[kind] = run(kind)
    # cross-model agreement of the SKO_Wcx_p5 discovery sets = model-agnosticism
    keys = list(res["rf"]["per_feature_rejection_rate_sko_p5"].keys())
    mat = {
        k: [res[m]["per_feature_rejection_rate_sko_p5"][k] for m in ("rf", "nn", "gb")]
        for k in keys
    }
    stable = [k for k, v in mat.items() if min(v) >= 0.8]
    res["cross_model"] = {
        "features_rejected_in_>=80pct_of_seeds_by_ALL_three_models": stable,
        "n_such_features": len(stable),
        "per_feature_rate_by_model": mat,
    }
    print("cross-model stable discoveries:", len(stable), stable)
    with open(os.path.join(OUT, "claim6_wdbc.json"), "w") as f:
        json.dump(res, f, indent=2)
    print("wrote", os.path.join(OUT, "claim6_wdbc.json"))