Implements real APS & dynamic scoring for Conformal Prediction

docs/api/calib/pyhealth.calib.predictionset.rst

@@ -1,24 +1,43 @@
 pyhealth.calib.predictionset
 ===================================
 
-Prediction set constructors that provide set-valued predictions with statistical 
-coverage guarantees. These methods are based on conformal prediction and related 
+Prediction set constructors that provide set-valued predictions with statistical
+coverage guarantees. These methods are based on conformal prediction and related
 techniques for uncertainty quantification.
 
+``BaseConformal``, ``LABEL``, ``ClusterLabel``, ``CovariateLabel``, and
+``NeighborhoodLabel`` all accept a ``score_type`` argument selecting the
+nonconformity/conformity score used for calibration and set construction:
+either ``"threshold"`` (the default, unchanged from prior releases) or
+``"aps"`` (Adaptive Prediction Sets, Romano, Sesia, and Candes 2020), which
+adapts the prediction set size to the model's per-input confidence. See
+:mod:`pyhealth.calib.predictionset.scores` for the exact score formulas.
+``SCRIB`` and ``FavMac`` are not included since their calibration
+procedures aren't a score-then-quantile pattern.
+
 Available Methods
 -----------------
 
 .. autosummary::
    :toctree: _autosummary
    :nosignatures:
 
+   pyhealth.calib.predictionset.BaseConformal
    pyhealth.calib.predictionset.LABEL
    pyhealth.calib.predictionset.SCRIB
    pyhealth.calib.predictionset.FavMac
    pyhealth.calib.predictionset.CovariateLabel
    pyhealth.calib.predictionset.ClusterLabel
    pyhealth.calib.predictionset.NeighborhoodLabel
 
+BaseConformal (Standard Split Conformal Prediction)
+----------------------------------------------------
+
+.. autoclass:: pyhealth.calib.predictionset.BaseConformal
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 LABEL (Least Ambiguous Set-valued Classifier)
 ----------------------------------------------
 
@@ -71,3 +90,14 @@ Helper Functions
 ----------------
 
 .. autofunction:: pyhealth.calib.predictionset.covariate.fit_kde
+
+Score Functions
+---------------
+
+Shared, pluggable nonconformity/conformity score implementations backing
+the ``score_type`` argument described above.
+
+.. autofunction:: pyhealth.calib.predictionset.scores.all_class_nc_scores
+.. autofunction:: pyhealth.calib.predictionset.scores.all_class_conformity_scores
+.. autofunction:: pyhealth.calib.predictionset.scores.true_class_nc_scores
+.. autofunction:: pyhealth.calib.predictionset.scores.true_class_conformity_scores

examples/conformal_eeg/test_tfm_tuev_inference.py

@@ -27,7 +27,7 @@
 TUEV_ROOT = "/srv/local/data/TUH/tuh_eeg_events/v2.0.0/edf/"
 
 REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-TOKENIZER_WEIGHTS  = os.path.join(REPO_ROOT, "weightfiles", "tfm_tokenizer_last.pth")
+TOKENIZER_WEIGHTS = os.path.join(REPO_ROOT, "weightfiles", "tfm_tokenizer_last.pth")
 CLASSIFIER_WEIGHTS_DIR = os.path.join(
     REPO_ROOT, "weightfiles", "TFM_Tokenizer_multiple_finetuned_on_TUEV"
 )

pyhealth/calib/predictionset/base_conformal/__init__.py

@@ -10,18 +10,30 @@
 Paper:
     Vovk, Vladimir, Alexander Gammerman, and Glenn Shafer.
     "Algorithmic learning in a random world." Springer, 2005.
-    
+
     Papadopoulos, Harris, Kostas Proedrou, Volodya Vovk, and Alex Gammerman.
     "Inductive confidence machines for regression." ECML 2002.
