scripts/semiknockoffs.py

"""Independent re-implementation of Semi-knockoffs (arXiv 2601.23124) and HRT.

Everything here is written from the paper's Algorithm boxes 1-4 (Appendix C) and
Eq. (1) (knockoff threshold).  No code from the authors' repository
(https://github.com/AngelReyero/loss_based_KO) was consulted or used.

Key objects
-----------
sko_statistic(...)        -> W^j_SKO  (Algorithm 3 oracle / Algorithm 4 estimated)
sko_pvalue(...)           -> Wilcoxon / sign-test p-value (Algorithm 1 / 2)
knockoff_threshold(W, q)  -> T_q of Eq. (1)  (knockoff+ form, with the "1 +")
hrt_pvalue(...)           -> Holdout Randomization Test (Tansey et al., 2022)
gaussian_nu / gaussian_rho-> ORACLE conditional expectations for a Gaussian
                             linear model, used to test the "Given nu_j, rho_j"
                             hypothesis of Theorems 3.3 / 3.4 exactly.
"""

from __future__ import annotations

import numpy as np
from scipy import stats
from sklearn.linear_model import Ridge, LinearRegression


# --------------------------------------------------------------------------
# losses
# --------------------------------------------------------------------------
def sq_loss(pred, y):
    return (np.asarray(pred).ravel() - np.asarray(y).ravel()) ** 2


def logloss(prob, y, eps=1e-9):
    p = np.clip(np.asarray(prob).ravel(), eps, 1 - eps)
    y = np.asarray(y).ravel()
    return -(y * np.log(p) + (1 - y) * np.log(1 - p))


# --------------------------------------------------------------------------
# oracle conditional expectations for the Gaussian linear model
#   X ~ N(0, Sigma),  y = beta' X + sigma * eps
# --------------------------------------------------------------------------
def gaussian_nu(X, j, Sigma):
    """nu_j(x^-j) = E[X^j | X^-j] = Sigma_{j,-j} Sigma_{-j,-j}^{-1} x^{-j}."""
    p = Sigma.shape[0]
    idx = [k for k in range(p) if k != j]
    S_mj = Sigma[np.ix_(idx, idx)]
    s_j = Sigma[j, idx]
    coef = np.linalg.solve(S_mj, s_j)
    return X[:, idx] @ coef


def gaussian_rho(X, y, j, Sigma, beta, sigma):
    """rho_j(x^-j, y) = E[X^j | X^-j, y] in the Gaussian linear model.

    (X, y) is jointly Gaussian: cov(X, y) = Sigma beta, var(y) = beta'Sigma beta + sigma^2.
    """
    p = Sigma.shape[0]
    idx = [k for k in range(p) if k != j]
    Sy = Sigma @ beta  # cov(X, y), length p
    vy = float(beta @ Sigma @ beta + sigma**2)
    # conditioning vector Z = (X^{-j}, y)
    C = np.zeros((p, p))  # cov of Z
    C[: p - 1, : p - 1] = Sigma[np.ix_(idx, idx)]
    C[: p - 1, -1] = Sy[idx]
    C[-1, : p - 1] = Sy[idx]
    C[-1, -1] = vy
    c = np.concatenate([Sigma[j, idx], [Sy[j]]])  # cov(X^j, Z)
    coef = np.linalg.solve(C, c)
    Z = np.column_stack([X[:, idx], np.asarray(y).ravel()])
    return Z @ coef


# --------------------------------------------------------------------------
# estimated conditional expectations (Algorithm 4, "Regression 1 / 2")
# --------------------------------------------------------------------------
def fit_nu_rho(X, y, j, alpha=1.0, learner="ridge"):
    """Return (nu_hat(X^-j), rho_hat(X^-j, y)) fitted on the SAME n samples.

    The paper motivates an l2-regularised ERM (Eq. 2 / Theorem 4.1); we use Ridge.
    """
    p = X.shape[1]
    idx = [k for k in range(p) if k != j]
    Xm = X[:, idx]
    Xy = np.column_stack([Xm, np.asarray(y).ravel()])
    if learner == "ridge":
        m1, m2 = Ridge(alpha=alpha), Ridge(alpha=alpha)
    else:
        m1, m2 = LinearRegression(), LinearRegression()
    m1.fit(Xm, X[:, j])
    m2.fit(Xy, X[:, j])
    return m1.predict(Xm), m2.predict(Xy)


# --------------------------------------------------------------------------
# core Semi-knockoff sampling (Algorithms 1-4)
# --------------------------------------------------------------------------
def _resample(eps, rng, scheme="perm"):
    n = len(eps)
    if scheme == "perm":  # Algorithm 1-4:  "Sample pi a permutation"
        return eps[rng.permutation(n)]
    elif scheme == "iid":  # proof E.2:      U_1..U_n  iid U{1..n}
        return eps[rng.integers(0, n, size=n)]
    raise ValueError(scheme)


def sko_losses(X, y, j, nu_pred, rho_pred, predict, loss, rng, n_perm=1, scheme="perm"):
    """Return (L1, L2), each of length n: the paired loss populations.

    L1_i = l(m(X~1_i), y_i) with X~1_i^j = nu(X_i^-j) + eps_{1, pi1(i)}
    L2_i = l(m(X~2_i), y_i) with X~2_i^j = rho(X_i^-j, y_i) + eps_{2, pi2(i)}

    n_perm > 1 performs the Rao-Blackwellisation described in Section 5
    ("multiple permutations per sample"): the losses are averaged over n_perm
    independent permutations before the paired test.
    """
    xj = X[:, j]
    e1 = xj - nu_pred
    e2 = xj - rho_pred
    L1 = np.zeros(len(y))
    L2 = np.zeros(len(y))
    for _ in range(n_perm):
        X1 = X.copy()
        X2 = X.copy()
        X1[:, j] = nu_pred + _resample(e1, rng, scheme)
        X2[:, j] = rho_pred + _resample(e2, rng, scheme)
        L1 += loss(predict(X1), y)
        L2 += loss(predict(X2), y)
    return L1 / n_perm, L2 / n_perm


