[CPBench] Benchmarking Wrapper

pyhealth/calib/base_classes.py

@@ -23,11 +23,34 @@ def to(self, device):
 
 
 class SetPredictor(ABC, torch.nn.Module):
+    """Base class for conformal-prediction-style prediction-set constructors.
+
+    Interface convention for `calibrate()`: every concrete subclass should
+    accept `calibrate(self, cal_dataset, train_dataset=None, test_dataset=None,
+    **kwargs)`. `cal_dataset` is required; `train_dataset`/`test_dataset` are
+    optional and only meaningful to methods that need extra data beyond the
+    calibration set (e.g. embeddings-based methods like ClusterLabel,
+    CovariateLabel, NeighborhoodLabel) -- methods that don't need them should
+    still accept and ignore the two keyword arguments rather than omitting
+    them from the signature.
+
+    This lets a generic caller (e.g. a benchmark harness) calibrate *any*
+    SetPredictor -- including ones not yet written -- with the same call:
+    `cp_model.calibrate(cal_dataset=cal, train_dataset=train, test_dataset=test)`,
+    without needing to special-case which extra data each method requires.
+    Methods that need embeddings should extract them internally (e.g. via
+    `pyhealth.calib.utils.extract_embeddings`) from whichever of
+    train_dataset/test_dataset they're given, rather than requiring the
+    caller to pre-compute and pass embeddings under method-specific kwarg
+    names -- pre-computed embeddings may still be accepted as additional
+    optional kwargs for callers who want to reuse/precompute them.
+    """
+
     def __init__(self, model, **kwargs) -> None:
         super().__init__()
         self.model = model
 
-    def calibrate(self, cal_dataset):
+    def calibrate(self, cal_dataset, train_dataset=None, test_dataset=None):
         ...
 
     def forward(self, **kwargs) -> Dict[str, torch.Tensor]:

pyhealth/calib/predictionset/base_conformal/__init__.py

@@ -15,7 +15,7 @@
     "Inductive confidence machines for regression." ECML 2002.
 """
 
-from typing import Dict, Union
+from typing import Dict, Optional, Union
 
 import numpy as np
 import torch
@@ -209,11 +209,22 @@ def _compute_nc_scores(
 
         return scores
 
-    def calibrate(self, cal_dataset: IterableDataset):
+    def calibrate(
+        self,
+        cal_dataset: IterableDataset,
+        train_dataset: Optional[IterableDataset] = None,
+        test_dataset: Optional[IterableDataset] = None,
+    ):
         """Calibrate the thresholds for prediction set construction.
 
         Args:
             cal_dataset: Calibration set (held-out validation data)
