diff --git a/docs/api/calib/pyhealth.calib.predictionset.rst b/docs/api/calib/pyhealth.calib.predictionset.rst index 4b29a22fe..18705f58f 100644 --- a/docs/api/calib/pyhealth.calib.predictionset.rst +++ b/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 ----------------- diff --git a/examples/conformal_label_binary.py b/examples/conformal_label_binary.py new file mode 100644 index 000000000..921cd93f9 --- /dev/null +++ b/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() diff --git a/pyhealth/calib/predictionset/base_conformal/__init__.py b/pyhealth/calib/predictionset/base_conformal/__init__.py index 9dde35db5..fbd0fab7f 100644 --- a/pyhealth/calib/predictionset/base_conformal/__init__.py +++ b/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 ) diff --git a/pyhealth/calib/predictionset/cluster/cluster_label.py b/pyhealth/calib/predictionset/cluster/cluster_label.py index 0c719973c..36218a61b 100644 --- a/pyhealth/calib/predictionset/cluster/cluster_label.py +++ b/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 diff --git a/pyhealth/calib/predictionset/cluster/neighborhood_label.py b/pyhealth/calib/predictionset/cluster/neighborhood_label.py index 36fdbaa0a..8baf97c3d 100644 --- a/pyhealth/calib/predictionset/cluster/neighborhood_label.py +++ b/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) diff --git a/pyhealth/calib/predictionset/covariate/covariate_label.py b/pyhealth/calib/predictionset/covariate/covariate_label.py index 3482e4e91..e19789a1e 100644 --- a/pyhealth/calib/predictionset/covariate/covariate_label.py +++ b/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 ) diff --git a/pyhealth/calib/predictionset/label.py b/pyhealth/calib/predictionset/label.py index b77934a28..10eb2e118 100644 --- a/pyhealth/calib/predictionset/label.py +++ b/pyhealth/calib/predictionset/label.py @@ -20,7 +20,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__ = ["LABEL"] @@ -98,14 +98,14 @@ def __init__( **kwargs, ) -> None: super().__init__(model, **kwargs) - if model.mode != "multiclass": + if model.mode not in ("multiclass", "binary"): 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 + self.mode = self.model.mode # multiclass or binary for param in model.parameters(): param.requires_grad = False self.model.eval() @@ -131,8 +131,11 @@ def calibrate(self, cal_dataset: Subset): ) y_prob = cal_dataset["y_prob"] y_true = cal_dataset["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 = cal_dataset["y_prob"].shape + K = y_prob.shape[1] # NC scores: higher = less conforming nc_scores = true_class_nc_scores( y_prob, y_true, score_type=self.score_type, rng=self.rng @@ -156,6 +159,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 ) diff --git a/pyhealth/calib/predictionset/scrib/__init__.py b/pyhealth/calib/predictionset/scrib/__init__.py index c0e703bb2..9bd191cbf 100644 --- a/pyhealth/calib/predictionset/scrib/__init__.py +++ b/pyhealth/calib/predictionset/scrib/__init__.py @@ -13,7 +13,7 @@ import torch from pyhealth.calib.base_classes import SetPredictor -from pyhealth.calib.utils import prepare_numpy_dataset +from pyhealth.calib.utils import binary_to_2col, prepare_numpy_dataset from pyhealth.models import BaseModel from . import quicksearch as qs @@ -251,9 +251,9 @@ def __init__( **kwargs, ) -> None: super().__init__(model, **kwargs) - if model.mode != "multiclass": + if model.mode not in ("multiclass", "binary"): raise NotImplementedError() - self.mode = self.model.mode # multiclass + self.mode = self.model.mode # multiclass or binary for param in model.parameters(): param.requires_grad = False self.model.eval() @@ -285,11 +285,18 @@ def calibrate(self, cal_dataset): cal_dataset = prepare_numpy_dataset( self.model, cal_dataset, ["y_prob", "y_true"], debug=self.debug ) + # Binary: re-present as a 2-class problem so the coordinate-descent + # threshold search runs over one threshold per class, as for multiclass. + y_prob = cal_dataset["y_prob"] + y_true = cal_dataset["y_true"] + if self.mode == "binary": + y_prob = binary_to_2col(y_prob) + y_true = np.asarray(y_true).reshape(-1).astype(int) if self.loss_name == CLASSPECIFIC_LOSSFUNC: - assert len(self.risk) == cal_dataset["y_prob"].shape[1] + assert len(self.risk) == y_prob.shape[1] best_ts, _ = _CoordDescent.search( - cal_dataset["y_prob"], - cal_dataset["y_true"], + y_prob, + y_true, self.risk, self.loss_name, loss_kwargs=self.loss_kwargs, @@ -306,14 +313,22 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: :rtype: Dict[str, torch.Tensor] """ ret = self.model(**kwargs) - y_predset = ret["y_prob"] > self.t + # Binary: build a 2-column probability [P(0), P(1)] for the set so it + # ranges over both classes; y_prob itself stays native. Use the same + # expansion precision as calibration when comparing with thresholds. + prob = ret["y_prob"] + if self.mode == "binary": + prob = torch.as_tensor( + binary_to_2col(prob.detach().cpu().numpy()), device=prob.device + ) + y_predset = prob > self.t if self.fill_max: # Match the calibration-time assumption: when no class clears # its threshold, fall back to the max-predicted class instead # of returning an empty set. empty = y_predset.sum(dim=1) == 0 if empty.any(): - argmax_idx = ret["y_prob"].argmax(dim=1) + argmax_idx = prob.argmax(dim=1) y_predset[empty, argmax_idx[empty]] = True ret["y_predset"] = y_predset return ret diff --git a/pyhealth/calib/utils.py b/pyhealth/calib/utils.py index fd7ffffc0..5f6a0b261 100644 --- a/pyhealth/calib/utils.py +++ b/pyhealth/calib/utils.py @@ -23,6 +23,23 @@ def one_hot_np(labels, K): return new_labels +def binary_to_2col(y_prob): + """Expand binary ``P(y=1)`` to two columns ``[P(y=0), P(y=1)]``. + + A prediction set ranges over both labels ``{0, 1}``, so it needs one + probability column per class. Turns shape ``(N,)`` or ``(N, 1)`` into + ``(N, 2)`` (numpy). + + Examples: + >>> from pyhealth.calib.utils import binary_to_2col + >>> binary_to_2col([0.2, 0.9]) + array([[0.8, 0.2], + [0.1, 0.9]]) + """ + p = np.asarray(y_prob, dtype=float).reshape(-1, 1) + return np.hstack([1.0 - p, p]) + + class LogLoss(torch.nn.Module): """Cross entropy, but takes in the probability instead of the logits""" diff --git a/pyhealth/metrics/binary.py b/pyhealth/metrics/binary.py index ea7d125f7..63bf9151e 100644 --- a/pyhealth/metrics/binary.py +++ b/pyhealth/metrics/binary.py @@ -4,6 +4,7 @@ import sklearn.metrics as sklearn_metrics import pyhealth.metrics.calibration as calib +import pyhealth.metrics.prediction_set as pset def binary_metrics_fn( @@ -11,6 +12,7 @@ def binary_metrics_fn( y_prob: np.ndarray, metrics: Optional[List[str]] = None, threshold: float = 0.5, + y_predset: np.ndarray | None = None, ) -> Dict[str, float]: """Computes metrics for binary classification. @@ -28,6 +30,17 @@ def binary_metrics_fn( - jaccard: Jaccard similarity coefficient score - ECE: Expected Calibration Error (with 20 equal-width bins). Check :func:`pyhealth.metrics.calibration.ece_confidence_binary`. - ECE_adapt: adaptive ECE (with 20 equal-size bins). Check :func:`pyhealth.metrics.calibration.ece_confidence_binary`. + + Conformal-prediction set metrics (require ``y_predset`` of shape + ``(n_samples, 2)``, one column per class): + - rejection_rate: Frequency of prediction sets with cardinality != 1. Check :func:`pyhealth.metrics.prediction_set.rejection_rate`. + - set_size: Average size of the prediction sets. Check :func:`pyhealth.metrics.prediction_set.size`. + - miscoverage_ps: Prob(k not in prediction set). Check :func:`pyhealth.metrics.prediction_set.miscoverage_ps`. + - miscoverage_mean_ps: The average (across classes) of miscoverage_ps. + - miscoverage_overall_ps: Prob(Y not in prediction set). Check :func:`pyhealth.metrics.prediction_set.miscoverage_overall_ps`. + - error_ps: Same as miscoverage_ps, but restricted to un-rejected samples. Check :func:`pyhealth.metrics.prediction_set.error_ps`. + - error_mean_ps: The average (across classes) of error_ps. + - error_overall_ps: Same as miscoverage_overall_ps, but restricted to un-rejected samples. Check :func:`pyhealth.metrics.prediction_set.error_overall_ps`. If no metrics are specified, pr_auc, roc_auc and f1 are computed by default. This function calls sklearn.metrics functions to compute the metrics. For @@ -58,6 +71,11 @@ def binary_metrics_fn( y_pred[y_pred >= threshold] = 1 y_pred[y_pred < threshold] = 0 + if y_predset is not None: + # Set metrics index class columns; binary model labels can be float + # tensors with shape (N, 1). Keep scalar metrics' inputs unchanged. + y_true_set = np.asarray(y_true).reshape(-1).astype(int) + output = {} for metric in metrics: if metric == "pr_auc": @@ -91,6 +109,12 @@ def binary_metrics_fn( output[metric] = calib.ece_confidence_binary( y_prob, y_true, bins=20, adaptive=metric.endswith("_adapt") ) + elif metric in pset.PREDICTION_SET_METRICS: + if y_predset is None: + continue + output[metric] = pset.compute_prediction_set_metric( + metric, y_predset, y_true_set + ) else: raise ValueError(f"Unknown metric for binary classification: {metric}") return output diff --git a/pyhealth/metrics/multiclass.py b/pyhealth/metrics/multiclass.py index 5b6db3a01..3f3243271 100644 --- a/pyhealth/metrics/multiclass.py +++ b/pyhealth/metrics/multiclass.py @@ -81,16 +81,6 @@ def multiclass_metrics_fn( """ if metrics is None: metrics = ["accuracy", "f1_macro", "f1_micro"] - prediction_set_metrics = [ - "rejection_rate", - "set_size", - "miscoverage_mean_ps", - "miscoverage_ps", - "miscoverage_overall_ps", - "error_mean_ps", - "error_ps", - "error_overall_ps", - ] y_pred = np.argmax(y_prob, axis=-1) output = {} @@ -167,26 +157,13 @@ def multiclass_metrics_fn( adaptive=metric.endswith("_adapt"), threshold=thres, ) - elif metric in prediction_set_metrics: + elif metric in pset.PREDICTION_SET_METRICS: if y_predset is None: continue - if metric == "rejection_rate": - output[metric] = pset.rejection_rate(y_predset) - elif metric == "set_size": - output[metric] = pset.size(y_predset) - elif metric == "miscoverage_mean_ps": - output[metric] = pset.miscoverage_ps(y_predset, y_true).mean() - elif metric == "miscoverage_ps": - output[metric] = pset.miscoverage_ps(y_predset, y_true) - elif metric == "miscoverage_overall_ps": - output[metric] = pset.miscoverage_overall_ps(y_predset, y_true) - elif metric == "error_mean_ps": - output[metric] = pset.error_ps(y_predset, y_true).mean() - elif metric == "error_ps": - output[metric] = pset.error_ps(y_predset, y_true) - elif metric == "error_overall_ps": - output[metric] = pset.error_overall_ps(y_predset, y_true) - + output[metric] = pset.compute_prediction_set_metric( + metric, y_predset, y_true + ) + elif metric == "hits@n": argsort = np.argsort(-y_prob, axis=1) ranking = np.array([np.where(argsort[i] == y_true[i])[0][0] for i in range(len(y_true))]) + 1 diff --git a/pyhealth/metrics/prediction_set.py b/pyhealth/metrics/prediction_set.py index 2b6f71705..c652c1e85 100644 --- a/pyhealth/metrics/prediction_set.py +++ b/pyhealth/metrics/prediction_set.py @@ -1,5 +1,52 @@ import numpy as np +#: Names of the conformal prediction-set metrics. Any classification +#: metrics_fn that receives a ``y_predset`` dispatches these through +#: :func:`compute_prediction_set_metric`. +PREDICTION_SET_METRICS = [ + "rejection_rate", + "set_size", + "miscoverage_mean_ps", + "miscoverage_ps", + "miscoverage_overall_ps", + "error_mean_ps", + "error_ps", + "error_overall_ps", +] + + +def compute_prediction_set_metric(metric: str, y_predset: np.ndarray, y_true: np.ndarray): + """Compute a single prediction-set metric by name. + + Shared by the binary and multiclass metrics functions so the dispatch + lives in one place. ``metric`` must be one of :data:`PREDICTION_SET_METRICS`. + + Examples: + >>> import numpy as np + >>> from pyhealth.metrics.prediction_set import compute_prediction_set_metric + >>> y_predset = np.asarray([[1, 0], [1, 1], [0, 1]]) + >>> y_true = np.asarray([0, 1, 1]) + >>> float(compute_prediction_set_metric("set_size", y_predset, y_true)) + 1.3333333333333333 + """ + if metric == "rejection_rate": + return rejection_rate(y_predset) + if metric == "set_size": + return size(y_predset) + if metric == "miscoverage_mean_ps": + return miscoverage_ps(y_predset, y_true).mean() + if metric == "miscoverage_ps": + return miscoverage_ps(y_predset, y_true) + if metric == "miscoverage_overall_ps": + return miscoverage_overall_ps(y_predset, y_true) + if metric == "error_mean_ps": + return error_ps(y_predset, y_true).mean() + if metric == "error_ps": + return error_ps(y_predset, y_true) + if metric == "error_overall_ps": + return error_overall_ps(y_predset, y_true) + raise ValueError(f"Unknown prediction-set metric: {metric}") + def size(y_pred:np.ndarray): """Average size of the prediction set. diff --git a/tests/core/test_binary_predictionset.py b/tests/core/test_binary_predictionset.py new file mode 100644 index 000000000..61b5d4a79 --- /dev/null +++ b/tests/core/test_binary_predictionset.py @@ -0,0 +1,145 @@ +"""Binary-mode conformal prediction across the SetPredictor subclasses. + +Each method re-presents a binary base model as a 2-class problem, so its +``y_predset`` must come out ``(N, 2)`` boolean and be scorable with the binary +prediction-set metrics. FavMac is intentionally excluded: it is multilabel-only. +""" + +import unittest + +import numpy as np +import torch + +from pyhealth.calib.predictionset import ( + LABEL, + SCRIB, + BaseConformal, + ClusterLabel, + CovariateLabel, + NeighborhoodLabel, +) +from pyhealth.calib.utils import binary_to_2col, extract_embeddings +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.metrics import binary_metrics_fn +from pyhealth.models import MLP + + +class TestBinaryToCol(unittest.TestCase): + """The (N,) / (N,1) -> (N,2) probability expansion.""" + + def test_shapes_and_values(self): + out = binary_to_2col([0.2, 0.9, 0.5]) + self.assertEqual(out.shape, (3, 2)) + np.testing.assert_allclose(out, [[0.8, 0.2], [0.1, 0.9], [0.5, 0.5]]) + # (N, 1) input is accepted and gives the same result. + np.testing.assert_allclose( + binary_to_2col(np.array([[0.2], [0.9], [0.5]])), out + ) + + +class TestBinaryPredictionSet(unittest.TestCase): + def setUp(self): + np.random.seed(0) + torch.manual_seed(0) + + self.samples = [ + { + "patient_id": f"patient-{i}", + "visit_id": f"visit-{i}", + "conditions": [f"cond-{i % 5}", f"cond-{(i + 1) % 5}"], + "procedures": [1.0 * i, 2.0, 3.5, 4.0], + "label": i % 2, + } + for i in range(12) + ] + self.dataset = create_sample_dataset( + samples=self.samples, + input_schema={"conditions": "sequence", "procedures": "tensor"}, + output_schema={"label": "binary"}, + dataset_name="test-binary", + ) + self.model = MLP( + dataset=self.dataset, + feature_keys=["conditions", "procedures"], + label_key="label", + mode="binary", + ) + self.model.eval() + + self.train_ds = self.dataset.subset([0, 1, 2, 3, 4, 5]) + self.cal_ds = self.dataset.subset([6, 7, 8, 9, 10, 11]) + + def _embeddings(self, ds): + return extract_embeddings(self.model, ds, batch_size=32, device="cpu") + + def _assert_binary_set(self, out): + """Every method must yield an (N, 2) bool set, native (N, 1) y_prob, + native binary labels, and metrics that compute through binary_metrics_fn.""" + self.assertEqual(out["y_predset"].dtype, torch.bool) + self.assertEqual(out["y_predset"].dim(), 2) + self.assertEqual(out["y_predset"].shape[1], 2) + self.assertEqual(out["y_prob"].shape[1], 1) + self.assertEqual(out["y_true"].shape, out["y_prob"].shape) + self.assertEqual(out["y_true"].dtype, torch.float32) + + res = binary_metrics_fn( + out["y_true"].numpy(), + out["y_prob"].numpy().reshape(-1), + metrics=["set_size", "rejection_rate", "miscoverage_ps"], + y_predset=out["y_predset"].numpy(), + ) + self.assertIn("rejection_rate", res) + self.assertLessEqual(res["set_size"], 2.0) + self.assertEqual(len(res["miscoverage_ps"]), 2) + + def _forward(self, cal_model): + loader = get_dataloader(self.dataset, batch_size=12, shuffle=False) + with torch.no_grad(): + return cal_model(**next(iter(loader))) + + def test_label(self): + m = LABEL(self.model, alpha=0.3) + self.assertEqual(m.mode, "binary") + m.calibrate(cal_dataset=self.cal_ds) + self.assertIsInstance(m.t, torch.Tensor) + self._assert_binary_set(self._forward(m)) + + def test_base_conformal(self): + m = BaseConformal(self.model, alpha=0.3, score_type="threshold") + m.calibrate(cal_dataset=self.cal_ds) + self._assert_binary_set(self._forward(m)) + + def test_scrib(self): + m = SCRIB(self.model, risk=0.3) + m.calibrate(cal_dataset=self.cal_ds) + self._assert_binary_set(self._forward(m)) + + def test_covariate_label(self): + m = CovariateLabel(self.model, alpha=0.3) + # Custom-weights path (uniform weights) avoids needing a shifted set. + m.calibrate( + cal_dataset=self.cal_ds, + cal_weights=np.ones(len(self.cal_ds)), + ) + self._assert_binary_set(self._forward(m)) + + def test_cluster_label(self): + m = ClusterLabel(self.model, alpha=0.3, n_clusters=2, random_state=0) + m.calibrate( + cal_dataset=self.cal_ds, + train_embeddings=self._embeddings(self.train_ds), + cal_embeddings=self._embeddings(self.cal_ds), + ) + self._assert_binary_set(self._forward(m)) + + def test_neighborhood_label(self): + m = NeighborhoodLabel(self.model, alpha=0.3, k_neighbors=3, lambda_L=50.0) + m.calibrate( + cal_dataset=self.cal_ds, + cal_embeddings=self._embeddings(self.cal_ds), + ) + self._assert_binary_set(self._forward(m)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_cluster_label.py b/tests/core/test_cluster_label.py index ca2ab382d..64b66af56 100644 --- a/tests/core/test_cluster_label.py +++ b/tests/core/test_cluster_label.py @@ -166,9 +166,8 @@ def test_initialization_with_array_alpha(self): self.assertIsInstance(cluster_model.alpha, np.ndarray) np.testing.assert_array_equal(cluster_model.alpha, alpha_per_class) - def test_initialization_non_multiclass_raises_error(self): - """Test that non-multiclass models raise an error.""" - # Create a binary classification dataset + def test_binary_mode_supported(self): + """Binary base models are supported (re-presented as 2-class).""" binary_samples = [ { "patient_id": "patient-0", @@ -198,12 +197,13 @@ def test_initialization_non_multiclass_raises_error(self): mode="binary", ) + cluster_model = ClusterLabel(model=binary_model, alpha=0.1, n_clusters=2) + self.assertEqual(cluster_model.mode, "binary") + + # A genuinely unsupported mode still raises. + binary_model.mode = "regression" with self.assertRaises(NotImplementedError): - ClusterLabel( - model=binary_model, - alpha=0.1, - n_clusters=2, - ) + ClusterLabel(model=binary_model, alpha=0.1, n_clusters=2) def test_initialization_invalid_n_clusters_raises_error(self): """Test that invalid n_clusters (non-positive or non-int) raises ValueError.""" diff --git a/tests/core/test_covariate_label.py b/tests/core/test_covariate_label.py index 6f2bfd04a..bc2bff700 100644 --- a/tests/core/test_covariate_label.py +++ b/tests/core/test_covariate_label.py @@ -131,9 +131,8 @@ def test_initialization_with_array_alpha(self): self.assertIsInstance(cal_model.alpha, np.ndarray) np.testing.assert_array_equal(cal_model.alpha, alpha_per_class) - def test_initialization_non_multiclass_raises_error(self): - """Test that non-multiclass models raise an error.""" - # Create a binary classification dataset with both labels + def test_binary_mode_supported(self): + """Binary base models are supported (re-presented as 2-class).""" binary_samples = [ { "patient_id": "patient-0", @@ -163,6 +162,16 @@ def test_initialization_non_multiclass_raises_error(self): mode="binary", ) + cal_model = CovariateLabel( + model=binary_model, + alpha=0.1, + kde_test=self.kde_test, + kde_cal=self.kde_cal, + ) + self.assertEqual(cal_model.mode, "binary") + + # A genuinely unsupported mode still raises. + binary_model.mode = "regression" with self.assertRaises(NotImplementedError): CovariateLabel( model=binary_model, diff --git a/tests/core/test_neighborhood_label.py b/tests/core/test_neighborhood_label.py index 0c812d141..e1d6c39d7 100644 --- a/tests/core/test_neighborhood_label.py +++ b/tests/core/test_neighborhood_label.py @@ -73,7 +73,8 @@ def test_initialization_invalid_k_neighbors_raises(self): with self.assertRaises(ValueError): NeighborhoodLabel(model=self.model, alpha=0.1, k_neighbors=2.5) - def test_initialization_non_multiclass_raises(self): + def test_binary_mode_supported(self): + """Binary base models are supported (re-presented as 2-class).""" binary_samples = [ {"patient_id": "a", "visit_id": "a", "conditions": ["c"], "procedures": [1.0], "label": 0}, {"patient_id": "b", "visit_id": "b", "conditions": ["d"], "procedures": [2.0], "label": 1}, @@ -87,6 +88,11 @@ def test_initialization_non_multiclass_raises(self): binary_model = MLP( dataset=binary_ds, feature_keys=["conditions"], label_key="label", mode="binary" ) + ncp = NeighborhoodLabel(model=binary_model, alpha=0.1, k_neighbors=2) + self.assertEqual(ncp.mode, "binary") + + # A genuinely unsupported mode still raises. + binary_model.mode = "regression" with self.assertRaises(NotImplementedError): NeighborhoodLabel(model=binary_model, alpha=0.1, k_neighbors=2)