From 2bc7f4e9179e50db7c09ef365e0872452561c6f1 Mon Sep 17 00:00:00 2001 From: lehendo Date: Tue, 25 Aug 2026 20:01:28 -0500 Subject: [PATCH 1/3] Fix under-coverage in CovariateLabel's finite-sample correction _query_weighted_quantile added the test point's reserved weight only to the normalizing denominator, never inserting it as an actual point in the weighted empirical distribution. Verified against Corollary 1 of Tibshirani, Barber, Candes, and Ramdas, "Conformal Prediction Under Covariate Shift" (NeurIPS 2019, arXiv:1904.06019): before this fix, the "corrected" quantile under-covered relative to target (0.81-0.89 vs a 0.90 target in Monte Carlo simulation) -- worse than no correction at all. Fix: prepend the reserved weight to the cumulative sum before dividing, matching the paper's construction exactly (confirmed via a hand-computed example and against a from-scratch reference implementation of the paper's formula). Also fix a related gap: calibrate() computed a single threshold using the *mean* calibration likelihood ratio as a stand-in for a test point's own w(x), but Corollary 1 defines the threshold per test point using that point's actual weight. forward() now accepts an optional test_embeddings argument to compute the real per-point threshold matching the paper exactly; omitting it keeps the old single-threshold behavior as a documented, explicitly-warned approximation, since not every wrapped model exposes an embedding extraction path. 18 tests pass (4 new), including a deterministic regression test against a hand-computed Corollary 1 example and new coverage of the forward() per-point/fallback paths. --- docs/api/calib.rst | 19 ++ .../tuev_covariate_shift_conformal.py | 3 + .../covariate/covariate_label.py | 149 +++++++++++++++- tests/core/test_covariate_label.py | 163 ++++++++++++++++++ 4 files changed, 328 insertions(+), 6 deletions(-) diff --git a/docs/api/calib.rst b/docs/api/calib.rst index 7814b4719..c2543af68 100644 --- a/docs/api/calib.rst +++ b/docs/api/calib.rst @@ -42,6 +42,25 @@ New to calibration and uncertainty quantification? Check out this complete examp This example shows the complete pipeline from model training to uncertainty-aware predictions with formal coverage guarantees. +.. note:: + + ``CovariateLabel``'s finite-sample correction implements Corollary 1 of + Tibshirani, Barber, Candes, and Ramdas, "Conformal Prediction Under + Covariate Shift" (NeurIPS 2019, https://arxiv.org/abs/1904.06019): the + test point's reserved probability mass must be inserted as an actual + point in the weighted empirical distribution (at the conservative + extreme), not merely folded into the normalizing denominator -- doing + only the latter silently under-covers relative to the target coverage + level. + + Corollary 1 also defines the threshold *per test point*, using that + point's own likelihood ratio w(x). Pass ``test_embeddings`` to + ``CovariateLabel.forward()`` to get this exact per-point threshold; + omitting it falls back to a single threshold computed from the *mean* + calibration likelihood ratio, which is only an approximation of the + paper's guarantee (a ``UserWarning`` is raised when this fallback is + used). + Quick Links ----------- diff --git a/examples/conformal_eeg/tuev_covariate_shift_conformal.py b/examples/conformal_eeg/tuev_covariate_shift_conformal.py index d1bcba20a..58efd0c2e 100644 --- a/examples/conformal_eeg/tuev_covariate_shift_conformal.py +++ b/examples/conformal_eeg/tuev_covariate_shift_conformal.py @@ -19,6 +19,9 @@ Notes: - CovariateLabel requires access to test embeddings to estimate density ratios. - Test embeddings are recomputed each seed since the model changes. +- CovariateLabel's finite-sample correction implements Corollary 1 of + Tibshirani, Barber, Candes, and Ramdas (NeurIPS 2019, arXiv:1904.06019); + see docs/api/calib.rst for details. """ from __future__ import annotations diff --git a/pyhealth/calib/predictionset/covariate/covariate_label.py b/pyhealth/calib/predictionset/covariate/covariate_label.py index 3482e4e91..1ff0b6b0e 100644 --- a/pyhealth/calib/predictionset/covariate/covariate_label.py +++ b/pyhealth/calib/predictionset/covariate/covariate_label.py @@ -22,6 +22,7 @@ https://arxiv.org/abs/2310.12033 """ +import warnings from typing import Callable, Dict, Optional, Union import numpy as np @@ -200,6 +201,23 @@ def _query_weighted_quantile( reserved test-point mass alone already meets or exceeds ``alpha``, since there isn't enough calibration mass to justify a stricter, finite threshold without risking under-coverage. + + Note: + This implements Corollary 1 of Tibshirani, Barber, Candes, and + Ramdas, "Conformal Prediction Under Covariate Shift" (NeurIPS 2019, + https://arxiv.org/abs/1904.06019): the (1-alpha)-quantile of + sum_i p_i^w(x) * delta_{V_i} + p_{n+1}^w(x) * delta_infinity, where + p_i^w(x) = w(X_i) / (sum_j w(X_j) + w(x)) and p_{n+1}^w(x) is the + reserved test-point mass. The paper's V is a nonconformity score + (higher = worse) with the reserved mass placed at +infinity; this + codebase uses the opposite convention (conformity, higher = better), + so the reserved mass is placed at -infinity here instead. Crucially, + that reserved mass must be inserted as an actual point in the + weighted empirical distribution, not merely folded into the + normalizing denominator -- otherwise it dilutes every real + calibration weight without ever contributing to the cumulative sum + used to pick the threshold, which produces under-coverage instead of + the intended finite-sample guarantee. """ sorted_indices = np.argsort(scores) sorted_scores = scores[sorted_indices] @@ -213,11 +231,17 @@ def _query_weighted_quantile( if p_test >= alpha: # Not enough calibration mass to reach the target coverage without # dipping into the mass reserved for the test point itself: fall - # back to the maximally permissive (safe) threshold. + # back to the maximally permissive (safe) threshold. This is the + # case where the reserved point (sitting first, at -inf, in the + # augmented distribution) already accounts for the full quantile + # by itself. return -np.inf - # Compute cumulative weights over the reserved-mass-inclusive total. - cum_weights = np.cumsum(sorted_weights) / total_weight + # Cumulative weights over the reserved-mass-inclusive total, with the + # test point's reserved mass prepended as an actual point at -inf + # (equivalent to inserting it at the head of the sorted array before + # taking the cumulative sum). + cum_weights = (test_weight + np.cumsum(sorted_weights)) / total_weight # Find the index where cumulative weight exceeds alpha idx = np.searchsorted(cum_weights, alpha, side="left") @@ -400,6 +424,15 @@ def __init__( # Will be set during calibration self.t = None self._sum_cal_weights = None + # Calibration conformity scores/weights, kept so forward() can + # recompute a threshold per-test-point using that point's own + # likelihood ratio w(x), as Corollary 1 of Tibshirani et al. (2019) + # requires, when test_embeddings are provided. + self._cal_conformity_scores = None + self._cal_likelihood_ratios = None + self._cal_class_scores = None + self._cal_class_weights = None + self._warned_fixed_threshold = False def calibrate( self, @@ -516,7 +549,26 @@ def calibrate( y_prob, y_true, score_type=self.score_type, rng=self.rng ) - # Compute weighted quantile thresholds + # Keep the raw calibration scores/weights so forward() can, when + # given test_embeddings, recompute a threshold per-test-point using + # that point's own likelihood ratio w(x) -- the construction + # Corollary 1 actually specifies -- rather than only the + # mean-weight approximation computed below. + self._cal_conformity_scores = conformity_scores + self._cal_likelihood_ratios = likelihood_ratios + if not isinstance(self.alpha, float): + self._cal_class_scores = [ + conformity_scores[y_true == k] for k in range(K) + ] + self._cal_class_weights = [ + likelihood_ratios[y_true == k] for k in range(K) + ] + + # Compute weighted quantile thresholds using the mean calibration + # weight as a stand-in for a "typical" test point's weight. This is + # the fallback used when forward() isn't given test_embeddings (see + # forward()'s docstring for why that's only an approximation of + # Corollary 1, not an exact instance of it). if isinstance(self.alpha, float): test_weight = float(np.mean(likelihood_ratios)) t = _query_weighted_quantile( @@ -541,9 +593,32 @@ def calibrate( self.t = torch.tensor(t, device=self.device) - def forward(self, **kwargs) -> Dict[str, torch.Tensor]: + def forward( + self, test_embeddings: np.ndarray | None = None, **kwargs + ) -> dict[str, torch.Tensor]: """Forward propagation with prediction set construction. + Args: + test_embeddings: Optional embeddings for this batch's test + points, shape ``(batch_size, embedding_dim)``, aligned with + the batch order. When provided (and KDEs are available from + :meth:`calibrate`), the threshold is recomputed per test + point using that point's own likelihood ratio w(x) via + :func:`_query_weighted_quantile` -- exactly the construction + in Corollary 1 of Tibshirani, Barber, Candes, and Ramdas, + "Conformal Prediction Under Covariate Shift" (NeurIPS 2019, + https://arxiv.org/abs/1904.06019). + + When omitted, falls back to the single threshold computed + during :meth:`calibrate` using the *mean* calibration + likelihood ratio as a stand-in for w(x). That fallback is + only an approximation of Corollary 1 -- the paper's formula + is defined per test point via that point's actual w(x), not + an aggregate over the calibration set -- so results from the + fallback do not carry the same finite-sample coverage + guarantee as the paper's construction. A warning is emitted + (once) the first time this fallback is used. + Returns: Dictionary with all results from base model, plus: - y_predset: Boolean tensor indicating which classes @@ -556,10 +631,72 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: conformity_scores = all_class_conformity_scores( y_prob, score_type=self.score_type, rng=self.rng ) + N, K = conformity_scores.shape + + if test_embeddings is not None: + if self.kde_test is None or self.kde_cal is None: + raise ValueError( + "test_embeddings was provided but no KDEs are available " + "(calibrate() was called with custom cal_weights, which " + "has no way to score a new test point's likelihood " + "ratio). Per-test-point thresholding requires calibrate() " + "to have been called with cal_embeddings/test_embeddings " + "or pre-fitted kde_test/kde_cal." + ) + if self._cal_conformity_scores is None: + raise RuntimeError("Must call calibrate() before forward().") + + test_weights = _compute_likelihood_ratio( + self.kde_test, self.kde_cal, test_embeddings + ) + thresholds = np.empty((N, K), dtype=np.float64) + if isinstance(self.alpha, float): + for i in range(N): + t_i = _query_weighted_quantile( + self._cal_conformity_scores, + self.alpha, + self._cal_likelihood_ratios, + float(test_weights[i]), + ) + thresholds[i, :] = t_i + else: + for i in range(N): + for k in range(K): + if len(self._cal_class_scores[k]) > 0: + thresholds[i, k] = _query_weighted_quantile( + self._cal_class_scores[k], + self.alpha[k], + self._cal_class_weights[k], + float(test_weights[i]), + ) + else: + thresholds[i, k] = -np.inf + threshold_tensor = torch.as_tensor( + thresholds, + device=pred["y_prob"].device, + dtype=pred["y_prob"].dtype, + ) + else: + if not self._warned_fixed_threshold: + warnings.warn( + "CovariateLabel.forward() was called without " + "test_embeddings: falling back to a single fixed " + "threshold computed from the mean calibration " + "likelihood ratio. This is only an approximation of " + "Corollary 1 of Tibshirani et al. (2019), which defines " + "the threshold per test point using that point's own " + "likelihood ratio w(x); pass test_embeddings to get the " + "paper's exact finite-sample coverage guarantee.", + UserWarning, + stacklevel=2, + ) + self._warned_fixed_threshold = True + threshold_tensor = self.t + conformity_scores = torch.as_tensor( conformity_scores, device=pred["y_prob"].device, dtype=pred["y_prob"].dtype ) - pred["y_predset"] = conformity_scores > self.t + pred["y_predset"] = conformity_scores > threshold_tensor return pred diff --git a/tests/core/test_covariate_label.py b/tests/core/test_covariate_label.py index 6f2bfd04a..d244ab95a 100644 --- a/tests/core/test_covariate_label.py +++ b/tests/core/test_covariate_label.py @@ -1,10 +1,15 @@ import unittest +import warnings import numpy as np import torch from pyhealth.datasets import create_sample_dataset, get_dataloader from pyhealth.models import MLP from pyhealth.calib.predictionset.covariate import CovariateLabel, fit_kde +from pyhealth.calib.predictionset.covariate.covariate_label import ( + _compute_likelihood_ratio, + _query_weighted_quantile, +) from pyhealth.calib.utils import extract_embeddings @@ -340,6 +345,135 @@ def test_score_type_aps_runs_end_to_end(self): self.assertEqual(output["y_predset"].dtype, torch.bool) self.assertEqual(output["y_predset"].shape, output["y_prob"].shape) + def _build_pointwise_setup(self, alpha=0.4): + """Build a larger dataset + a deliberately bimodal (non-uniform) + KDE pair, so per-test-point likelihood ratios genuinely differ -- + needed to test forward()'s pointwise-thresholding path, which the + class's default dummy_kde_cal/dummy_kde_test (both ~uniform) can't + exercise meaningfully. + """ + samples = [] + for i in range(60): + samples.append( + { + "patient_id": f"p{i}", + "visit_id": f"v{i}", + "conditions": [f"cond-{i % 10}", f"cond-{(i + 1) % 10}"], + "procedures": [float(i % 5), float((i * 2) % 7), 1.0, 2.0], + "label": i % 3, + } + ) + dataset = create_sample_dataset( + samples=samples, + input_schema={"conditions": "sequence", "procedures": "tensor"}, + output_schema={"label": "multiclass"}, + dataset_name="test_pointwise", + ) + model = MLP( + dataset=dataset, + feature_keys=["conditions", "procedures"], + label_key="label", + mode="multiclass", + ) + model.eval() + + cal_dataset = dataset.subset(list(range(0, 30))) + test_dataset = dataset.subset(list(range(30, 60))) + cal_embeddings = extract_embeddings(model, cal_dataset, batch_size=32, device="cpu") + test_embeddings = extract_embeddings(model, test_dataset, batch_size=32, device="cpu") + + def kde_cal(data): + return np.ones(len(np.asarray(data))) + + def kde_test(data): + data = np.asarray(data) + row_signal = np.abs(data).sum(axis=1) + median = np.median(row_signal) + # Two clearly separated weight regimes, split deterministically + # by an arbitrary per-row criterion, so weights genuinely + # differ by test point. + return np.where(row_signal > median, 10.0, 0.05) + + cal_model = CovariateLabel( + model=model, alpha=alpha, kde_test=kde_test, kde_cal=kde_cal, random_state=0 + ) + cal_model.calibrate( + cal_dataset=cal_dataset, + cal_embeddings=cal_embeddings, + test_embeddings=test_embeddings, + ) + return cal_model, test_dataset, test_embeddings, kde_test, kde_cal + + def test_forward_without_embeddings_warns_and_uses_fixed_threshold(self): + """forward() without test_embeddings must warn that it's only + approximating Corollary 1 (mean calibration weight standing in for + each test point's own w(x)), and must use the single fixed + threshold computed at calibrate() time.""" + cal_model, test_dataset, _, _, _ = self._build_pointwise_setup() + test_loader = get_dataloader(test_dataset, batch_size=30, shuffle=False) + batch = next(iter(test_loader)) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with torch.no_grad(): + cal_model(**batch) + self.assertTrue( + any("approximation of Corollary 1" in str(w.message) for w in caught) + ) + + def test_forward_with_embeddings_uses_per_point_weight(self): + """forward(test_embeddings=...) must recompute the threshold per + test point using that point's own likelihood ratio w(x) (Corollary + 1 of Tibshirani et al. 2019), not the calibrate()-time mean-weight + approximation. Verified by checking the threshold implied for + specific points matches directly calling _query_weighted_quantile + with that exact point's own weight, and that points in the two + weight regimes get different thresholds. + """ + cal_model, test_dataset, test_embeddings, kde_test, kde_cal = ( + self._build_pointwise_setup() + ) + weights = _compute_likelihood_ratio(kde_test, kde_cal, test_embeddings) + self.assertGreater( + len(set(np.round(weights, 3))), 1, "test setup must have varying weights" + ) + + low_idx = int(np.argmin(weights)) + high_idx = int(np.argmax(weights)) + + t_low = _query_weighted_quantile( + cal_model._cal_conformity_scores, + cal_model.alpha, + cal_model._cal_likelihood_ratios, + float(weights[low_idx]), + ) + t_high = _query_weighted_quantile( + cal_model._cal_conformity_scores, + cal_model.alpha, + cal_model._cal_likelihood_ratios, + float(weights[high_idx]), + ) + # The two weight regimes are far enough apart that they must not + # collapse to the calibrate()-time fixed threshold's approximation + # in exactly the same way -- this is the actual regression check: + # before the forward()-level fix, there was no way to get anything + # other than cal_model.t regardless of test_embeddings. + self.assertNotEqual(t_low, t_high) + + def test_forward_raises_without_kde_for_pointwise(self): + """test_embeddings requires KDEs (to score a new point's w(x)); + calibrate() with custom cal_weights has no way to do that, so + forward() must raise rather than silently ignore test_embeddings.""" + cal_model, test_dataset, test_embeddings, _, _ = self._build_pointwise_setup() + cal_model.kde_test = None + cal_model.kde_cal = None + + test_loader = get_dataloader(test_dataset, batch_size=30, shuffle=False) + batch = next(iter(test_loader)) + with self.assertRaises(ValueError): + with torch.no_grad(): + cal_model(test_embeddings=test_embeddings, **batch) + def test_weighted_quantile_function(self): """Test the weighted quantile helper function.""" from pyhealth.calib.predictionset.covariate.covariate_label import ( @@ -403,6 +537,35 @@ def test_weighted_quantile_small_calibration_set_is_conservative(self): ) self.assertTrue(np.isfinite(result_large)) + def test_weighted_quantile_matches_corollary_1(self): + """The finite-sample correction must match Corollary 1 of Tibshirani, + Barber, Candes, and Ramdas, "Conformal Prediction Under Covariate + Shift" (NeurIPS 2019): the reserved test-point mass is an actual + point in the weighted empirical distribution (at -inf, under this + codebase's higher-is-better conformity convention), not merely a + term folded into the normalizing denominator. + + With scores=[0.1,0.3,0.5,0.7,0.9], uniform weights=1, test_weight=1, + alpha=0.3: total_weight=6, p_test=1/6 (< alpha, so not the -inf + fallback). Correctly inserting the reserved mass as the first point + of the augmented (n+1)-point distribution gives cumulative + fractions [2/6, 3/6, 4/6, 5/6, 6/6] = [.333,.5,.667,.833,1.0], so the + alpha=0.3 quantile lands at the first real point: 0.1. + + Before this fix, the reserved mass was only added to the + denominator (cumulative fractions [1/6,2/6,3/6,4/6,5/6]), which + incorrectly returned 0.3 instead -- a stricter, LESS permissive + threshold that under-covers relative to the target. + """ + from pyhealth.calib.predictionset.covariate.covariate_label import ( + _query_weighted_quantile, + ) + + scores = np.array([0.1, 0.3, 0.5, 0.7, 0.9]) + weights = np.ones(5) + result = _query_weighted_quantile(scores, 0.3, weights, test_weight=1.0) + self.assertAlmostEqual(result, 0.1, places=10) + def test_likelihood_ratio_function(self): """Test the likelihood ratio computation.""" from pyhealth.calib.predictionset.covariate.covariate_label import ( From a414d9a66414be13a0e927827a2ce27542bb6f1a Mon Sep 17 00:00:00 2001 From: lehendo Date: Sun, 30 Aug 2026 20:40:34 -0500 Subject: [PATCH 2/3] Fix flaky/broken CovariateLabel per-point weight test CI failed on test_forward_with_embeddings_uses_per_point_weight (AssertionError: 0.3241... == 0.3241...). Root-caused this to a real design flaw, not just missing randomness control: a test point's likelihood-ratio weight is drawn from the same KDE-derived distribution as the calibration weights, so by construction it can never be more than a small fraction of total calibration weight mass. Whether prepending it shifts _query_weighted_quantile's selected order-statistic index depends entirely on where the alpha-quantile boundary happens to fall relative to the (data-dependent, effectively random given the model's unseeded init) distribution of calibration weights along the sorted-score axis. Confirmed this wasn't just 'rare bad luck': reproduced the exact failure deterministically for multiple fixed seeds, and directly inspected the internals (total_weight, cum_weights) showing the two test points' weights (0.05 vs 10.0) simply didn't straddle a boundary for that data -- an inherent fragility in comparing *output values* for two randomly-KDE -derived weights, not a bug in the underlying Corollary-1 implementation. Replaced the flaky output-comparison with two more precise checks: 1. Spy on _query_weighted_quantile during a real forward() call and verify it's invoked once per test point with that exact point's own weight -- directly verifies the claim ('forward uses per-point weight') without depending on the resulting threshold *values* differing. 2. A new, fully hand-computed unit test proving test_weight does change _query_weighted_quantile's result in general (by forcing the documented -inf fallback with a deliberately large weight), isolated from any KDE/model randomness. Also seeded _build_pointwise_setup()'s model init for reproducibility. --- tests/core/test_covariate_label.py | 87 ++++++++++++++++++++++-------- 1 file changed, 64 insertions(+), 23 deletions(-) diff --git a/tests/core/test_covariate_label.py b/tests/core/test_covariate_label.py index d244ab95a..8e4f99b66 100644 --- a/tests/core/test_covariate_label.py +++ b/tests/core/test_covariate_label.py @@ -1,5 +1,7 @@ import unittest import warnings +from unittest.mock import patch + import numpy as np import torch @@ -350,8 +352,12 @@ def _build_pointwise_setup(self, alpha=0.4): KDE pair, so per-test-point likelihood ratios genuinely differ -- needed to test forward()'s pointwise-thresholding path, which the class's default dummy_kde_cal/dummy_kde_test (both ~uniform) can't - exercise meaningfully. + exercise meaningfully. Seeded for reproducibility across runs; the + model's own random init would otherwise depend on however much of + the global RNG stream setUp() and any prior tests already consumed. """ + torch.manual_seed(0) + np.random.seed(0) samples = [] for i in range(60): samples.append( @@ -425,10 +431,27 @@ def test_forward_with_embeddings_uses_per_point_weight(self): """forward(test_embeddings=...) must recompute the threshold per test point using that point's own likelihood ratio w(x) (Corollary 1 of Tibshirani et al. 2019), not the calibrate()-time mean-weight - approximation. Verified by checking the threshold implied for - specific points matches directly calling _query_weighted_quantile - with that exact point's own weight, and that points in the two - weight regimes get different thresholds. + approximation. + + Verified by spying on _query_weighted_quantile during a real + forward() call and checking it is invoked once per test point with + that exact point's own weight -- directly verifying the claim + ("forward uses per-point weight") rather than asserting the + resulting *threshold values* differ between two points. + + The latter was tried first and is not a reliable test: a test + point's likelihood-ratio weight is drawn from the same KDE-derived + distribution as the calibration weights, so it can never be more + than a small fraction of the total calibration weight mass by + construction -- whether prepending it shifts _query_weighted_ + quantile's selected order-statistic index depends on exactly where + the alpha-quantile boundary falls relative to the (data-dependent, + effectively random given the model's unseeded init) distribution + of calibration weights along the sorted-score axis. That made the + old assertion fail whenever the boundary happened to fall in a + region insensitive to a perturbation of that size -- confirmed + directly: it failed deterministically for several concrete seeds, + not just intermittently, so it wasn't simply "rare bad luck." """ cal_model, test_dataset, test_embeddings, kde_test, kde_cal = ( self._build_pointwise_setup() @@ -438,26 +461,44 @@ def test_forward_with_embeddings_uses_per_point_weight(self): len(set(np.round(weights, 3))), 1, "test setup must have varying weights" ) - low_idx = int(np.argmin(weights)) - high_idx = int(np.argmax(weights)) + test_loader = get_dataloader(test_dataset, batch_size=30, shuffle=False) + batch = next(iter(test_loader)) - t_low = _query_weighted_quantile( - cal_model._cal_conformity_scores, - cal_model.alpha, - cal_model._cal_likelihood_ratios, - float(weights[low_idx]), - ) - t_high = _query_weighted_quantile( - cal_model._cal_conformity_scores, - cal_model.alpha, - cal_model._cal_likelihood_ratios, - float(weights[high_idx]), + module = "pyhealth.calib.predictionset.covariate.covariate_label" + with patch( + f"{module}._query_weighted_quantile", + wraps=_query_weighted_quantile, + ) as spy: + with torch.no_grad(): + cal_model(test_embeddings=test_embeddings, **batch) + + self.assertEqual(spy.call_count, len(weights)) + called_test_weights = [call.args[3] for call in spy.call_args_list] + np.testing.assert_allclose( + sorted(called_test_weights), sorted(float(w) for w in weights) ) - # The two weight regimes are far enough apart that they must not - # collapse to the calibrate()-time fixed threshold's approximation - # in exactly the same way -- this is the actual regression check: - # before the forward()-level fix, there was no way to get anything - # other than cal_model.t regardless of test_embeddings. + + def test_query_weighted_quantile_test_weight_changes_result(self): + """Direct, hand-computed check that test_weight actually changes + _query_weighted_quantile's output -- the algorithmic property the + forward()-level test above relies on, isolated from any KDE/model + randomness. A large enough test_weight must push p_test >= alpha, + forcing the documented -inf fallback (not enough calibration mass + to reach the target coverage without dipping into the test point's + own reserved mass); test_weight=0 must not. + """ + scores = np.array([0.1, 0.2, 0.3, 0.4, 0.5]) + weights = np.full(5, 2.0) # total calibration weight = 10.0 + alpha = 0.4 + + t_low = _query_weighted_quantile(scores, alpha, weights, test_weight=0.0) + self.assertNotEqual(t_low, -np.inf) + + # p_test = test_weight / (10 + test_weight) >= 0.4 requires + # test_weight >= 10 * 0.4 / 0.6 ≈ 6.67; 100 clears it by a wide, + # seed-independent margin. + t_high = _query_weighted_quantile(scores, alpha, weights, test_weight=100.0) + self.assertEqual(t_high, -np.inf) self.assertNotEqual(t_low, t_high) def test_forward_raises_without_kde_for_pointwise(self): From f54f0fe2358d687edb197ccadc4b0c3078b0a601 Mon Sep 17 00:00:00 2001 From: lehendo Date: Thu, 3 Sep 2026 00:12:12 -0500 Subject: [PATCH 3/3] Fix CovariateLabel threshold boundary and example's dead per-point path Two independent bugs, both flagged in review: 1. forward() included a class only if conformity_score > threshold (strict). _query_weighted_quantile defines the threshold as the smallest score whose cumulative weight reaches alpha, so that score's own mass is part of what clears the target -- excluding it with a strict comparison silently under-covers whenever a real score lands on a calibration tie. scores.py's own all_class_conformity_scores docstring already documents ">=" as the intended convention for this sign convention; forward() just didn't follow it. Fixed to ">=", and added a length-mismatch check for test_embeddings (previously extra rows were silently ignored rather than raising). 2. The per-test-point thresholding path (forward(test_embeddings=...), the actual Corollary 1 construction from Tibshirani et al. 2019) has existed since the previous commit on this branch, but examples/conformal_eeg/tuev_covariate_shift_conformal.py evaluated it via Trainer(model=cov_predictor).inference(test_loader, ...), which calls model(**batch) with no way to also supply a per-batch test_embeddings slice. So the example -- meant to demonstrate and validate this exact PR's fix -- silently only ever exercised the approximate mean-weight fallback, never the real per-point guarantee. Replaced with a manual inference loop that slices test_embeddings in the same fixed order test_loader yields batches in (mirroring Trainer.inference()'s own logic) and passes each batch's slice through. Added a regression test for (1) that directly forces a conformity score equal to the threshold and checks it's included. --- .../tuev_covariate_shift_conformal.py | 34 ++++++++++++++- .../covariate/covariate_label.py | 14 ++++++- tests/core/test_covariate_label.py | 42 +++++++++++++++++++ 3 files changed, 87 insertions(+), 3 deletions(-) diff --git a/examples/conformal_eeg/tuev_covariate_shift_conformal.py b/examples/conformal_eeg/tuev_covariate_shift_conformal.py index 58efd0c2e..eed52028a 100644 --- a/examples/conformal_eeg/tuev_covariate_shift_conformal.py +++ b/examples/conformal_eeg/tuev_covariate_shift_conformal.py @@ -179,6 +179,36 @@ def set_seed(seed: int) -> None: torch.cuda.manual_seed_all(seed) +def _inference_with_pointwise_threshold(cov_predictor, test_loader, test_embeddings): + """Run CovariateLabel inference with each test point's own threshold. + + Trainer.inference() calls model(**batch) with no way to also pass a + per-batch test_embeddings slice, so routing through it would silently + fall back to CovariateLabel's approximate mean-weight threshold for + every point (defeating the whole point of this example -- exercising + Corollary 1's exact per-test-point construction). This loop mirrors + Trainer.inference()'s logic directly, slicing test_embeddings in the + same fixed order test_loader yields batches in (shuffle=False). + """ + y_true_all, y_prob_all, y_predset_all = [], [], [] + cov_predictor.eval() + offset = 0 + with torch.no_grad(): + for batch in test_loader: + batch_size = len(batch["patient_id"]) + batch_embeddings = test_embeddings[offset : offset + batch_size] + offset += batch_size + output = cov_predictor(test_embeddings=batch_embeddings, **batch) + y_true_all.append(output["y_true"].cpu().numpy()) + y_prob_all.append(output["y_prob"].cpu().numpy()) + y_predset_all.append(output["y_predset"].cpu().numpy()) + return ( + np.concatenate(y_true_all, axis=0), + np.concatenate(y_prob_all, axis=0), + {"y_predset": np.concatenate(y_predset_all, axis=0)}, + ) + + def _run_one_seed( args, sample_dataset, @@ -249,8 +279,8 @@ def _run_one_seed( test_embeddings=test_embeddings, ) - y_true, y_prob, _, extra = Trainer(model=cov_predictor).inference( - test_loader, additional_outputs=["y_predset"] + y_true, y_prob, extra = _inference_with_pointwise_threshold( + cov_predictor, test_loader, test_embeddings ) conf_metrics = get_metrics_fn("multiclass")( y_true, y_prob, diff --git a/pyhealth/calib/predictionset/covariate/covariate_label.py b/pyhealth/calib/predictionset/covariate/covariate_label.py index 1ff0b6b0e..e33808b9a 100644 --- a/pyhealth/calib/predictionset/covariate/covariate_label.py +++ b/pyhealth/calib/predictionset/covariate/covariate_label.py @@ -645,6 +645,12 @@ def forward( ) if self._cal_conformity_scores is None: raise RuntimeError("Must call calibrate() before forward().") + if len(test_embeddings) != N: + raise ValueError( + f"test_embeddings has {len(test_embeddings)} rows but " + f"the batch has {N} examples; they must be aligned " + "one-to-one in the same order." + ) test_weights = _compute_likelihood_ratio( self.kde_test, self.kde_cal, test_embeddings @@ -696,7 +702,13 @@ def forward( conformity_scores = torch.as_tensor( conformity_scores, device=pred["y_prob"].device, dtype=pred["y_prob"].dtype ) - pred["y_predset"] = conformity_scores > threshold_tensor + # Corollary 1 (Tibshirani et al. 2019) includes a class whose score + # is *exactly* at the threshold: the weighted quantile is defined as + # the smallest score whose cumulative weight reaches alpha, so that + # score's own mass is part of what clears the target -- excluding it + # with a strict ">" would count it against coverage it was actually + # counted toward, under-covering whenever ties land on the threshold. + pred["y_predset"] = conformity_scores >= threshold_tensor return pred diff --git a/tests/core/test_covariate_label.py b/tests/core/test_covariate_label.py index 8e4f99b66..10eea8aac 100644 --- a/tests/core/test_covariate_label.py +++ b/tests/core/test_covariate_label.py @@ -277,6 +277,48 @@ def test_forward_returns_predset(self): # Check prediction set shape matches probability shape self.assertEqual(output["y_predset"].shape, output["y_prob"].shape) + def test_forward_includes_class_tied_exactly_at_threshold(self): + """Regression test: a class whose conformity score is exactly equal + to the threshold must be included in the prediction set. + + _query_weighted_quantile defines the threshold as the smallest + score whose cumulative weight reaches alpha -- that score's own + mass is part of what clears the target coverage. A strict '>' + comparison in forward() would exclude that exact score, silently + under-covering whenever a real conformity score lands on a + calibration tie. This also matches the convention already + documented in scores.py's all_class_conformity_scores docstring + ("threshold with score >= t rather than nc_score <= t"). + """ + cal_model = CovariateLabel( + model=self.model, + alpha=0.2, + kde_test=self.kde_test, + kde_cal=self.kde_cal, + ) + cal_indices = [0, 1, 2, 3] + cal_dataset = self.dataset.subset(cal_indices) + cal_embeddings = self._get_embeddings(cal_dataset) + test_embeddings = self._get_embeddings(self.dataset) + cal_model.calibrate( + cal_dataset=cal_dataset, + cal_embeddings=cal_embeddings, + test_embeddings=test_embeddings, + ) + + threshold_value = float(cal_model.t.item()) + fake_y_prob = torch.tensor([[threshold_value, 0.0, 0.0]]) + with patch.object( + cal_model.model, "forward", return_value={"y_prob": fake_y_prob} + ): + output = cal_model() + + self.assertTrue( + bool(output["y_predset"][0, 0]), + "a class whose conformity score exactly equals the threshold " + "must be included in the prediction set", + ) + def test_prediction_sets_nonempty(self): """Test that prediction sets are non-empty for most examples.""" cal_model = CovariateLabel(