+            train_dataset: Unused by this method. Accepted only so callers
+                that generically calibrate any SetPredictor (e.g. CPBench)
+                can pass the same (cal_dataset, train_dataset, test_dataset)
+                arguments to every CP method without branching on which
+                extra data each one needs.
+            test_dataset: Unused by this method (see train_dataset).
         """
         # Get predictions and true labels
         cal_dataset_dict = prepare_numpy_dataset(

pyhealth/calib/predictionset/cluster/cluster_label.py

@@ -137,6 +137,8 @@ def __init__(
     def calibrate(
         self,
         cal_dataset: IterableDataset,
+        train_dataset: Optional[IterableDataset] = None,
+        test_dataset: Optional[IterableDataset] = None,
         train_embeddings: Optional[np.ndarray] = None,
         cal_embeddings: Optional[np.ndarray] = None,
         batch_size: int = 32,
@@ -151,14 +153,23 @@ def calibrate(
 
         Args:
             cal_dataset: Calibration set
+            train_dataset: Training set to extract embeddings from, used
+                only if train_embeddings is not already provided. ClusterLabel
+                needs a pool of training-set embeddings to fit clusters on --
+                one of train_dataset or train_embeddings is required.
+            test_dataset: Unused by this method. Accepted only so callers
+                that generically calibrate any SetPredictor (e.g. CPBench)
+                can pass the same (cal_dataset, train_dataset, test_dataset)
+                arguments to every CP method without branching on which
+                extra data each one needs.
             train_embeddings: Optional pre-computed training embeddings
-                of shape (n_train, embedding_dim). If not provided, will
-                be extracted from the model (requires train_dataset parameter).
+                of shape (n_train, embedding_dim). If not provided, will be
+                extracted from train_dataset.
             cal_embeddings: Optional pre-computed calibration embeddings
                 of shape (n_cal, embedding_dim). If not provided, will be
                 extracted from cal_dataset.
             batch_size: Batch size for embedding extraction when
-                cal_embeddings is not provided. Default is 32.
+                cal_embeddings/train_embeddings are not provided. Default 32.
 
         Note:
             Either provide embeddings directly or ensure the model supports
@@ -186,9 +197,16 @@ def calibrate(
             cal_embeddings = np.asarray(cal_embeddings)
 
         if train_embeddings is None:
-            raise ValueError(
-                "train_embeddings must be provided. "
-                "Extract embeddings from training set using extract_embeddings()."
+            if train_dataset is None:
+                raise ValueError(
+                    "ClusterLabel needs a pool of training-set embeddings to "
+                    "fit clusters on. Provide either train_embeddings "
+                    "directly, or train_dataset so they can be extracted "
+                    "automatically."
+                )
+            print("Extracting embeddings from training set...")
+            train_embeddings = extract_embeddings(
+                self.model, train_dataset, batch_size=batch_size, device=self.device
             )
         else:
             train_embeddings = np.asarray(train_embeddings)

pyhealth/calib/predictionset/cluster/neighborhood_label.py

@@ -107,6 +107,8 @@ def __init__(
     def calibrate(
         self,
         cal_dataset: IterableDataset,
+        train_dataset: Optional[IterableDataset] = None,
+        test_dataset: Optional[IterableDataset] = None,
         cal_embeddings: Optional[np.ndarray] = None,
         batch_size: int = 32,
     ) -> None:
@@ -120,6 +122,12 @@ def calibrate(
         Args:
             cal_dataset: Calibration dataset (for labels and predictions if
                 cal_embeddings not provided).
+            train_dataset: Unused by this method. Accepted only so callers
+                that generically calibrate any SetPredictor (e.g. CPBench)
+                can pass the same (cal_dataset, train_dataset, test_dataset)
+                arguments to every CP method without branching on which
+                extra data each one needs.
+            test_dataset: Unused by this method (see train_dataset).
             cal_embeddings: Optional precomputed calibration embeddings
                 (n_cal, embedding_dim). If None, extracted from cal_dataset.
             batch_size: Batch size for embedding extraction when cal_embeddings

pyhealth/calib/predictionset/covariate/covariate_label.py

@@ -30,7 +30,7 @@
 
 from pyhealth.calib.base_classes import SetPredictor
 from pyhealth.calib.calibration.kcal.kde import RBFKernelMean
-from pyhealth.calib.utils import prepare_numpy_dataset
+from pyhealth.calib.utils import extract_embeddings, prepare_numpy_dataset
 from pyhealth.datasets import get_dataloader
 from pyhealth.models import BaseModel
 
@@ -354,50 +354,71 @@ def __init__(
     def calibrate(
         self,
         cal_dataset: IterableDataset,
+        train_dataset: Optional[IterableDataset] = None,
+        test_dataset: Optional[IterableDataset] = None,
         cal_embeddings: Optional[np.ndarray] = None,
         test_embeddings: Optional[np.ndarray] = None,
         cal_weights: Optional[np.ndarray] = None,
+        batch_size: int = 32,
     ):
         """Calibrate the thresholds with covariate shift correction.
 
         This method supports three approaches for handling covariate shift:
-
-        1. **KDE-based (CoDrug approach)**: Provide cal_embeddings and 
-           test_embeddings (and optionally kde_test/kde_cal). The method will
+
+        1. **KDE-based (CoDrug approach)**: Provide cal_embeddings/test_embeddings
+           (or cal_dataset/test_dataset, from which they'll be extracted
+           automatically) and, optionally, kde_test/kde_cal. The method will
            use kernel density estimation to compute likelihood ratios.
-           
+
         2. **Custom weights**: Directly provide cal_weights computed from your
            own covariate shift correction method (e.g., importance sampling,
            propensity scores, discriminator-based methods, etc.).
-           
+
         3. **Pre-fitted KDEs**: Provide kde_test and kde_cal during initialization