def sko_statistic(
    X, y, j, nu_pred, rho_pred, predict, loss, rng, n_perm=1, scheme="perm"
):
    """W^j_SKO = (1/n) sum_i [ l(m(X~1_i), y_i) - l(m(X~2_i), y_i) ]."""
    L1, L2 = sko_losses(
        X, y, j, nu_pred, rho_pred, predict, loss, rng, n_perm=n_perm, scheme=scheme
    )
    return float(np.mean(L1 - L2))


def sko_pvalue(
    X,
    y,
    j,
    nu_pred,
    rho_pred,
    predict,
    loss,
    rng,
    test="wilcoxon",
    n_perm=1,
    scheme="perm",
):
    """Algorithm 1 / 2: one-sided nonparametric paired test, H1: L1 > L2."""
    L1, L2 = sko_losses(
        X, y, j, nu_pred, rho_pred, predict, loss, rng, n_perm=n_perm, scheme=scheme
    )
    d = L1 - L2
    if test == "wilcoxon":
        nz = d[d != 0]
        if len(nz) == 0:
            return 1.0
        return float(
            stats.wilcoxon(nz, alternative="greater", zero_method="wilcox").pvalue
        )
    elif test == "sign":
        k = int(np.sum(d > 0))
        m = int(np.sum(d != 0))
        if m == 0:
            return 1.0
        return float(stats.binomtest(k, m, 0.5, alternative="greater").pvalue)
    raise ValueError(test)


# --------------------------------------------------------------------------
# knockoff threshold, Eq. (1)
# --------------------------------------------------------------------------
def knockoff_threshold(W, q):
    W = np.asarray(W, dtype=float)
    cands = np.sort(np.unique(np.abs(W[W != 0])))
    for t in cands:
        num = 1 + np.sum(W <= -t)
        den = max(1, np.sum(W >= t))
        if num / den <= q:
            return float(t)
    return np.inf


def knockoff_select(W, q):
    t = knockoff_threshold(W, q)
    return np.where(np.asarray(W) >= t)[0]


def fdp(selected, support):
    selected = np.asarray(selected, dtype=int)
    if len(selected) == 0:
        return 0.0
    false = np.sum(~np.isin(selected, support))
    return float(false) / len(selected)


def power(selected, support):
    support = np.asarray(support, dtype=int)
    if len(support) == 0:
        return np.nan
    return float(np.sum(np.isin(support, selected))) / len(support)


# --------------------------------------------------------------------------
# HRT (Tansey, Veitch, Zhang, Rabadan & Blei, JCGS 2022)
#   - REQUIRES a train/test split: the model is fit on train, the randomisation
#     test is run on the held-out test set.
# --------------------------------------------------------------------------
def hrt_pvalue(X_te, y_te, j, cond_mean_te, resid_pool, predict, loss, rng, K=200):
    """Holdout randomization test p-value for feature j.

    cond_mean_te : nu_hat(X_te^{-j}) fitted on the TRAIN split
    resid_pool   : residual pool X^j - nu_hat(X^{-j}) from the TRAIN split
    """
    t0 = float(np.mean(loss(predict(X_te), y_te)))
    n = X_te.shape[0]
    cnt = 0
    for _ in range(K):
        Xk = X_te.copy()
        Xk[:, j] = cond_mean_te + resid_pool[rng.integers(0, len(resid_pool), size=n)]
        tk = float(np.mean(loss(predict(Xk), y_te)))
        if tk <= t0:
            cnt += 1
    return (1.0 + cnt) / (K + 1.0)


# --------------------------------------------------------------------------
# data generators from the paper
# --------------------------------------------------------------------------
def ar1_cov(p, rho):
    i = np.arange(p)
    return rho ** np.abs(i[:, None] - i[None, :])


def gen_adjacent(n, p, rng, rho=0.6, sparsity=0.25, noise=1.0):
    """Figure 4 setting: X ~ N(0, Sigma), Sigma_ij = 0.6^|i-j|,
    y = beta'X + eps, first 0.25p coordinates of beta in [1, 2], rest 0."""
    Sigma = ar1_cov(p, rho)
    L = np.linalg.cholesky(Sigma)
    X = rng.standard_normal((n, p)) @ L.T
    k = int(round(sparsity * p))
    beta = np.zeros(p)
    beta[:k] = rng.uniform(1.0, 2.0, size=k)
    y = X @ beta + noise * rng.standard_normal(n)
    return X, y, beta, Sigma, np.arange(k)


def gen_masked(n, p, rng, rho=0.6):
    """Figure 5 setting: one relevant coordinate l, y = X_l + 0.5 eps1;
    a correlated NULL variable X_{l-1} = X_l + 0.5 eps2."""
    Sigma = ar1_cov(p, rho)
    L = np.linalg.cholesky(Sigma)
    X = rng.standard_normal((n, p)) @ L.T
    l = int(rng.integers(1, p))
    y = X[:, l] + 0.5 * rng.standard_normal(n)
    X[:, l - 1] = X[:, l] + 0.5 * rng.standard_normal(n)
    return X, y, l, np.array([l])