diff --git a/docs/api/calib/pyhealth.calib.predictionset.rst b/docs/api/calib/pyhealth.calib.predictionset.rst index fe445ea1b..740b1c87c 100644 --- a/docs/api/calib/pyhealth.calib.predictionset.rst +++ b/docs/api/calib/pyhealth.calib.predictionset.rst @@ -1,10 +1,20 @@ 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 ----------------- @@ -12,6 +22,7 @@ Available Methods :toctree: _autosummary :nosignatures: + pyhealth.calib.predictionset.BaseConformal pyhealth.calib.predictionset.LABEL pyhealth.calib.predictionset.SCRIB pyhealth.calib.predictionset.FavMac @@ -19,6 +30,14 @@ Available Methods 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 diff --git a/examples/conformal_eeg/test_tfm_tuev_inference.py b/examples/conformal_eeg/test_tfm_tuev_inference.py index d25eff7ac..d8350d46c 100644 --- a/examples/conformal_eeg/test_tfm_tuev_inference.py +++ b/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" ) diff --git a/pyhealth/calib/predictionset/base_conformal/__init__.py b/pyhealth/calib/predictionset/base_conformal/__init__.py index 54a47ea55..9dde35db5 100644 --- a/pyhealth/calib/predictionset/base_conformal/__init__.py +++ b/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,6 +175,12 @@ 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__( @@ -167,6 +188,7 @@ def __init__( 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 diff --git a/pyhealth/calib/predictionset/cluster/cluster_label.py b/pyhealth/calib/predictionset/cluster/cluster_label.py index f56a325fa..0c719973c 100644 --- a/pyhealth/calib/predictionset/cluster/cluster_label.py +++ b/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 diff --git a/pyhealth/calib/predictionset/cluster/neighborhood_label.py b/pyhealth/calib/predictionset/cluster/neighborhood_label.py index 2d9f2dc6d..36fdbaa0a 100644 --- a/pyhealth/calib/predictionset/cluster/neighborhood_label.py +++ b/pyhealth/calib/predictionset/cluster/neighborhood_label.py @@ -12,6 +12,11 @@ from pyhealth.calib.base_classes import SetPredictor from pyhealth.calib.predictionset.base_conformal import _query_weighted_quantile +from pyhealth.calib.predictionset.scores import ( + SUPPORTED_SCORE_TYPES, + all_class_conformity_scores, + true_class_conformity_scores, +) from pyhealth.calib.utils import extract_embeddings, prepare_numpy_dataset from pyhealth.models import BaseModel @@ -33,6 +38,12 @@ class NeighborhoodLabel(SetPredictor): k_neighbors: Number of nearest calibration neighbors. Default 50. lambda_L: Temperature for exponential weights; smaller => more localization. Default 100.0. + score_type: Conformity score to use: "threshold" (default, + conformity score = p(true class), Sadinle, Lei, and Wasserman + 2019) or "aps" (Adaptive Prediction Sets, Romano, Sesia, and + Candes 2020). See :mod:`pyhealth.calib.predictionset.scores`. + random_state: Optional int seed for the RNG used by + score_type="aps". Ignored for score_type="threshold". debug: If True, process fewer samples for faster iteration. Examples: @@ -61,6 +72,11 @@ class NeighborhoodLabel(SetPredictor): ... y_true, y_prob, metrics=["accuracy", "miscoverage_ps"], ... y_predset=extra["y_predset"] ... ) + >>> + >>> # Use APS instead of the default threshold score + >>> ncp_aps = NeighborhoodLabel( + ... model=model, alpha=0.1, k_neighbors=50, score_type="aps") + >>> ncp_aps.calibrate(cal_dataset=cal_ds, cal_embeddings=cal_embeddings) """ def __init__( @@ -69,6 +85,8 @@ def __init__( alpha: float, k_neighbors: int = 50, lambda_L: float = 100.0, + score_type: str = "threshold", + random_state: int | None = None, debug: bool = False, **kwargs, ) -> None: @@ -78,6 +96,11 @@ def __init__( raise NotImplementedError( "NeighborhoodLabel 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 @@ -87,6 +110,8 @@ def __init__( self.device = model.device self.debug = debug + self.score_type = score_type + self.rng = np.random.default_rng(random_state) if not (0.0 < alpha < 1.0): raise ValueError(f"alpha must be in (0, 1), got {alpha!r}") @@ -148,7 +173,9 @@ def calibrate( f"cal_dataset size {N}" ) - conformity_scores = y_prob[np.arange(N), y_true] + conformity_scores = true_class_conformity_scores( + y_prob, y_true, score_type=self.score_type, rng=self.rng + ) k = min(self.k_neighbors, N) self._nn = NearestNeighbors(n_neighbors=k, metric="euclidean").fit( @@ -223,7 +250,15 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: ) if pred["y_prob"].ndim > 1: th = th.view(-1, *([1] * (pred["y_prob"].ndim - 1))) - y_predset = pred["y_prob"] >= th + + y_prob_np = pred["y_prob"].detach().cpu().numpy() + conformity_scores = all_class_conformity_scores( + y_prob_np, score_type=self.score_type, rng=self.rng + ) + conformity_scores = torch.as_tensor( + conformity_scores, device=pred["y_prob"].device, dtype=pred["y_prob"].dtype + ) + y_predset = conformity_scores >= th # if threshold is high, include at least argmax empty = y_predset.sum(dim=1) == 0 if empty.any(): diff --git a/pyhealth/calib/predictionset/covariate/covariate_label.py b/pyhealth/calib/predictionset/covariate/covariate_label.py index e0d66ee90..3482e4e91 100644 --- a/pyhealth/calib/predictionset/covariate/covariate_label.py +++ b/pyhealth/calib/predictionset/covariate/covariate_label.py @@ -30,6 +30,11 @@ from pyhealth.calib.base_classes import SetPredictor from pyhealth.calib.calibration.kcal.kde import RBFKernelMean +from pyhealth.calib.predictionset.scores import ( + SUPPORTED_SCORE_TYPES, + all_class_conformity_scores, + true_class_conformity_scores, +) from pyhealth.calib.utils import prepare_numpy_dataset from pyhealth.datasets import get_dataloader from pyhealth.models import BaseModel @@ -263,6 +268,12 @@ class CovariateLabel(SetPredictor): distribution. Should be a callable that takes embeddings (numpy array) and returns density estimates. Used for KDE-based likelihood ratio weighting (CoDrug approach). + score_type: Conformity score to use: "threshold" (default, + conformity score = p(true class), Sadinle, Lei, and Wasserman + 2019) or "aps" (Adaptive Prediction Sets, Romano, Sesia, and + Candes 2020). See :mod:`pyhealth.calib.predictionset.scores`. + random_state: Optional int seed for the RNG used by + score_type="aps". Ignored for score_type="threshold". debug: Whether to use debug mode (processes fewer samples for faster iteration) @@ -330,6 +341,12 @@ class CovariateLabel(SetPredictor): >>> custom_weights = compute_custom_weights(val_data, test_data) >>> cal_model = CovariateLabel(model, alpha=0.1) >>> cal_model.calibrate(cal_dataset=val_data, cal_weights=custom_weights) + + **Example 3: APS instead of the default threshold score** + + >>> cal_model_aps = CovariateLabel(model, alpha=0.1, score_type="aps") + >>> cal_model_aps.calibrate(cal_dataset=val_data, + ... cal_embeddings=cal_embs, test_embeddings=test_embs) """ def __init__( @@ -338,6 +355,8 @@ def __init__( alpha: Union[float, np.ndarray], kde_test: Optional[Callable] = None, kde_cal: Optional[Callable] = None, + score_type: str = "threshold", + random_state: int | None = None, debug: bool = False, **kwargs, ) -> None: @@ -347,6 +366,11 @@ def __init__( raise NotImplementedError( "CovariateLabel 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 @@ -357,6 +381,8 @@ 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): @@ -485,8 +511,10 @@ def calibrate( # 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] + # Extract conformity scores (higher = more conforming) + conformity_scores = true_class_conformity_scores( + y_prob, y_true, score_type=self.score_type, rng=self.rng + ) # Compute weighted quantile thresholds if isinstance(self.alpha, float): @@ -523,8 +551,15 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: """ pred = self.model(**kwargs) - # Construct prediction set by thresholding probabilities - pred["y_predset"] = pred["y_prob"] > self.t + # Construct prediction set by thresholding conformity scores + y_prob = pred["y_prob"].detach().cpu().numpy() + conformity_scores = all_class_conformity_scores( + y_prob, score_type=self.score_type, rng=self.rng + ) + conformity_scores = torch.as_tensor( + conformity_scores, device=pred["y_prob"].device, dtype=pred["y_prob"].dtype + ) + pred["y_predset"] = conformity_scores > self.t return pred diff --git a/pyhealth/calib/predictionset/label.py b/pyhealth/calib/predictionset/label.py index 8ff87d070..b77934a28 100644 --- a/pyhealth/calib/predictionset/label.py +++ b/pyhealth/calib/predictionset/label.py @@ -9,14 +9,17 @@ """ -from typing import Dict, Union - import numpy as np import torch from torch.utils.data import Subset 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 prepare_numpy_dataset from pyhealth.models import BaseModel @@ -42,7 +45,16 @@ class LABEL(SetPredictor): :param model: A trained base model. :type model: BaseModel :param alpha: Target mis-coverage rate(s). - :type alpha: Union[float, np.ndarray] + :type alpha: float | np.ndarray + :param score_type: Nonconformity score to use: "threshold" (default, + the LAC score from Sadinle, Lei, and Wasserman 2019, NC score + = 1 - p(true class)) or "aps" (Adaptive Prediction Sets, Romano, + Sesia, and Candes 2020). See + :mod:`pyhealth.calib.predictionset.scores` for the exact formulas. + :type score_type: str + :param random_state: Optional int seed for the RNG used by + score_type="aps". Ignored for score_type="threshold". + :type random_state: int | None Examples: >>> from pyhealth.datasets import ISRUCDataset, split_by_patient, get_dataloader @@ -70,20 +82,37 @@ class LABEL(SetPredictor): ... y_predset=extra_output['y_predset']) ... ) {'accuracy': 0.709843241966832, 'miscoverage_ps': array([0.1499847 , 0.29997638, 0.14993964, 0.14994704, 0.14999252])} + >>> + >>> # Use APS instead of the default threshold score + >>> cal_model_aps = LABEL(model, 0.15, score_type="aps", random_state=0) + >>> cal_model_aps.calibrate(cal_dataset=test_data) """ def __init__( - self, model: BaseModel, alpha: Union[float, np.ndarray], debug=False, **kwargs + self, + model: BaseModel, + alpha: float | np.ndarray, + score_type: str = "threshold", + random_state: int | None = None, + debug=False, + **kwargs, ) -> None: super().__init__(model, **kwargs) if model.mode != "multiclass": raise NotImplementedError() + 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 # multiclass for param in model.parameters(): param.requires_grad = False self.model.eval() self.device = model.device self.debug = debug + self.score_type = score_type + self.rng = np.random.default_rng(random_state) if not isinstance(alpha, float): alpha = np.asarray(alpha) @@ -104,31 +133,37 @@ def calibrate(self, cal_dataset: Subset): y_true = cal_dataset["y_true"] N, K = cal_dataset["y_prob"].shape - # NC scores: 1 - p(true class); higher = less conforming + # NC scores: higher = less conforming + nc_scores = true_class_nc_scores( + y_prob, y_true, score_type=self.score_type, rng=self.rng + ) if isinstance(self.alpha, float): - t = _query_quantile( - 1.0 - y_prob[np.arange(N), y_true], self.alpha - ) + t = _query_quantile(nc_scores, self.alpha) else: t = [ - _query_quantile( - 1.0 - y_prob[y_true == k, k], self.alpha[k] - ) + _query_quantile(nc_scores[y_true == k], self.alpha[k]) for k in range(K) ] self.t = torch.tensor(t, device=self.device) - def forward(self, **kwargs) -> Dict[str, torch.Tensor]: + def forward(self, **kwargs) -> dict[str, torch.Tensor]: """Forward propagation (just like the original model). :return: A dictionary with all results from the base model, with the following updates: y_predset: a bool tensor representing the prediction for each class. - :rtype: Dict[str, torch.Tensor] + :rtype: dict[str, torch.Tensor] """ pred = self.model(**kwargs) - # Include class y if its NC score (1 - p(y)) <= NC threshold - 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 + pred["y_predset"] = nc_scores <= self.t return pred if __name__ == "__main__": diff --git a/pyhealth/calib/predictionset/scores.py b/pyhealth/calib/predictionset/scores.py new file mode 100644 index 000000000..ae4172d1d --- /dev/null +++ b/pyhealth/calib/predictionset/scores.py @@ -0,0 +1,221 @@ +"""Shared, pluggable conformity/nonconformity score functions. + +This module separates the *score* used by a conformal-prediction-set method +from the *calibration/thresholding procedure* it's plugged into. These are +two independent axes: the choice of score ("threshold"/LAC vs "aps") does +not depend on how the resulting scores get turned into a threshold (marginal +quantile, per-class quantile, per-cluster quantile, weighted quantile for +covariate shift, or localized weighted quantile for neighborhood methods). + +Supported score types: + + - "threshold" (a.k.a. LAC, Sadinle, Lei, and Wasserman 2019): the score + for class k is simply based on the model's predicted probability for + k. Simple, but not adaptive to how "peaked" or "flat" the predicted + distribution is. + + - "aps" (Adaptive Prediction Sets, Romano, Sesia, and Candes 2020; + score definition as restated in Angelopoulos, Bates, Malik, and + Jordan 2021, "Uncertainty Sets for Image Classifiers using Conformal + Prediction", Algorithm 2): the score for class k is the cumulative + sum of predicted probabilities for all classes ranked strictly above + k, plus a Uniform(0,1)-randomized fraction of k's own probability:: + + E(x, k) = sum_{j : pi(x,j) > pi(x,k)} pi(x,j) + U * pi(x,k) + + where pi(x, ·) are the model's predicted class probabilities and + U ~ Uniform(0,1) is drawn once per example and reused across every + candidate class k for that example (so the resulting prediction sets + are "nested": the set of included classes is always a prefix of the + classes sorted by decreasing probability). This adapts the set size + to the model's confidence for each individual input, which the + "threshold" score does not. + +Both scores are computed here in *nonconformity* convention (higher = less +conforming, i.e. 1 minus a probability-like quantity) since that's the +convention BaseConformal/LABEL/ClusterLabel use internally. A *conformity* +(higher = more conforming) variant is also provided for CovariateLabel/ +NeighborhoodLabel, which use the opposite sign convention internally; it is +simply `1 - nonconformity`, preserving the same ranking of examples either +way. +""" + +import numpy as np + +__all__ = [ + "SUPPORTED_SCORE_TYPES", + "all_class_conformity_scores", + "all_class_nc_scores", + "true_class_conformity_scores", + "true_class_nc_scores", +] + +SUPPORTED_SCORE_TYPES = ("threshold", "aps") + + +def _validate_score_type(score_type: str) -> None: + if score_type not in SUPPORTED_SCORE_TYPES: + raise ValueError( + f"Unknown score_type: {score_type!r}. Supported: " + f"{SUPPORTED_SCORE_TYPES}." + ) + + +def _aps_all_class_nc_scores( + y_prob: np.ndarray, + rng: np.random.Generator, + randomize: bool, +) -> np.ndarray: + """Computes the APS nonconformity score for every class, every row. + + E(x, k) = [sum of predicted probabilities for classes ranked strictly + above k] + U * p(x, k), with U ~ Uniform(0,1) drawn once per row and + reused across all classes in that row (Angelopoulos et al. 2021, + Algorithm 2/3, with the regularization term lambda=0, i.e. plain APS + rather than RAPS). + + Args: + y_prob: Predicted probabilities, shape (N, K). + rng: Random generator used to draw the per-row U ~ Uniform(0,1) + tie-breaking/adaptivity term. + randomize: If False, uses U=1 for every row (the conservative, + non-randomized variant: ties are broken by always including the + full probability mass of a class's own rank). If True (the + variant the APS/RAPS papers use for their reported results), + draws a genuine U ~ Uniform(0,1) per row. + + Returns: + Nonconformity scores of shape (N, K); higher means less conforming. + """ + n = y_prob.shape[0] + # Ties in probability are broken randomly by perturbing the sort key + # infinitesimally, per the APS paper's note that "label-ordering ties + # should be broken randomly" when probabilities aren't all distinct. + tie_break = rng.uniform(0.0, 1e-12, size=y_prob.shape) + order = np.argsort(-(y_prob + tie_break), axis=1) # descending, per row + sorted_probs = np.take_along_axis(y_prob, order, axis=1) + cumsum = np.cumsum(sorted_probs, axis=1) + # Sum of all classes ranked strictly above each rank r (0-indexed): + # cumsum up to and including r, minus r's own probability. + cumsum_excl_own = cumsum - sorted_probs + + if randomize: + u = rng.uniform(0.0, 1.0, size=(n, 1)) + else: + u = np.ones((n, 1)) + + sorted_scores = cumsum_excl_own + u * sorted_probs # (N, K), sorted order + + # Undo the sort to get back to original class-index order. + inverse_order = np.argsort(order, axis=1) + scores = np.take_along_axis(sorted_scores, inverse_order, axis=1) + return scores + + +def all_class_nc_scores( + y_prob: np.ndarray, + score_type: str = "threshold", + rng: np.random.Generator | None = None, + randomize: bool = True, +) -> np.ndarray: + """Nonconformity score (higher = less conforming) for every class. + + Args: + y_prob: Predicted probabilities, shape (N, K). + score_type: "threshold" (Sadinle, Lei, and Wasserman 2019) or "aps" + (Romano, Sesia, and Candes 2020). Default "threshold". + rng: Random generator, required (and only used) if score_type="aps" + and randomize=True. + randomize: Whether to use the randomized ("exact coverage") variant + of APS. Ignored for score_type="threshold". + + Returns: + Nonconformity scores of shape (N, K). + + Examples: + >>> import numpy as np + >>> from pyhealth.calib.predictionset.scores import all_class_nc_scores + >>> y_prob = np.array([[0.7, 0.2, 0.1], [0.3, 0.5, 0.2]]) + >>> all_class_nc_scores(y_prob, score_type="threshold") + array([[0.3, 0.8, 0.9], + [0.7, 0.5, 0.8]]) + >>> rng = np.random.default_rng(0) + >>> scores = all_class_nc_scores(y_prob, score_type="aps", rng=rng) + >>> np.round(scores, 2) + array([[0.42, 0.82, 0.96], + [0.72, 0.36, 0.95]]) + """ + _validate_score_type(score_type) + if score_type == "threshold": + return 1.0 - y_prob + # score_type == "aps" + if rng is None: + rng = np.random.default_rng() + return _aps_all_class_nc_scores(y_prob, rng, randomize) + + +def all_class_conformity_scores( + y_prob: np.ndarray, + score_type: str = "threshold", + rng: np.random.Generator | None = None, + randomize: bool = True, +) -> np.ndarray: + """Conformity score (higher = more conforming) for every class. + + Equivalent to ``1 - all_class_nc_scores(...)``: same ranking of + examples, just the sign convention used by CovariateLabel and + NeighborhoodLabel (which threshold with ``score >= t`` rather than + ``nc_score <= t``). + + Examples: + >>> import numpy as np + >>> from pyhealth.calib.predictionset.scores import all_class_conformity_scores + >>> y_prob = np.array([[0.7, 0.2, 0.1], [0.3, 0.5, 0.2]]) + >>> all_class_conformity_scores(y_prob, score_type="threshold") + array([[0.7, 0.2, 0.1], + [0.3, 0.5, 0.2]]) + """ + return 1.0 - all_class_nc_scores(y_prob, score_type, rng, randomize) + + +def true_class_nc_scores( + y_prob: np.ndarray, + y_true: np.ndarray, + score_type: str = "threshold", + rng: np.random.Generator | None = None, + randomize: bool = True, +) -> np.ndarray: + """Nonconformity score of the true class only, shape (N,). Used during + calibration, where only the true label's score is needed. + + Examples: + >>> import numpy as np + >>> from pyhealth.calib.predictionset.scores import true_class_nc_scores + >>> y_prob = np.array([[0.7, 0.2, 0.1], [0.3, 0.5, 0.2]]) + >>> y_true = np.array([0, 1]) + >>> true_class_nc_scores(y_prob, y_true, score_type="threshold") + array([0.3, 0.5]) + """ + scores = all_class_nc_scores(y_prob, score_type, rng, randomize) + n = len(y_true) + return scores[np.arange(n), y_true] + + +def true_class_conformity_scores( + y_prob: np.ndarray, + y_true: np.ndarray, + score_type: str = "threshold", + rng: np.random.Generator | None = None, + randomize: bool = True, +) -> np.ndarray: + """Conformity score of the true class only, shape (N,). + + Examples: + >>> import numpy as np + >>> from pyhealth.calib.predictionset.scores import true_class_conformity_scores + >>> y_prob = np.array([[0.7, 0.2, 0.1], [0.3, 0.5, 0.2]]) + >>> y_true = np.array([0, 1]) + >>> true_class_conformity_scores(y_prob, y_true, score_type="threshold") + array([0.7, 0.5]) + """ + return 1.0 - true_class_nc_scores(y_prob, y_true, score_type, rng, randomize) diff --git a/tests/core/test_cluster_label.py b/tests/core/test_cluster_label.py index 3b63ba7ed..ca2ab382d 100644 --- a/tests/core/test_cluster_label.py +++ b/tests/core/test_cluster_label.py @@ -377,6 +377,40 @@ def test_prediction_sets_nonempty(self): torch.all(set_sizes > 0), "Some prediction sets are empty" ) + def test_score_type_aps_runs_end_to_end(self): + """score_type='aps' should calibrate and produce non-empty, + correctly-typed prediction sets, just like the default 'threshold'.""" + cluster_model = ClusterLabel( + model=self.model, + alpha=0.3, + n_clusters=2, + random_state=42, + score_type="aps", + ) + + train_indices = [0, 1, 2, 3, 4, 5] + cal_indices = [6, 7, 8, 9, 10, 11] + train_dataset = self.dataset.subset(train_indices) + cal_dataset = self.dataset.subset(cal_indices) + + train_embeddings = self._get_embeddings(train_dataset) + cal_embeddings = self._get_embeddings(cal_dataset) + + cluster_model.calibrate( + cal_dataset=cal_dataset, + train_embeddings=train_embeddings, + cal_embeddings=cal_embeddings, + ) + + test_loader = get_dataloader(self.dataset, batch_size=2, shuffle=False) + with torch.no_grad(): + for data_batch in test_loader: + output = cluster_model(**data_batch) + self.assertEqual(output["y_predset"].dtype, torch.bool) + self.assertEqual(output["y_predset"].shape, output["y_prob"].shape) + set_sizes = output["y_predset"].sum(dim=1) + self.assertTrue(torch.all(set_sizes > 0)) + def test_calibrate_requires_train_embeddings(self): """Test that calibrate requires train_embeddings.""" cluster_model = ClusterLabel( diff --git a/tests/core/test_covariate_label.py b/tests/core/test_covariate_label.py index 14c38cd46..6f2bfd04a 100644 --- a/tests/core/test_covariate_label.py +++ b/tests/core/test_covariate_label.py @@ -307,6 +307,39 @@ def test_prediction_sets_nonempty(self): torch.all(set_sizes > 0), "Some prediction sets are empty" ) + def test_score_type_aps_runs_end_to_end(self): + """score_type='aps' should calibrate and produce non-empty, + correctly-typed prediction sets, just like the default 'threshold'.""" + cal_model = CovariateLabel( + model=self.model, + alpha=0.3, + kde_test=self.kde_test, + kde_cal=self.kde_cal, + score_type="aps", + random_state=42, + ) + + cal_indices = [0, 1, 2, 3] + cal_dataset = self.dataset.subset(cal_indices) + cal_embeddings = self._get_embeddings(cal_dataset) + test_embeddings = self._get_embeddings(self.dataset) + + cal_model.calibrate( + cal_dataset=cal_dataset, + cal_embeddings=cal_embeddings, + test_embeddings=test_embeddings, + ) + + test_indices = [4, 5] + test_dataset = self.dataset.subset(test_indices) + test_loader = get_dataloader(test_dataset, batch_size=2, shuffle=False) + + with torch.no_grad(): + for data_batch in test_loader: + output = cal_model(**data_batch) + self.assertEqual(output["y_predset"].dtype, torch.bool) + self.assertEqual(output["y_predset"].shape, output["y_prob"].shape) + def test_weighted_quantile_function(self): """Test the weighted quantile helper function.""" from pyhealth.calib.predictionset.covariate.covariate_label import ( diff --git a/tests/core/test_neighborhood_label.py b/tests/core/test_neighborhood_label.py index b33f4c4b0..0c812d141 100644 --- a/tests/core/test_neighborhood_label.py +++ b/tests/core/test_neighborhood_label.py @@ -132,6 +132,28 @@ def test_prediction_sets_nonempty_batch(self): set_sizes = out["y_predset"].sum(dim=1) self.assertTrue(torch.all(set_sizes > 0), "Prediction sets should be non-empty") + def test_score_type_aps_runs_end_to_end(self): + """score_type='aps' should calibrate and produce non-empty, + correctly-typed prediction sets, just like the default 'threshold' + (NeighborhoodLabel always guarantees non-empty sets via its own + argmax fallback, independent of score_type).""" + ncp = NeighborhoodLabel( + model=self.model, alpha=0.3, k_neighbors=2, lambda_L=100.0, + score_type="aps", random_state=42, + ) + cal_dataset = self.dataset.subset([2, 3, 4, 5]) + cal_emb = self._get_embeddings(cal_dataset) + ncp.calibrate(cal_dataset=cal_dataset, cal_embeddings=cal_emb) + + loader = get_dataloader(self.dataset, batch_size=2, shuffle=False) + with torch.no_grad(): + for batch in loader: + out = ncp(**batch) + self.assertEqual(out["y_predset"].dtype, torch.bool) + self.assertEqual(out["y_predset"].shape, out["y_prob"].shape) + set_sizes = out["y_predset"].sum(dim=1) + self.assertTrue(torch.all(set_sizes > 0)) + def test_calibrate_without_embeddings_extracts(self): ncp = NeighborhoodLabel(model=self.model, alpha=0.1, k_neighbors=2) cal_dataset = self.dataset.subset([3, 4, 5]) diff --git a/tests/core/test_scores.py b/tests/core/test_scores.py new file mode 100644 index 000000000..e6af15c69 --- /dev/null +++ b/tests/core/test_scores.py @@ -0,0 +1,163 @@ +"""Tests for pyhealth.calib.predictionset.scores: the shared score module +implementing both the "threshold" (LAC) and "aps" (Adaptive Prediction +Sets, Romano/Sesia/Candes 2020) nonconformity/conformity scores. +""" + +import unittest + +import numpy as np + +from pyhealth.calib.predictionset.scores import ( + SUPPORTED_SCORE_TYPES, + all_class_conformity_scores, + all_class_nc_scores, + true_class_conformity_scores, + true_class_nc_scores, +) + + +class TestScoresThreshold(unittest.TestCase): + """"threshold" is just 1 - p (and its complement), regardless of rng.""" + + def setUp(self): + self.y_prob = np.array([[0.7, 0.2, 0.1], [0.3, 0.5, 0.2]]) + self.y_true = np.array([0, 1]) + + def test_all_class_nc_scores(self): + np.testing.assert_allclose( + all_class_nc_scores(self.y_prob, score_type="threshold"), + 1.0 - self.y_prob, + ) + + def test_all_class_conformity_scores(self): + np.testing.assert_allclose( + all_class_conformity_scores(self.y_prob, score_type="threshold"), + self.y_prob, + ) + + def test_true_class_nc_scores(self): + np.testing.assert_allclose( + true_class_nc_scores(self.y_prob, self.y_true, score_type="threshold"), + [0.3, 0.5], + ) + + def test_true_class_conformity_scores(self): + np.testing.assert_allclose( + true_class_conformity_scores(self.y_prob, self.y_true, score_type="threshold"), + [0.7, 0.5], + ) + + def test_default_score_type_is_threshold(self): + """Backward compatibility: omitting score_type must match the old, + hardcoded 1 - p behavior every caller used before score_type + existed.""" + np.testing.assert_allclose( + all_class_nc_scores(self.y_prob), + 1.0 - self.y_prob, + ) + + +class TestScoresAPS(unittest.TestCase): + """Verify the APS score formula: E(x,k) = [sum of probs ranked above k] + + U * p(k), and its structural properties.""" + + def test_non_randomized_matches_hand_computation(self): + """With randomize=False (U=1), APS collapses to the cumulative sum + of sorted probabilities -- hand-computable exactly.""" + y_prob = np.array([[0.5, 0.3, 0.15, 0.05]]) + rng = np.random.default_rng(0) + scores = all_class_nc_scores(y_prob, score_type="aps", rng=rng, randomize=False) + # sorted descending: 0.5, 0.3, 0.15, 0.05 -> cumsum 0.5, 0.8, 0.95, 1.0 + np.testing.assert_allclose(scores, [[0.5, 0.8, 0.95, 1.0]]) + + def test_scores_bounded_in_unit_interval(self): + rng = np.random.default_rng(1) + n, k = 50, 6 + logits = rng.normal(size=(n, k)) + y_prob = np.exp(logits) / np.exp(logits).sum(1, keepdims=True) + scores = all_class_nc_scores(y_prob, score_type="aps", rng=rng) + self.assertTrue(np.all(scores >= 0.0)) + self.assertTrue(np.all(scores <= 1.0)) + + def test_higher_probability_class_has_lower_or_equal_nc_score(self): + """APS nonconformity score must be monotonically non-decreasing as + predicted probability decreases (higher-probability classes are + included in smaller/first-formed sets).""" + rng = np.random.default_rng(2) + y_prob = np.array([[0.6, 0.25, 0.1, 0.05]]) + scores = all_class_nc_scores(y_prob, score_type="aps", rng=rng, randomize=False)[0] + order = np.argsort(-y_prob[0]) + sorted_scores = scores[order] + self.assertTrue(np.all(np.diff(sorted_scores) >= -1e-12)) + + def test_reproducible_with_seeded_rng(self): + y_prob = np.array([[0.4, 0.35, 0.25]]) + s1 = all_class_nc_scores(y_prob, score_type="aps", rng=np.random.default_rng(42)) + s2 = all_class_nc_scores(y_prob, score_type="aps", rng=np.random.default_rng(42)) + np.testing.assert_allclose(s1, s2) + + def test_nc_and_conformity_are_complementary(self): + rng = np.random.default_rng(3) + y_prob = np.array([[0.5, 0.3, 0.2]]) + nc = all_class_nc_scores(y_prob, score_type="aps", rng=np.random.default_rng(3)) + conf = all_class_conformity_scores(y_prob, score_type="aps", rng=np.random.default_rng(3)) + np.testing.assert_allclose(nc, 1.0 - conf) + + def test_true_class_score_matches_all_class_indexing(self): + rng = np.random.default_rng(4) + y_prob = np.array([[0.5, 0.3, 0.2], [0.1, 0.6, 0.3]]) + y_true = np.array([1, 2]) + all_scores = all_class_nc_scores(y_prob, score_type="aps", rng=np.random.default_rng(4)) + true_scores = true_class_nc_scores(y_prob, y_true, score_type="aps", rng=np.random.default_rng(4)) + np.testing.assert_allclose(true_scores, all_scores[np.arange(2), y_true]) + + +class TestScoresCoverage(unittest.TestCase): + """The core statistical property: both score types must achieve + approximately the target marginal coverage under split conformal + calibration, for both marginal and class-conditional targets.""" + + def _query_quantile(self, nc_scores, alpha): + nc_scores = np.sort(nc_scores) + n = len(nc_scores) + loc = int(np.ceil((1 - alpha) * (n + 1))) - 1 + if loc >= n: + return np.inf + return float(nc_scores[loc]) + + def test_marginal_coverage_threshold_and_aps(self): + rng = np.random.default_rng(5) + n, k = 4000, 5 + logits = rng.normal(size=(n, k)) * 2 + y_prob = np.exp(logits) / np.exp(logits).sum(1, keepdims=True) + y_true = np.array([rng.choice(k, p=y_prob[i]) for i in range(n)]) + cal, test = slice(0, n // 2), slice(n // 2, n) + alpha = 0.1 + + for score_type in SUPPORTED_SCORE_TYPES: + cal_rng = np.random.default_rng(6) + nc_cal = true_class_nc_scores( + y_prob[cal], y_true[cal], score_type=score_type, rng=cal_rng + ) + t = self._query_quantile(nc_cal, alpha) + test_rng = np.random.default_rng(7) + nc_test = all_class_nc_scores(y_prob[test], score_type=score_type, rng=test_rng) + predset = nc_test <= t + covered = predset[np.arange(n // 2), y_true[test]] + coverage = covered.mean() + # Allow generous slack for finite-sample noise at N=2000. + self.assertGreaterEqual( + coverage, 1 - alpha - 0.05, + f"{score_type} marginal coverage {coverage:.3f} too far below target", + ) + + +class TestScoresValidation(unittest.TestCase): + def test_unknown_score_type_raises(self): + y_prob = np.array([[0.5, 0.5]]) + with self.assertRaises(ValueError): + all_class_nc_scores(y_prob, score_type="not_a_real_score_type") + + +if __name__ == "__main__": + unittest.main()