File size: 4,296 Bytes
fc329a3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | """Symmetrized full conformal with local normalization (Theorem 4.1)."""
import numpy as np
from sklearn.neighbors import NearestNeighbors
from .base import ConformalResult
from ..utils.simplex import ilr
def full_conformal(
R_cal: np.ndarray,
R_test: np.ndarray,
alpha: float,
U_cal: np.ndarray,
U_test: np.ndarray,
k: int = 20,
) -> ConformalResult:
"""Symmetrized full conformal prediction with local normalization.
For each test point j, constructs augmented set = cal ∪ {j}, computes
LOO sigma for all n_cal+1 points in the augmented set, then derives
the conformal p-value. Exact per Theorem 4.1.
Optimization: cal-point LOO sigmas are pre-computed on the cal set and
updated only when the test point falls among a cal point's k nearest
neighbors (rare for large n_cal).
Args:
R_cal: calibration residuals (n_cal,)
R_test: test residuals (n_test,)
alpha: miscoverage level
U_cal: calibration predictions (n_cal, K)
U_test: test predictions (n_test, K)
k: kNN neighbors for leave-one-out scale estimation
Returns:
ConformalResult with exact marginal coverage guarantee.
"""
n_cal = len(R_cal)
n_test = len(R_test)
if n_cal < 2:
radius = np.full(n_test, np.inf, dtype=float)
return ConformalResult(covered=np.ones(n_test, dtype=bool), radius=radius, threshold=alpha)
k_loo = min(k, n_cal - 1)
Z_cal = ilr(U_cal)
Z_test = ilr(U_test)
# Pre-compute cal LOO sigmas (reused across test points)
nn_cal = NearestNeighbors(n_neighbors=k_loo + 1).fit(Z_cal)
cal_dists, cal_nn_idx = nn_cal.kneighbors(Z_cal)
cal_loo_sigma = np.zeros(n_cal)
for i in range(n_cal):
neighbor_idx = cal_nn_idx[i][cal_nn_idx[i] != i][:k_loo]
cal_loo_sigma[i] = max(np.median(R_cal[neighbor_idx]), 1e-8)
# Pre-compute cal normalized scores (base, without augmentation effect)
S_cal_base = R_cal / cal_loo_sigma
covered = np.zeros(n_test, dtype=bool)
radius = np.zeros(n_test)
for j in range(n_test):
# For the test point: its LOO sigma is computed using cal as reference
# (in the augmented set, dropping itself leaves exactly cal)
test_dists, test_nn_idx = nn_cal.kneighbors(Z_test[j:j + 1])
test_nbrs = test_nn_idx[0][:k_loo]
sigma_test_j = max(np.median(R_cal[test_nbrs]), 1e-8)
# For cal points: check if adding the test point changes their LOO sigma
# A cal point's LOO sigma changes only if the test point is closer than
# its k-th nearest cal neighbor. For efficiency, compute the updated
# sigma only for affected cal points.
S_cal = S_cal_base.copy()
# Distance from each cal point to this test point
d_to_test = np.linalg.norm(Z_cal - Z_test[j], axis=1)
# k-th neighbor distance for each cal point (last column of cal_dists,
# but we used k_loo+1 neighbors including self, so k-th non-self is index k_loo)
kth_dist = cal_dists[:, k_loo] # distance to (k_loo+1)-th neighbor (including self)
affected = d_to_test < kth_dist
if np.any(affected):
for i in np.where(affected)[0]:
# Recompute LOO sigma: drop self from cal, add test point
orig_nbrs = cal_nn_idx[i][cal_nn_idx[i] != i][:k_loo]
# Replace the farthest neighbor with test point if closer
nbr_residuals = list(R_cal[orig_nbrs])
# Drop the farthest cal neighbor, add test residual
nbr_residuals[-1] = R_test[j]
new_sigma = max(np.median(nbr_residuals), 1e-8)
S_cal[i] = R_cal[i] / new_sigma
S_test_j = R_test[j] / sigma_test_j
# p-value: fraction of augmented scores >= test score
n_geq = np.sum(S_cal >= S_test_j) + 1 # +1 for test point itself
p_val = n_geq / (n_cal + 1)
covered[j] = p_val > alpha
# Radius: effective threshold
all_S = np.concatenate([S_cal, [S_test_j]])
q_idx = int(np.ceil((1 - alpha) * (n_cal + 1))) - 1
radius[j] = sigma_test_j * np.sort(all_S)[min(q_idx, len(all_S) - 1)]
return ConformalResult(covered=covered, radius=radius, threshold=alpha)
|