From 2bc7f4e9179e50db7c09ef365e0872452561c6f1 Mon Sep 17 00:00:00 2001 From: lehendo Date: Tue, 25 Aug 2026 20:01:28 -0500 Subject: [PATCH 1/2] 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/2] 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):