Covariate Shift Conformal Prediction Fixes (#1180)

* Covariate CP fixes

* small edits to pass checks

Author: lehendo

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)
 

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(

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}")

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,29 +482,29 @@ 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)
         conformity_scores = y_prob[np.arange(N), y_true]
 
         # 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 = []
             for k in range(K):
                 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)

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 (