+
+    Sadinle, Mauricio, Jing Lei, and Larry Wasserman. "Least ambiguous
+    set-valued classifiers with bounded error levels." Journal of the
+    American Statistical Association (2019). [score_type="threshold"]
+
+    Romano, Yaniv, Matteo Sesia, and Emmanuel Candes. "Classification with
+    valid and adaptive coverage." NeurIPS 2020. [score_type="aps"]
 """
 
-from typing import Dict, Union
+from typing import Union
 
 import numpy as np
 import torch
 from torch.utils.data import IterableDataset
 
 from pyhealth.calib.base_classes import SetPredictor
+from pyhealth.calib.predictionset.scores import (
+    SUPPORTED_SCORE_TYPES,
+    all_class_nc_scores,
+    true_class_nc_scores,
+)
 from pyhealth.calib.utils import prepare_numpy_dataset
 from pyhealth.models import BaseModel
 
@@ -109,17 +121,20 @@ 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. Currently only one score
-            is implemented:
+        score_type: Type of nonconformity score to use:
             - "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.
+            - "aps": Adaptive Prediction Sets (Romano, Sesia, and Candes
+              2020). NC score for class k is the cumulative sum of predicted
+              probabilities for classes ranked above k, plus a randomized
+              U * p(k) term (U ~ Uniform(0,1), one draw per example, shared
+              across all candidate classes for that example). Unlike
+              "threshold", this adapts the prediction set size to how
+              peaked or flat the model's predicted distribution is for each
+              individual input. See :mod:`pyhealth.calib.predictionset.scores`
+              for the exact formula.
+        random_state: Optional int seed for the RNG used by score_type="aps"
+            (the U ~ Uniform(0,1) draws). Ignored for score_type="threshold".
         debug: Whether to use debug mode (processes fewer samples)
 
     Examples:
@@ -160,13 +175,20 @@ class BaseConformal(SetPredictor):
         >>> conformal_model_cc = BaseConformal(
         ...     model, alpha=[0.1, 0.15, 0.1, 0.1, 0.1])
         >>> conformal_model_cc.calibrate(cal_dataset=val_data)
+        >>>
+        >>> # Use APS instead of the default threshold score (adapts set
+        >>> # size to how confident the model is on each individual input)
+        >>> conformal_model_aps = BaseConformal(
+        ...     model, alpha=0.1, score_type="aps", random_state=0)
+        >>> conformal_model_aps.calibrate(cal_dataset=val_data)
     """
 
     def __init__(
         self,
         model: BaseModel,
         alpha: Union[float, np.ndarray],
         score_type: str = "threshold",
+        random_state: int | None = None,
         debug: bool = False,
         **kwargs,
     ) -> None:
