Support binary classification in conformal prediction set methods

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

@@ -15,6 +15,18 @@ adapts the prediction set size to the model's per-input confidence. See
 ``SCRIB`` and ``FavMac`` are not included since their calibration
 procedures aren't a score-then-quantile pattern.
 
+All of these methods (``SCRIB`` included) also accept **binary** base models in
+addition to multiclass. A binary model emits a single positive-class
+probability, which is expanded internally to a two-class layout
+(:func:`pyhealth.calib.utils.binary_to_2col`) so the prediction set ranges over
+both classes. Score binary results with
+:func:`pyhealth.metrics.binary_metrics_fn`, which accepts ``y_predset`` and
+computes the conformal set metrics (``set_size``, ``rejection_rate``,
+``miscoverage_ps``, ...). See
+``examples/conformal_label_binary.py`` for training, calibration,
+and evaluation with both score types on synthetic MIMIC-III. ``FavMac`` remains
+multilabel-only.
+
 Available Methods
 -----------------
 

examples/conformal_label_binary.py

@@ -0,0 +1,107 @@
+"""Binary conformal prediction (LABEL) with RNN on synthetic MIMIC-III.
+
+This script trains a binary readmission model, calibrates LABEL using threshold
+and APS scores, and evaluates prediction sets on held-out patients.
+
+Run from the repository root:
+    python -m examples.conformal_label_binary
+
+The public synthetic dataset is downloaded automatically. Training, validation,
+calibration, and testing use separate patients. Results illustrate the workflow
+on synthetic data; they are not estimates of clinical performance.
+"""
+
+import tempfile
+
+import torch
+
+from pyhealth.calib.predictionset import LABEL
+from pyhealth.datasets import (
+    MIMIC3Dataset,
+    get_dataloader,
+    split_by_patient_conformal,
+)
+from pyhealth.metrics import binary_metrics_fn
+from pyhealth.models import RNN
+from pyhealth.tasks import ReadmissionPredictionMIMIC3
+from pyhealth.trainer import Trainer
+
+if __name__ == "__main__":
+    torch.manual_seed(42)
+    cache_dir = tempfile.TemporaryDirectory()
+
+    # STEP 1: Load dataset
+    base_dataset = MIMIC3Dataset(
+        root="https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III",
+        tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
+        cache_dir=cache_dir.name,
+        dev=True,
+        num_workers=1,
+    )
+    base_dataset.stats()
+
+    # STEP 2: Set task
+    # Must include minors to get any readmission samples on the synthetic dataset
+    task = ReadmissionPredictionMIMIC3(exclude_minors=False)
+    sample_dataset = base_dataset.set_task(task)
+
+    # STEP 3: Reserve calibration patients separately from model validation.
+    train_dataset, val_dataset, cal_dataset, test_dataset = split_by_patient_conformal(
+        sample_dataset, [0.5, 0.1, 0.2, 0.2], seed=42
+    )
+    print(
+        f"Samples: train={len(train_dataset)}, validation={len(val_dataset)}, "
+        f"calibration={len(cal_dataset)}, test={len(test_dataset)}"
+    )
+    train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)
+    val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)
+    test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)
+
+    # STEP 4: Define model
+    model = RNN(
+        dataset=sample_dataset,
+    )
+
+    # STEP 5: Train
+    # Tiny synthetic splits may contain one class, so monitor accuracy, not AUC.
+    trainer = Trainer(model=model, metrics=["accuracy"], enable_logging=False)
+    trainer.train(
+        train_dataloader=train_dataloader,
+        val_dataloader=val_dataloader,
+        epochs=1,
+        monitor="accuracy",
+    )
+
+    # STEP 6: Evaluate the base model on held-out test patients.
+    print("Base model metrics:", trainer.evaluate(test_dataloader))
+
+    # STEP 7: Calibrate binary prediction sets with either supported score.
+    # A 70% target allows a finite quantile with only four calibration samples.
+    alpha = 0.3
+    for score_type in ("threshold", "aps"):
+        predictor = LABEL(model, alpha=alpha, score_type=score_type, random_state=42)
+        predictor.calibrate(cal_dataset=cal_dataset)
+
+        # Collect prediction sets explicitly; Trainer.evaluate reports scalar
+        # classification metrics but does not collect y_predset.
+        y_true, y_prob, _, extra = Trainer(
+            model=predictor, device=trainer.device, enable_logging=False
+        ).inference(test_dataloader, additional_outputs=["y_predset"])
+        metrics = binary_metrics_fn(
+            y_true,
+            y_prob,
+            metrics=[
+                "accuracy", "set_size", "rejection_rate", "miscoverage_overall_ps"
+            ],
+            y_predset=extra["y_predset"],
+        )
+        print(f"\nLABEL ({score_type}), target coverage: {1 - alpha:.0%}")
+        print("Metrics:", metrics)
+        print(f"Empirical coverage: {1 - metrics['miscoverage_overall_ps']:.3f}")
+        print("Probability shape:", y_prob.shape)  # (N, 1), still binary
+        print("Prediction-set shape:", extra["y_predset"].shape)  # (N, 2)
+        print("First five sets [class 0, class 1]:")
+        print(extra["y_predset"][:5])
+
+    sample_dataset.close()
+    cache_dir.cleanup()

