diff --git a/docs/api/calib/pyhealth.calib.predictionset.rst b/docs/api/calib/pyhealth.calib.predictionset.rst index 740b1c87c..54e4969f4 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). Fitting K-means on train+cal embeddings combined + (rather than train only, with calibration points assigned via + ``.predict()``) was checked via Monte Carlo simulation and found not + to introduce measurable coverage bias -- 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..9ba261f4e 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,43 @@ 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`` and ``cal_embeddings`` + combined (see ``calibrate()``), so calibration points do influence + the cluster centroids used to assign their own threshold, unlike a + strict split-conformal setup where the category function would be + fit on data disjoint from calibration. This was checked empirically + (Monte Carlo simulation across cluster-count regimes, including + calibration-set-dominated fits) and found not to introduce + measurable coverage bias -- unlike a k-nearest-neighbors-based + category function (see + :class:`~pyhealth.calib.predictionset.NeighborhoodLabel`), where + a similar self-inclusion effect *is* a hard, always-occurring + artifact (a query point is trivially its own nearest neighbor), a + single calibration point's leverage on a K-means centroid -- an + average over many points -- is negligible in practice. This is a + deliberate, verified design choice, not an oversight. + + 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 diff --git a/tests/core/test_cluster_label.py b/tests/core/test_cluster_label.py index ca2ab382d..522d1ebbe 100644 --- a/tests/core/test_cluster_label.py +++ b/tests/core/test_cluster_label.py @@ -514,5 +514,115 @@ 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. + + Also specifically regression-tests the design choice documented in + ClusterLabel's docstring: fitting KMeans on train+cal combined (as + calibrate() does, via .labels_ for calibration points) versus fitting + on train only and using .predict() for calibration points (the + stricter split-conformal-consistent alternative) should not produce a + measurably different coverage outcome. + """ + + def _run_trial(self, rng, n_train, n_cal, n_test, n_kmeans_clusters, + alpha, use_predict_for_cal): + 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) + + if use_predict_for_cal: + km = KMeans(n_clusters=n_kmeans_clusters, random_state=0, n_init=10) + km.fit(train_emb) + cal_cluster = km.predict(cal_emb) + else: + all_emb = np.concatenate([train_emb, cal_emb], axis=0) + km = KMeans(n_clusters=n_kmeans_clusters, random_state=0, n_init=10) + km.fit(all_emb) + cal_cluster = km.labels_[n_train:] + + 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 (KMeans fit on + train+cal, .labels_ for calibration points) 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, + use_predict_for_cal=False) + 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}", + ) + + def test_combined_fit_matches_train_only_fit_coverage(self): + """Regression test for the documented design choice: fitting KMeans + on train+cal combined (current calibrate() behavior) must not + produce measurably worse coverage than fitting on train only and + using .predict() for calibration points -- verified here at a + scale (thousands of trials-worth of test points) a full + model-based Monte Carlo test can't practically reach.""" + rng_combined = np.random.default_rng(123) + rng_train_only = np.random.default_rng(123) + alpha = 0.1 + n_trials = 40 + + combined = [ + self._run_trial(rng_combined, n_train=20, n_cal=300, n_test=2000, + n_kmeans_clusters=3, alpha=alpha, + use_predict_for_cal=False) + for _ in range(n_trials) + ] + train_only = [ + self._run_trial(rng_train_only, n_train=20, n_cal=300, n_test=2000, + n_kmeans_clusters=3, alpha=alpha, + use_predict_for_cal=True) + for _ in range(n_trials) + ] + + # Both should be close to target; neither should be a full + # standard-error below the other -- i.e. combined-fit isn't + # measurably worse than the "stricter" alternative. + self.assertGreaterEqual(np.mean(combined), 1 - alpha - 0.03) + self.assertGreaterEqual(np.mean(train_only), 1 - alpha - 0.03) + self.assertAlmostEqual(np.mean(combined), np.mean(train_only), delta=0.03) + + if __name__ == "__main__": unittest.main()