diff --git a/docs/api/calib/pyhealth.calib.predictionset.rst b/docs/api/calib/pyhealth.calib.predictionset.rst index 740b1c87c..fdff98e6a 100644 --- a/docs/api/calib/pyhealth.calib.predictionset.rst +++ b/docs/api/calib/pyhealth.calib.predictionset.rst @@ -73,6 +73,21 @@ CovariateLabel (Covariate Shift Adaptive) ClusterLabel (K-means Cluster-based Conformal) ---------------------------------------------- +.. note:: + + ClusterLabel is an instance of Mondrian conformal prediction (Vovk, + Lindsay, Nouretdinov, and Gammerman, "Mondrian confidence machine," + Technical report, Royal Holloway University of London, 2003) using + K-means clusters as the category function -- not itself a specific + published method, but a pyhealth-original combination of a standard + technique with the Mondrian framework. Coverage holds independently + *within each cluster* (a strictly stronger guarantee than plain + marginal coverage). K-means is fit on training embeddings only, and + calibration/test points are assigned to clusters out-of-sample via + ``.predict()`` -- the category function must be independent of the + calibration data for the Mondrian guarantee to hold; see the class + docstring's ``Note`` for details. + .. autoclass:: pyhealth.calib.predictionset.ClusterLabel :members: :undoc-members: diff --git a/examples/conformal_eeg/tuev_kmeans_conformal.py b/examples/conformal_eeg/tuev_kmeans_conformal.py index faad50eaa..08bbfd7e6 100644 --- a/examples/conformal_eeg/tuev_kmeans_conformal.py +++ b/examples/conformal_eeg/tuev_kmeans_conformal.py @@ -18,6 +18,12 @@ Notes: - ClusterLabel uses K-means clustering on embeddings to compute cluster-specific thresholds. - Different K values can be tested to find the optimal cluster count. +- This is Mondrian conformal prediction (Vovk, Lindsay, Nouretdinov, and + Gammerman 2003) with K-means clusters as the category function: coverage + holds independently within each cluster, not just marginally across the + whole population. As with the other classes in this module, this assumes + exchangeability between calibration and test embeddings and does not + correct for covariate shift. """ from __future__ import annotations diff --git a/pyhealth/calib/predictionset/cluster/cluster_label.py b/pyhealth/calib/predictionset/cluster/cluster_label.py index 0c719973c..131ef77c1 100644 --- a/pyhealth/calib/predictionset/cluster/cluster_label.py +++ b/pyhealth/calib/predictionset/cluster/cluster_label.py @@ -6,8 +6,25 @@ similar patients into clusters and computes separate calibration thresholds for each cluster, enabling cluster-aware prediction sets. -This serves as a baseline approach for future personalized/dynamic conformal -prediction methods that use patient similarity for calibration set construction. +This is an instance of Mondrian conformal prediction (Vovk, Lindsay, +Nouretdinov, and Gammerman 2003) using K-means-defined clusters as the +category/taxonomy function -- not itself a specific published method, but a +pyhealth-original combination of a standard technique (K-means) with the +general Mondrian conformal prediction framework. It serves as a baseline for +future personalized/dynamic conformal prediction methods that use patient +similarity for calibration set construction. + +Paper: + Vovk, Vladimir, Alexander Gammerman, and Glenn Shafer. + "Algorithmic learning in a random world." Springer, 2005. + + Vovk, Vladimir, David Lindsay, Ilia Nouretdinov, and Alex Gammerman. + "Mondrian confidence machine." Technical report, Royal Holloway + University of London, 2003. (Introduces category-conditional -- + "Mondrian" -- conformal prediction, of which per-cluster calibration + is an instance: each cluster is a Mondrian "category," and the + guarantee below holds independently within each one, not just on + average across the population.) """ from typing import Dict, Optional, Union @@ -39,8 +56,39 @@ class ClusterLabel(SetPredictor): At inference time, test samples are assigned to their nearest cluster and use the cluster-specific threshold. - This approach is simpler than KDE-based methods and serves as a baseline - for more advanced personalized conformal prediction approaches. + This is Mondrian conformal prediction (Vovk, Lindsay, Nouretdinov, and + Gammerman 2003) with K-means clusters as the category function, so the + coverage guarantee holds independently *within each cluster*, not just + marginally: + + - For marginal alpha (float): P(Y not in C(X) | cluster=c) <= alpha, + for every cluster c -- which implies, but is stronger than, the + overall marginal guarantee P(Y not in C(X)) <= alpha. + - For class-conditional alpha (array): P(Y not in C(X) | Y=k, + cluster=c) <= alpha[k], for every class k and cluster c. + + This approach is simpler than KDE-based methods (see + :class:`~pyhealth.calib.predictionset.CovariateLabel`) and serves as a + baseline for more advanced personalized conformal prediction approaches. + + Note: + K-means is fit on ``train_embeddings`` only (see ``calibrate()``); + calibration and test points are then assigned to clusters + out-of-sample via ``.predict()``. This is required by the Mondrian + guarantee above: the category function (cluster membership) must be + independent of the calibration data used to compute each category's + threshold. Fitting on calibration data too (even combined with + training data) would let a calibration point influence the cluster + boundary used to assign its own threshold -- the same kind of + self-inclusion leak that :class:`~pyhealth.calib.predictionset.NeighborhoodLabel` + avoids via leave-one-out for its k-nearest-neighbor category + function. + + As with the other classes in this module, this assumes the + calibration and test embeddings are exchangeable; it does not + correct for covariate shift (see + :class:`~pyhealth.calib.predictionset.CovariateLabel` for a method + that does). Args: model: A trained base model that supports embedding extraction @@ -172,10 +220,10 @@ def calibrate( """Calibrate cluster-specific thresholds. This method: - 1. Combines train and calibration embeddings for clustering - 2. Fits K-means on the combined embeddings - 3. Assigns calibration samples to clusters - 4. Computes cluster-specific calibration thresholds + 1. Fits K-means on the training embeddings only + 2. Assigns calibration samples to clusters out-of-sample via + ``.predict()`` + 3. Computes cluster-specific calibration thresholds Args: cal_dataset: Calibration set @@ -221,26 +269,26 @@ def calibrate( else: train_embeddings = np.asarray(train_embeddings) - # Combine embeddings for clustering - print(f"Combining embeddings: train={train_embeddings.shape}, cal={cal_embeddings.shape}") - all_embeddings = np.concatenate([train_embeddings, cal_embeddings], axis=0) - print(f"Total embeddings for clustering: {all_embeddings.shape}") - - # Fit K-means on combined embeddings - print(f"Fitting K-means with {self.n_clusters} clusters...") + # Fit K-means on training embeddings only. Calibration points must + # not influence the cluster boundaries used to assign their own + # threshold -- the Mondrian conformal guarantee (Vovk, Lindsay, + # Nouretdinov, and Gammerman 2003) requires the category function + # (here, K-means cluster membership) to be independent of the + # calibration data. Calibration (and test) points are then assigned + # out-of-sample via .predict(), exactly mirroring how a real test + # point is assigned at inference time. + print(f"Fitting K-means with {self.n_clusters} clusters on training embeddings...") self.kmeans_model = KMeans( n_clusters=self.n_clusters, random_state=self.random_state, n_init=10, ) - self.kmeans_model.fit(all_embeddings) + self.kmeans_model.fit(train_embeddings) - # Assign calibration samples to clusters - # Note: cal_embeddings start at index len(train_embeddings) in all_embeddings - cal_start_idx = len(train_embeddings) - cal_cluster_labels = self.kmeans_model.labels_[cal_start_idx:] + # Assign calibration samples to clusters out-of-sample + cal_cluster_labels = self.kmeans_model.predict(cal_embeddings) - print(f"Cluster assignments: {np.bincount(cal_cluster_labels)}") + print(f"Cluster assignments: {np.bincount(cal_cluster_labels, minlength=self.n_clusters)}") # Compute non-conformity scores (higher = less conforming) conformity_scores = true_class_nc_scores( diff --git a/tests/core/test_cluster_label.py b/tests/core/test_cluster_label.py index ca2ab382d..760f36f1b 100644 --- a/tests/core/test_cluster_label.py +++ b/tests/core/test_cluster_label.py @@ -514,5 +514,136 @@ def test_model_device_handling(self): self.assertEqual(output["y_predset"].device.type, device.type) +class TestClusterLabelCoverage(unittest.TestCase): + """Monte Carlo verification of ClusterLabel's core statistical claim: + per-cluster (Mondrian) coverage, at scale a full trained-model pipeline + can't practically reach. Exercises the same calibration logic + ClusterLabel.calibrate()/forward() use (KMeans + _query_quantile), + directly, the same way test_scores.py's TestScoresCoverage does for + the shared score module. + """ + + def _run_trial(self, rng, n_train, n_cal, n_test, n_kmeans_clusters, alpha): + from sklearn.cluster import KMeans + from pyhealth.calib.predictionset.base_conformal import _query_quantile + + n_true_clusters = 3 + embed_dim = 5 + centers = rng.normal(scale=8.0, size=(n_true_clusters, embed_dim)) + beta_params = [(2, 8), (5, 5), (8, 2)] + + def sample(n): + true_c = rng.integers(0, n_true_clusters, size=n) + emb = centers[true_c] + rng.normal(scale=1.0, size=(n, embed_dim)) + scores = np.array([rng.beta(*beta_params[c]) for c in true_c]) + return emb, scores + + train_emb, _ = sample(n_train) + cal_emb, cal_scores = sample(n_cal) + test_emb, test_scores = sample(n_test) + + # Mirrors ClusterLabel.calibrate(): fit K-means on training + # embeddings only, assign calibration/test points out-of-sample. + km = KMeans(n_clusters=n_kmeans_clusters, random_state=0, n_init=10) + km.fit(train_emb) + cal_cluster = km.predict(cal_emb) + test_cluster = km.predict(test_emb) + + thresholds = {} + for c in range(n_kmeans_clusters): + mask = cal_cluster == c + thresholds[c] = ( + _query_quantile(cal_scores[mask], alpha) if mask.sum() > 0 else np.inf + ) + t_test = np.array([thresholds[c] for c in test_cluster]) + return (test_scores <= t_test).mean() + + def test_per_cluster_coverage_matches_target(self): + """The core claim: ClusterLabel's calibration logic (K-means fit on + training embeddings only, calibration points assigned via + .predict()) should achieve approximately the target 1-alpha + coverage, matching the standard split-conformal quantile guarantee + applied within each Mondrian category (cluster).""" + rng = np.random.default_rng(42) + alpha = 0.1 + coverages = [ + self._run_trial(rng, n_train=600, n_cal=300, n_test=2000, + n_kmeans_clusters=3, alpha=alpha) + for _ in range(30) + ] + mean_coverage = np.mean(coverages) + self.assertGreaterEqual( + mean_coverage, 1 - alpha - 0.03, + f"Mean coverage {mean_coverage:.4f} too far below target {1 - alpha}", + ) + + +class TestClusterLabelKMeansFitIsOutOfSample(unittest.TestCase): + """Regression test: calibrate() must not let calibration data influence + the K-means cluster boundaries used to assign calibration points' own + thresholds. K-means must be fit on train_embeddings only, and + calibration points assigned via out-of-sample .predict() -- the + Mondrian conformal guarantee (Vovk, Lindsay, Nouretdinov, and Gammerman + 2003) requires the category function to be independent of the + calibration data it's later evaluated against. + """ + + def setUp(self): + np.random.seed(0) + torch.manual_seed(0) + self.samples = [ + { + "patient_id": f"p{i}", + "visit_id": f"v{i}", + "conditions": [f"c{i}"], + "procedures": [float(i % 3)], + "label": i % 3, + } + for i in range(12) + ] + self.dataset = create_sample_dataset( + samples=self.samples, + input_schema={"conditions": "sequence", "procedures": "sequence"}, + output_schema={"label": "multiclass"}, + dataset_name="test_cluster_out_of_sample", + ) + self.model = MLP( + dataset=self.dataset, + feature_keys=["conditions", "procedures"], + label_key="label", + mode="multiclass", + ) + + def test_kmeans_fit_receives_only_train_embeddings(self): + from unittest.mock import patch + from sklearn.cluster import KMeans + + train_ds = self.dataset.subset(list(range(6))) + cal_ds = self.dataset.subset(list(range(6, 12))) + train_embeddings = extract_embeddings(self.model, train_ds, batch_size=32) + cal_embeddings = extract_embeddings(self.model, cal_ds, batch_size=32) + + cluster_predictor = ClusterLabel(model=self.model, alpha=0.2, n_clusters=2) + + with patch.object(KMeans, "fit", autospec=True) as mock_fit, \ + patch.object(KMeans, "predict", autospec=True, return_value=np.zeros(6, dtype=int)) as mock_predict: + mock_fit.side_effect = lambda self, X, *a, **kw: setattr( + self, "cluster_centers_", np.zeros((2, X.shape[1])) + ) + cluster_predictor.calibrate( + cal_dataset=cal_ds, + train_embeddings=train_embeddings, + cal_embeddings=cal_embeddings, + ) + + fit_X = mock_fit.call_args.args[1] + self.assertEqual(fit_X.shape[0], len(train_embeddings)) + np.testing.assert_array_equal(fit_X, train_embeddings) + + predict_X = mock_predict.call_args.args[1] + self.assertEqual(predict_X.shape[0], len(cal_embeddings)) + np.testing.assert_array_equal(predict_X, cal_embeddings) + + if __name__ == "__main__": unittest.main()