@@ -176,6 +198,11 @@ def __init__(
             raise NotImplementedError(
                 "BaseConformal only supports multiclass classification"
             )
+        if score_type not in SUPPORTED_SCORE_TYPES:
+            raise ValueError(
+                f"Unknown score_type: {score_type!r}. Supported: "
+                f"{SUPPORTED_SCORE_TYPES}."
+            )
 
         self.mode = self.model.mode
 
@@ -187,6 +214,7 @@ def __init__(
         self.device = model.device
         self.debug = debug
         self.score_type = score_type
+        self.rng = np.random.default_rng(random_state)
 
         # Store alpha
         if not isinstance(alpha, float):
@@ -208,13 +236,9 @@ def _compute_nc_scores(
         Returns:
             Non-conformity scores of shape (N,) — higher means less conforming.
         """
-        N = len(y_true)
-        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}")
-
-        return scores
+        return true_class_nc_scores(
+            y_prob, y_true, score_type=self.score_type, rng=self.rng
+        )
 
     def calibrate(self, cal_dataset: IterableDataset):
         """Calibrate the thresholds for prediction set construction.
@@ -268,7 +292,7 @@ def calibrate(self, cal_dataset: IterableDataset):
         if self.debug:
             print(f"Calibrated thresholds: {self.t}")
 
-    def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
+    def forward(self, **kwargs) -> dict[str, torch.Tensor]:
         """Forward propagation with prediction set construction.
 
         Returns:
@@ -284,8 +308,15 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
 
         pred = self.model(**kwargs)
 
-        # Include class y if its NC score (1 - p(y)) <= NC threshold self.t
-        pred["y_predset"] = (1.0 - pred["y_prob"]) <= self.t
+        y_prob = pred["y_prob"].detach().cpu().numpy()
+        nc_scores = all_class_nc_scores(
+            y_prob, score_type=self.score_type, rng=self.rng
+        )
+        nc_scores = torch.as_tensor(
+            nc_scores, device=pred["y_prob"].device, dtype=pred["y_prob"].dtype
+        )
+        # Include class y if its NC score <= NC threshold self.t
+        pred["y_predset"] = nc_scores <= self.t
 
         return pred
 

pyhealth/calib/predictionset/cluster/cluster_label.py

@@ -19,6 +19,11 @@
 
 from pyhealth.calib.base_classes import SetPredictor
 from pyhealth.calib.predictionset.base_conformal import _query_quantile
+from pyhealth.calib.predictionset.scores import (
+    SUPPORTED_SCORE_TYPES,
+    all_class_nc_scores,
+    true_class_nc_scores,
+)
 from pyhealth.calib.utils import extract_embeddings, prepare_numpy_dataset
 from pyhealth.models import BaseModel
 
@@ -44,7 +49,13 @@ class ClusterLabel(SetPredictor):
             - float: marginal coverage P(Y not in C(X)) <= alpha
             - array: class-conditional P(Y not in C(X) | Y=k) <= alpha[k]
         n_clusters: Number of K-means clusters. Default is 5.
-        random_state: Random seed for K-means clustering. Default is 42.
+        random_state: Random seed for K-means clustering, and (if
+            score_type="aps") for the score's U ~ Uniform(0,1) draws.
+            Default is 42.
+        score_type: Nonconformity score to use: "threshold" (default,
+            NC score = 1 - p(true class), Sadinle, Lei, and Wasserman 2019)
+            or "aps" (Adaptive Prediction Sets, Romano, Sesia, and Candes
+            2020). See :mod:`pyhealth.calib.predictionset.scores`.
         debug: Whether to use debug mode (processes fewer samples for
             faster iteration)
 
@@ -89,6 +100,15 @@ class ClusterLabel(SetPredictor):
         ...     y_true, y_prob, metrics=["accuracy", "miscoverage_ps"],
         ...     y_predset=extra["y_predset"]
         ... )
+        >>>
+        >>> # Use APS instead of the default threshold score
+        >>> cluster_predictor_aps = ClusterLabel(
+        ...     model=model, alpha=0.1, n_clusters=5, score_type="aps")
+        >>> cluster_predictor_aps.calibrate(
+        ...     cal_dataset=cal_ds,
+        ...     train_embeddings=train_embeddings,
+        ...     cal_embeddings=cal_embeddings,
+        ... )
     """
 
     def __init__(
@@ -97,6 +117,7 @@ def __init__(
         alpha: Union[float, np.ndarray],
         n_clusters: int = 5,
         random_state: int = 42,
+        score_type: str = "threshold",
         debug: bool = False,
         **kwargs,
     ) -> None:
@@ -106,6 +127,11 @@ def __init__(
             raise NotImplementedError(
                 "ClusterLabel only supports multiclass classification"
             )
+        if score_type not in SUPPORTED_SCORE_TYPES:
+            raise ValueError(
+                f"Unknown score_type: {score_type!r}. Supported: "
+                f"{SUPPORTED_SCORE_TYPES}."
+            )
 
         self.mode = self.model.mode
 
@@ -116,6 +142,7 @@ def __init__(
 
         self.device = model.device
         self.debug = debug
+        self.score_type = score_type
 
         # Store alpha
         if not isinstance(alpha, float):
@@ -129,6 +156,7 @@ def __init__(
             )
         self.n_clusters = n_clusters
         self.random_state = random_state
+        self.rng = np.random.default_rng(random_state)
 
         # Will be set during calibration
         self.kmeans_model = None
@@ -215,7 +243,9 @@ def calibrate(
         print(f"Cluster assignments: {np.bincount(cal_cluster_labels)}")
 
         # Compute non-conformity scores (higher = less conforming)
-        conformity_scores = 1.0 - y_prob[np.arange(N), y_true]
+        conformity_scores = true_class_nc_scores(
+            y_prob, y_true, score_type=self.score_type, rng=self.rng
+        )
 
         # Compute cluster-specific thresholds
         self.cluster_thresholds = {}
@@ -313,8 +343,15 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
             )
             cluster_thresholds = cluster_thresholds.view(view_shape)
 
-        # Include class y if its NC score (1 - p(y)) <= NC threshold
-        pred["y_predset"] = (1.0 - pred["y_prob"]) <= cluster_thresholds
+        # Include class y if its NC score <= NC threshold
+        y_prob_np = pred["y_prob"].detach().cpu().numpy()
+        nc_scores = all_class_nc_scores(
+            y_prob_np, score_type=self.score_type, rng=self.rng
+        )
+        nc_scores = torch.as_tensor(
+            nc_scores, device=pred["y_prob"].device, dtype=pred["y_prob"].dtype
+        )
+        pred["y_predset"] = nc_scores <= cluster_thresholds
         pred.pop("embed", None)  # do not expose internal embedding to caller
         return pred