-           along with cal_embeddings here.
+           along with cal_embeddings (or cal_dataset) here.
 
         Args:
-            cal_dataset: Calibration set
+            cal_dataset: Calibration set. Also used to extract cal_embeddings
+                automatically when cal_embeddings is not explicitly given.
+            train_dataset: Unused by this method. Accepted only so callers
+                that generically calibrate any SetPredictor (e.g. CPBench)
+                can pass the same (cal_dataset, train_dataset, test_dataset)
+                arguments to every CP method without branching on which
+                extra data each one needs.
+            test_dataset: Test set to extract test_embeddings from, used only
+                if test_embeddings is not already provided and no KDEs are
+                set. CovariateLabel needs test-side density information to
+                correct for covariate shift -- one of test_dataset,
+                test_embeddings, or cal_weights is required (unless kde_test/
+                kde_cal were already provided at __init__).
             cal_embeddings: Optional pre-computed calibration embeddings
-                of shape (n_cal, embedding_dim). If provided along with
-                test_embeddings and KDEs are not set, will be used to
-                compute likelihood ratios via KDE (CoDrug approach).
+                of shape (n_cal, embedding_dim). If not provided, extracted
+                from cal_dataset.
             test_embeddings: Optional pre-computed test embeddings
-                of shape (n_test, embedding_dim). Used with cal_embeddings
-                for KDE-based likelihood ratio computation.
+                of shape (n_test, embedding_dim). If not provided, extracted
+                from test_dataset when given.
             cal_weights: Optional custom weights for calibration samples
                 of shape (n_cal,). If provided, these weights will be used
                 directly instead of computing likelihood ratios via KDE.
                 Weights should represent importance weights or likelihood ratios
                 p_test(x) / p_cal(x). These will be normalized internally.
+            batch_size: Batch size for embedding extraction when
+                cal_embeddings/test_embeddings are not provided. Default 32.
 
         Note:
             You must provide ONE of:
             1. cal_weights (custom weights), OR
             2. kde_test and kde_cal during initialization, OR
-            3. cal_embeddings and test_embeddings here
+            3. test_embeddings or test_dataset here
 
         Examples:
-            >>> # Approach 1: KDE-based (CoDrug)
-            >>> model.calibrate(cal_dataset, cal_embeddings, test_embeddings)
-            >>> 
+            >>> # Approach 1: KDE-based (CoDrug), embeddings extracted for you
+            >>> model.calibrate(cal_dataset, test_dataset=test_dataset)
+            >>>
+            >>> # Approach 1b: KDE-based (CoDrug), pre-computed embeddings
+            >>> model.calibrate(cal_dataset, cal_embeddings=cal_embeddings,
+            ...     test_embeddings=test_embeddings)
+            >>>
             >>> # Approach 2: Custom weights (e.g., from importance sampling)
             >>> custom_weights = compute_importance_weights(cal_data, test_data)
             >>> model.calibrate(cal_dataset, cal_weights=custom_weights)
