diff --git a/docs/api/calib.rst b/docs/api/calib.rst index 7599b2a10..7814b4719 100644 --- a/docs/api/calib.rst +++ b/docs/api/calib.rst @@ -21,7 +21,7 @@ confidence levels: - :class:`~pyhealth.calib.predictionset.LABEL`: Conformal prediction with bounded error - :class:`~pyhealth.calib.predictionset.SCRIB`: Class-specific risk control - :class:`~pyhealth.calib.predictionset.FavMac`: Value-maximizing sets with cost control -- :class:`~pyhealth.calib.predictionset.CovariateLabel`: Covariate shift adaptive conformal +- :class:`~pyhealth.calib.predictionset.CovariateLabel`: Covariate shift adaptive conformal prediction with a finite-sample correction for the calibration/test weighting - :class:`~pyhealth.calib.predictionset.ClusterLabel`: K-means cluster-based conformal prediction - :class:`~pyhealth.calib.predictionset.NeighborhoodLabel`: Neighborhood Conformal Prediction (NCP) diff --git a/examples/cxr/covid19cxr_conformal.py b/examples/cxr/covid19cxr_conformal.py index 1d1bf8508..47e2eebbc 100644 --- a/examples/cxr/covid19cxr_conformal.py +++ b/examples/cxr/covid19cxr_conformal.py @@ -171,7 +171,8 @@ print("\nCreating CovariateLabel predictor...") covariate_predictor = CovariateLabel(model=resnet, alpha=alpha) -# Calibrate with embeddings (KDEs will be fitted automatically) +# Calibrate with embeddings (KDEs will be fitted automatically). The +# calibration weights include a finite-sample correction for the test point. print("Calibrating CovariateLabel predictor...") print(" - Fitting KDEs for covariate shift correction...") covariate_predictor.calibrate( diff --git a/pyhealth/calib/predictionset/base_conformal/__init__.py b/pyhealth/calib/predictionset/base_conformal/__init__.py index 3451f9062..54a47ea55 100644 --- a/pyhealth/calib/predictionset/base_conformal/__init__.py +++ b/pyhealth/calib/predictionset/base_conformal/__init__.py @@ -109,9 +109,17 @@ class BaseConformal(SetPredictor): alpha: Target miscoverage rate(s). Can be: - float: marginal coverage P(Y not in C(X)) <= alpha - array: class-conditional P(Y not in C(X) | Y=k) <= alpha[k] - score_type: Type of conformity score to use. Options: - - "aps": Adaptive Prediction Sets (default, uses probability scores) - - "threshold": Simple threshold on probabilities + score_type: Type of conformity score to use. Currently only one score + is implemented: + - "threshold" (default): NC score = 1 - p(true class), the score + from Sadinle, Lei, and Wasserman (2019) ("LABEL"). + - "aps": accepted as a backward-compatible alias for + "threshold". Despite the name, this does **not** implement + Adaptive Prediction Sets (Romano, Sesia, and Candes 2020) -- + that method uses a different score (cumulative sorted class + probabilities) which is not implemented here. If you need + genuine APS, do not rely on this option; it is kept only so + existing calls with ``score_type="aps"`` keep working. debug: Whether to use debug mode (processes fewer samples) Examples: @@ -158,7 +166,7 @@ def __init__( self, model: BaseModel, alpha: Union[float, np.ndarray], - score_type: str = "aps", + score_type: str = "threshold", debug: bool = False, **kwargs, ) -> None: @@ -201,8 +209,7 @@ def _compute_nc_scores( Non-conformity scores of shape (N,) — higher means less conforming. """ N = len(y_true) - if self.score_type == "aps" or self.score_type == "threshold": - # NC score = 1 - p(true class); higher = less conforming + if self.score_type == "threshold" or self.score_type == "aps": scores = 1.0 - y_prob[np.arange(N), y_true] else: raise ValueError(f"Unknown score_type: {self.score_type}") diff --git a/pyhealth/calib/predictionset/covariate/covariate_label.py b/pyhealth/calib/predictionset/covariate/covariate_label.py index f90430be8..e0d66ee90 100644 --- a/pyhealth/calib/predictionset/covariate/covariate_label.py +++ b/pyhealth/calib/predictionset/covariate/covariate_label.py @@ -170,25 +170,49 @@ def _compute_likelihood_ratio( def _query_weighted_quantile( - scores: np.ndarray, alpha: float, weights: np.ndarray + scores: np.ndarray, + alpha: float, + weights: np.ndarray, + test_weight: float = 0.0, ) -> float: - """Compute weighted quantile of scores. + """Compute the weighted conformal quantile of scores. + + Implements the finite-sample correction for weighted conformal prediction + under covariate shift (Tibshirani et al. 2019). Args: - scores: Array of conformity scores - alpha: Quantile level (between 0 and 1) - weights: Weights for each score + scores: Array of conformity scores (higher = more conforming). + alpha: Quantile level (between 0 and 1). + weights: Un-normalized weights (likelihood ratios) for each score. + test_weight: Un-normalized weight representing the test point. + Reserves ``test_weight / (sum(weights) + test_weight)`` of + probability mass at conformity ``-inf``. Default 0.0 recovers the + old, uncorrected behavior (kept for backward compatibility with + direct callers of this helper). Returns: - The weighted alpha-quantile of scores + The weighted alpha-quantile of scores. Returns ``-inf`` if the + reserved test-point mass alone already meets or exceeds ``alpha``, + since there isn't enough calibration mass to justify a stricter, + finite threshold without risking under-coverage. """ - # Sort scores and corresponding weights sorted_indices = np.argsort(scores) sorted_scores = scores[sorted_indices] sorted_weights = weights[sorted_indices] - # Compute cumulative weights - cum_weights = np.cumsum(sorted_weights) / np.sum(sorted_weights) + total_weight = np.sum(sorted_weights) + test_weight + if total_weight <= 0: + return -np.inf + + p_test = test_weight / total_weight + if p_test >= alpha: + # Not enough calibration mass to reach the target coverage without + # dipping into the mass reserved for the test point itself: fall + # back to the maximally permissive (safe) threshold. + return -np.inf + + # Compute cumulative weights over the reserved-mass-inclusive total. + cum_weights = np.cumsum(sorted_weights) / total_weight # Find the index where cumulative weight exceeds alpha idx = np.searchsorted(cum_weights, alpha, side="left") @@ -197,7 +221,7 @@ def _query_weighted_quantile( if idx >= len(sorted_scores): idx = len(sorted_scores) - 1 - return sorted_scores[idx] + return float(sorted_scores[idx]) class CovariateLabel(SetPredictor): @@ -458,8 +482,7 @@ def calibrate( self.kde_test, self.kde_cal, X ) - # Normalize weights - weights = likelihood_ratios / np.sum(likelihood_ratios) + # Keep weights un-normalized here self._sum_cal_weights = np.sum(likelihood_ratios) # Extract conformity scores (probabilities of true class) @@ -467,8 +490,10 @@ def calibrate( # Compute weighted quantile thresholds if isinstance(self.alpha, float): - # Marginal coverage: single threshold - t = _query_weighted_quantile(conformity_scores, self.alpha, weights) + test_weight = float(np.mean(likelihood_ratios)) + t = _query_weighted_quantile( + conformity_scores, self.alpha, likelihood_ratios, test_weight + ) else: # Class-conditional coverage: one threshold per class t = [] @@ -476,11 +501,10 @@ def calibrate( mask = y_true == k if np.sum(mask) > 0: class_scores = conformity_scores[mask] - class_weights = weights[mask] - # Renormalize class weights - class_weights = class_weights / np.sum(class_weights) + class_weights = likelihood_ratios[mask] + class_test_weight = float(np.mean(class_weights)) t_k = _query_weighted_quantile( - class_scores, self.alpha[k], class_weights + class_scores, self.alpha[k], class_weights, class_test_weight ) else: # If no calibration examples, use -inf (include all) diff --git a/tests/core/test_covariate_label.py b/tests/core/test_covariate_label.py index b5aa0aaf6..14c38cd46 100644 --- a/tests/core/test_covariate_label.py +++ b/tests/core/test_covariate_label.py @@ -317,12 +317,59 @@ def test_weighted_quantile_function(self): weights = np.array([0.1, 0.2, 0.3, 0.2, 0.2]) alpha = 0.5 + # Default test_weight=0.0 preserves the old (uncorrected) behavior + # for any direct caller that doesn't opt into the finite-sample + # correction. quantile = _query_weighted_quantile(scores, alpha, weights) self.assertIsInstance(quantile, (float, np.floating)) self.assertGreaterEqual(quantile, scores.min()) self.assertLessEqual(quantile, scores.max()) + def test_weighted_quantile_reserves_test_point_mass(self): + """The finite-sample correction should recover the standard (N+1) + reserved-mass fraction in the no-shift limit (uniform weights, + test_weight = mean of calibration weights).""" + from pyhealth.calib.predictionset.covariate.covariate_label import ( + _query_weighted_quantile, + ) + + N = 6 + weights = np.ones(N) + test_weight = float(np.mean(weights)) + p_test = test_weight / (np.sum(weights) + test_weight) + + self.assertAlmostEqual(p_test, 1.0 / (N + 1), places=10) + + def test_weighted_quantile_small_calibration_set_is_conservative(self): + """With very few calibration examples relative to the requested + alpha, the corrected quantile should fall back to -inf (maximally + permissive / safe) rather than returning an overconfident finite + threshold, since there isn't enough calibration mass to support the + target coverage without dipping into the reserved test-point mass.""" + from pyhealth.calib.predictionset.covariate.covariate_label import ( + _query_weighted_quantile, + ) + + scores = np.array([0.4, 0.6]) # N=2 + weights = np.array([1.1, 1.1]) + test_weight = float(np.mean(weights)) + + # 1/(N+1) = 1/3 ~= 0.333 > alpha=0.3, so there isn't enough + # calibration mass to safely support this target. + result = _query_weighted_quantile(scores, 0.3, weights, test_weight) + self.assertEqual(result, -np.inf) + + # A larger, adequately-sized calibration set at the same alpha + # should NOT hit the same fallback. + scores_large = np.linspace(0.0, 1.0, 50) + weights_large = np.ones(50) + test_weight_large = float(np.mean(weights_large)) + result_large = _query_weighted_quantile( + scores_large, 0.3, weights_large, test_weight_large + ) + self.assertTrue(np.isfinite(result_large)) + def test_likelihood_ratio_function(self): """Test the likelihood ratio computation.""" from pyhealth.calib.predictionset.covariate.covariate_label import (