From b1964a30fe6a8673531b3ec29ff5102b7a7e983d Mon Sep 17 00:00:00 2001 From: lehendo Date: Sat, 29 Aug 2026 19:28:14 -0500 Subject: [PATCH 1/2] Audit ClusterLabel: cite Mondrian conformal prediction, document verified K-means design choice, add coverage tests ClusterLabel had no paper citation at all, despite being a legitimate, correctly-implemented instance of Mondrian conformal prediction (Vovk, Lindsay, Nouretdinov, and Gammerman, 'Mondrian confidence machine,' 2003) using K-means clusters as the category function. Added that citation plus Vovk/Gammerman/Shafer 2005, and documented the actual guarantee this class provides: per-cluster coverage (P(Y not in C(X) | cluster=c) <= alpha for every cluster), which is strictly stronger than plain marginal coverage -- the previous docstring didn't call this out. Investigated a theoretical concern by analogy to the NeighborhoodLabel self-inclusion bug fixed earlier this session: calibrate() fits KMeans on train+cal embeddings combined, so calibration points influence the very 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. Verified via Monte Carlo simulation (including a calibration-set-dominated stress test, 10 train vs 300 cal points) that this does NOT introduce measurable coverage bias, unlike k-NN's self- inclusion (a hard, always-occurring artifact) -- a single point's leverage on a K-means centroid, an average over many points, is negligible. Documented this as a deliberate, verified design choice rather than leaving it unexamined, and added regression tests locking in both the per-cluster coverage guarantee and the train+cal-vs-train-only equivalence finding. --- .../calib/pyhealth.calib.predictionset.rst | 15 +++ .../conformal_eeg/tuev_kmeans_conformal.py | 6 + .../predictionset/cluster/cluster_label.py | 60 +++++++++- tests/core/test_cluster_label.py | 110 ++++++++++++++++++ 4 files changed, 187 insertions(+), 4 deletions(-) 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() From fba49ce791e392f6c1abd6b29fa7a618d7fb0e0f Mon Sep 17 00:00:00 2001 From: lehendo Date: Thu, 3 Sep 2026 00:06:46 -0500 Subject: [PATCH 2/2] Fit ClusterLabel's K-means on training embeddings only calibrate() fit K-means on concatenate([train_embeddings, cal_embeddings]) and read calibration points' cluster assignments off kmeans_model.labels_ -- i.e. in-sample. Each calibration point's own presence in the fit could shift the cluster boundary used to assign its own threshold, breaking the Mondrian conformal guarantee's requirement (Vovk, Lindsay, Nouretdinov, and Gammerman 2003) that the category function be independent of the calibration data it's evaluated against. A prior commit on this branch (b1964a3) investigated this exact concern via Monte Carlo simulation and found no measurable coverage bias, given K-means centroids are a smooth average over many points (unlike NeighborhoodLabel's k-NN self-match, which is an exact, always-occurring distance-0 hit). That's a reasonable empirical argument, but the textbook-correct fix -- fit K-means on train_embeddings only, assign calibration/test points via out-of-sample .predict() -- is simpler code (no concatenation bookkeeping) and needs no ongoing empirical justification, so there's no reason to keep the weaker guarantee. Replaced the "combined vs. train-only fit" Monte Carlo comparison test (whose premise no longer applies, since only train-only fitting remains) with a regression test spying on KMeans.fit/.predict to directly verify calibrate() fits on train_embeddings only and assigns calibration points out-of-sample. --- .../calib/pyhealth.calib.predictionset.rst | 10 +- .../predictionset/cluster/cluster_label.py | 60 ++++---- tests/core/test_cluster_label.py | 129 ++++++++++-------- 3 files changed, 108 insertions(+), 91 deletions(-) diff --git a/docs/api/calib/pyhealth.calib.predictionset.rst b/docs/api/calib/pyhealth.calib.predictionset.rst index 54e4969f4..fdff98e6a 100644 --- a/docs/api/calib/pyhealth.calib.predictionset.rst +++ b/docs/api/calib/pyhealth.calib.predictionset.rst @@ -82,11 +82,11 @@ ClusterLabel (K-means Cluster-based Conformal) 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. + 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: diff --git a/pyhealth/calib/predictionset/cluster/cluster_label.py b/pyhealth/calib/predictionset/cluster/cluster_label.py index 9ba261f4e..131ef77c1 100644 --- a/pyhealth/calib/predictionset/cluster/cluster_label.py +++ b/pyhealth/calib/predictionset/cluster/cluster_label.py @@ -72,21 +72,17 @@ class ClusterLabel(SetPredictor): 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. + 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 @@ -224,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 @@ -273,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 522d1ebbe..760f36f1b 100644 --- a/tests/core/test_cluster_label.py +++ b/tests/core/test_cluster_label.py @@ -521,17 +521,9 @@ class TestClusterLabelCoverage(unittest.TestCase): 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): + 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 @@ -550,16 +542,11 @@ def sample(n): 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:] - + # 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 = {} @@ -572,17 +559,16 @@ def sample(n): 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).""" + """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, - use_predict_for_cal=False) + n_kmeans_clusters=3, alpha=alpha) for _ in range(30) ] mean_coverage = np.mean(coverages) @@ -591,37 +577,72 @@ def test_per_cluster_coverage_matches_target(self): 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) +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) - # 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) + 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__":