@@ -425,37 +446,41 @@ def calibrate(
             likelihood_ratios = np.asarray(cal_weights, dtype=np.float64)
             print("Using custom calibration weights")
         else:
-            # Use KDE-based approach (CoDrug method)
+            # Use KDE-based approach (CoDrug method). cal_dataset is always
+            # available, so cal-side embeddings can always be obtained.
+            if cal_embeddings is None:
+                print("Extracting embeddings from calibration set...")
+                cal_embeddings = extract_embeddings(
+                    self.model, cal_dataset, batch_size=batch_size, device=self.device
+                )
+            else:
+                cal_embeddings = np.asarray(cal_embeddings)
+
             # Check if we have KDEs
             if self.kde_test is None or self.kde_cal is None:
-                if cal_embeddings is None or test_embeddings is None:
+                if test_embeddings is None and test_dataset is not None:
+                    print("Extracting embeddings from test set...")
+                    test_embeddings = extract_embeddings(
+                        self.model, test_dataset, batch_size=batch_size,
+                        device=self.device,
+                    )
+                if test_embeddings is None:
                     raise ValueError(
-                        "Must provide ONE of:\n"
+                        "CovariateLabel needs test-side density information "
+                        "to correct for covariate shift. Provide ONE of:\n"
                         "  1. cal_weights (custom weights), OR\n"
                         "  2. kde_test and kde_cal during __init__, OR\n"
-                        "  3. cal_embeddings and test_embeddings during calibrate()"
+                        "  3. test_embeddings or test_dataset during calibrate()"
                     )
 
-                # Fit KDEs if embeddings provided
+                # Fit KDEs on the (now available) embeddings
                 print("Fitting KDEs on provided embeddings (CoDrug approach)...")
                 self.kde_cal, self.kde_test = fit_kde(cal_embeddings, test_embeddings)
 
-            # Use provided embeddings or extract from calibration data
-            if cal_embeddings is not None:
-                X = cal_embeddings
-            else:
-                # KDEs should already be provided in this case
-                # We just need to get the embeddings for likelihood ratio
-                # This assumes the model outputs embeddings
-                raise NotImplementedError(
-                    "Automatic embedding extraction not yet supported. "
-                    "Please provide cal_embeddings and test_embeddings."
-                )
-
             # Compute likelihood ratios using KDE
             print("Computing likelihood ratios via KDE...")
             likelihood_ratios = _compute_likelihood_ratio(
-                self.kde_test, self.kde_cal, X
+                self.kde_test, self.kde_cal, cal_embeddings
             )
 
         # Normalize weights

pyhealth/calib/predictionset/favmac/__init__.py

@@ -197,11 +197,17 @@ def __init__(
         self.target_cost = target_cost
         self.delta = delta
 
-    def calibrate(self, cal_dataset):
+    def calibrate(self, cal_dataset, train_dataset=None, test_dataset=None):
         """Calibrate the cost-control procedure.
 
         :param cal_dataset: Calibration set.
         :type cal_dataset: Subset
+        :param train_dataset: Unused by this method. Accepted only so
+            callers that generically calibrate any SetPredictor (e.g.
+            CPBench) can pass the same (cal_dataset, train_dataset,
+            test_dataset) arguments to every CP method without branching on
+            which extra data each one needs.
+        :param test_dataset: Unused by this method (see train_dataset).
         """
         _cal_data = prepare_numpy_dataset(
             self.model, cal_dataset, ["logit", "y_true"], debug=self.debug

pyhealth/calib/predictionset/label.py

@@ -9,7 +9,7 @@
 
 """
 
-from typing import Dict, Union
+from typing import Dict, Optional, Union
 
 import numpy as np
 import torch
@@ -91,11 +91,22 @@ def __init__(
 
         self.t = None
 
-    def calibrate(self, cal_dataset: Subset):
+    def calibrate(
+        self,
+        cal_dataset: Subset,
+        train_dataset: Optional[Subset] = None,
+        test_dataset: Optional[Subset] = None,
+    ):
         """Calibrate the thresholds used to construct the prediction set.
 
         :param cal_dataset: Calibration set.
         :type cal_dataset: Subset
+        :param train_dataset: Unused by this method. Accepted only so
+            callers that generically calibrate any SetPredictor (e.g.
+            CPBench) can pass the same (cal_dataset, train_dataset,
+            test_dataset) arguments to every CP method without branching on
+            which extra data each one needs.
+        :param test_dataset: Unused by this method (see train_dataset).
         """
         cal_dataset = prepare_numpy_dataset(
             self.model, cal_dataset, ["y_prob", "y_true"], debug=self.debug

pyhealth/calib/predictionset/scrib/__init__.py

@@ -267,11 +267,17 @@ def __init__(
 
         self.t = None
 
-    def calibrate(self, cal_dataset):
+    def calibrate(self, cal_dataset, train_dataset=None, test_dataset=None):
         """Calibrate/Search for the thresholds used to construct the prediction set.
 
         :param cal_dataset: Calibration set.
         :type cal_dataset: Subset
+        :param train_dataset: Unused by this method. Accepted only so
+            callers that generically calibrate any SetPredictor (e.g.
+            CPBench) can pass the same (cal_dataset, train_dataset,
+            test_dataset) arguments to every CP method without branching on
+            which extra data each one needs.
+        :param test_dataset: Unused by this method (see train_dataset).
         """
         cal_dataset = prepare_numpy_dataset(
             self.model, cal_dataset, ["y_prob", "y_true"], debug=self.debug