ClusterLabel audit + tests

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:

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

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(

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()