pyhealth/calib/predictionset/base_conformal/__init__.py

@@ -34,7 +34,7 @@
     all_class_nc_scores,
     true_class_nc_scores,
 )
-from pyhealth.calib.utils import prepare_numpy_dataset
+from pyhealth.calib.utils import binary_to_2col, prepare_numpy_dataset
 from pyhealth.models import BaseModel
 
 __all__ = ["BaseConformal"]
@@ -194,9 +194,9 @@ def __init__(
     ) -> None:
         super().__init__(model, **kwargs)
 
-        if model.mode != "multiclass":
+        if model.mode not in ("multiclass", "binary"):
             raise NotImplementedError(
-                "BaseConformal only supports multiclass classification"
+                "BaseConformal only supports multiclass and binary classification"
             )
         if score_type not in SUPPORTED_SCORE_TYPES:
             raise ValueError(
@@ -256,6 +256,9 @@ def calibrate(self, cal_dataset: IterableDataset):
 
         y_prob = cal_dataset_dict["y_prob"]
         y_true = cal_dataset_dict["y_true"]
+        if self.mode == "binary":
+            y_prob = binary_to_2col(y_prob)
+            y_true = np.asarray(y_true).reshape(-1).astype(int)
         N, K = y_prob.shape
 
         # Compute non-conformity scores (higher = less conforming)
@@ -309,6 +312,10 @@ def forward(self, **kwargs) -> dict[str, torch.Tensor]:
         pred = self.model(**kwargs)
 
         y_prob = pred["y_prob"].detach().cpu().numpy()
+        # Binary: expand to 2 columns so the set ranges over both classes
+        # (y_prob itself stays native).
+        if self.mode == "binary":
+            y_prob = binary_to_2col(y_prob)
         nc_scores = all_class_nc_scores(
             y_prob, score_type=self.score_type, rng=self.rng
         )

pyhealth/calib/predictionset/cluster/cluster_label.py

@@ -24,7 +24,11 @@
     all_class_nc_scores,
     true_class_nc_scores,
 )
-from pyhealth.calib.utils import extract_embeddings, prepare_numpy_dataset
+from pyhealth.calib.utils import (
+    binary_to_2col,
+    extract_embeddings,
+    prepare_numpy_dataset,
+)
 from pyhealth.models import BaseModel
 
 __all__ = ["ClusterLabel"]
@@ -123,9 +127,9 @@ def __init__(
     ) -> None:
         super().__init__(model, **kwargs)
 
-        if model.mode != "multiclass":
+        if model.mode not in ("multiclass", "binary"):
             raise NotImplementedError(
-                "ClusterLabel only supports multiclass classification"
+                "ClusterLabel only supports multiclass and binary classification"
             )
         if score_type not in SUPPORTED_SCORE_TYPES:
             raise ValueError(
@@ -202,6 +206,9 @@ def calibrate(
 
         y_prob = cal_dataset_dict["y_prob"]
         y_true = cal_dataset_dict["y_true"]
+        if self.mode == "binary":
+            y_prob = binary_to_2col(y_prob)
+            y_true = np.asarray(y_true).reshape(-1).astype(int)
         N, K = y_prob.shape
 
         # Extract embeddings if not provided
@@ -334,23 +341,28 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
             cluster_thresholds, device=self.device, dtype=pred["y_prob"].dtype
         )
 
-        # Broadcast thresholds to match y_prob shape (batch_size, n_classes).
-        # Marginal: thresholds are (batch_size,) -> view to (batch_size, 1, ...).
-        # Class-conditional: thresholds are already (batch_size, K), no view.
-        if pred["y_prob"].ndim > 1 and cluster_thresholds.ndim == 1:
-            view_shape = (cluster_thresholds.shape[0],) + (1,) * (
-                pred["y_prob"].ndim - 1
-            )
-            cluster_thresholds = cluster_thresholds.view(view_shape)
-
-        # Include class y if its NC score <= NC threshold
+        # Compute NC scores; binary: expand y_prob to 2 columns first so the
+        # set ranges over both classes (y_prob itself stays native).
         y_prob_np = pred["y_prob"].detach().cpu().numpy()
+        if self.mode == "binary":
+            y_prob_np = binary_to_2col(y_prob_np)
         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
         )
+
+        # Broadcast thresholds to match nc_scores shape (batch_size, n_classes).
+        # Marginal: thresholds are (batch_size,) -> view to (batch_size, 1, ...).
+        # Class-conditional: thresholds are already (batch_size, K), no view.
+        if nc_scores.ndim > 1 and cluster_thresholds.ndim == 1:
+            view_shape = (cluster_thresholds.shape[0],) + (1,) * (
+                nc_scores.ndim - 1
+            )
+            cluster_thresholds = cluster_thresholds.view(view_shape)
+
+        # Include class y if its NC score <= NC threshold
         pred["y_predset"] = nc_scores <= cluster_thresholds
         pred.pop("embed", None)  # do not expose internal embedding to caller
         return pred

pyhealth/calib/predictionset/cluster/neighborhood_label.py

@@ -17,7 +17,11 @@
     all_class_conformity_scores,
     true_class_conformity_scores,
 )
-from pyhealth.calib.utils import extract_embeddings, prepare_numpy_dataset
+from pyhealth.calib.utils import (
+    binary_to_2col,
+    extract_embeddings,
+    prepare_numpy_dataset,
+)
 from pyhealth.models import BaseModel
 
 __all__ = ["NeighborhoodLabel"]
@@ -92,9 +96,9 @@ def __init__(
     ) -> None:
         super().__init__(model, **kwargs)
 
-        if model.mode != "multiclass":
+        if model.mode not in ("multiclass", "binary"):
             raise NotImplementedError(
-                "NeighborhoodLabel only supports multiclass classification"
+                "NeighborhoodLabel only supports multiclass and binary classification"
             )
         if score_type not in SUPPORTED_SCORE_TYPES:
             raise ValueError(
@@ -158,6 +162,9 @@ def calibrate(
         )
         y_prob = cal_dict["y_prob"]
         y_true = cal_dict["y_true"]
+        if self.mode == "binary":
+            y_prob = binary_to_2col(y_prob)
+            y_true = np.asarray(y_true).reshape(-1).astype(int)
         N = y_prob.shape[0]
 
         if cal_embeddings is None:
@@ -248,21 +255,27 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
         th = torch.as_tensor(
             thresholds, device=self.device, dtype=pred["y_prob"].dtype
         )
-        if pred["y_prob"].ndim > 1:
-            th = th.view(-1, *([1] * (pred["y_prob"].ndim - 1)))
 
+        # Compute conformity scores; binary: expand y_prob to 2 columns first so
+        # the set ranges over both classes (y_prob itself stays native).
         y_prob_np = pred["y_prob"].detach().cpu().numpy()
+        if self.mode == "binary":
+            y_prob_np = binary_to_2col(y_prob_np)
         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
         )
+        if conformity_scores.ndim > 1:
+            th = th.view(-1, *([1] * (conformity_scores.ndim - 1)))
         y_predset = conformity_scores >= th
-        # if threshold is high, include at least argmax
+        # if threshold is high, include at least the highest-probability class
         empty = y_predset.sum(dim=1) == 0
         if empty.any():
-            argmax_idx = pred["y_prob"].argmax(dim=1)
+            argmax_idx = torch.as_tensor(
+                y_prob_np.argmax(axis=1), device=pred["y_prob"].device
+            )
             y_predset[empty, argmax_idx[empty]] = True
         pred["y_predset"] = y_predset
         pred.pop("embed", None)

pyhealth/calib/predictionset/covariate/covariate_label.py

@@ -35,7 +35,7 @@
     all_class_conformity_scores,
     true_class_conformity_scores,
 )
-from pyhealth.calib.utils import prepare_numpy_dataset
+from pyhealth.calib.utils import binary_to_2col, prepare_numpy_dataset
 from pyhealth.datasets import get_dataloader
 from pyhealth.models import BaseModel
 
@@ -362,9 +362,9 @@ def __init__(
     ) -> None:
         super().__init__(model, **kwargs)
 
-        if model.mode != "multiclass":
+        if model.mode not in ("multiclass", "binary"):
             raise NotImplementedError(
-                "CovariateLabel only supports multiclass classification"
+                "CovariateLabel only supports multiclass and binary classification"
             )
         if score_type not in SUPPORTED_SCORE_TYPES:
             raise ValueError(
@@ -462,6 +462,9 @@ def calibrate(
 
         y_prob = cal_dataset_dict["y_prob"]
         y_true = cal_dataset_dict["y_true"]
+        if self.mode == "binary":
+            y_prob = binary_to_2col(y_prob)
+            y_true = np.asarray(y_true).reshape(-1).astype(int)
         N, K = y_prob.shape
 
         # Determine weights: either custom or KDE-based
@@ -551,8 +554,11 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
         """
         pred = self.model(**kwargs)
 
-        # Construct prediction set by thresholding conformity scores
+        # Construct prediction set by thresholding conformity scores; binary:
+        # expand y_prob to 2 columns first (y_prob itself stays native).
         y_prob = pred["y_prob"].detach().cpu().numpy()
+        if self.mode == "binary":
+            y_prob = binary_to_2col(y_prob)
         conformity_scores = all_class_conformity_scores(
             y_prob, score_type=self.score_type, rng=self.rng
         )