diff --git a/docs/api/calib/pyhealth.calib.predictionset.rst b/docs/api/calib/pyhealth.calib.predictionset.rst index 740b1c87c..4b29a22fe 100644 --- a/docs/api/calib/pyhealth.calib.predictionset.rst +++ b/docs/api/calib/pyhealth.calib.predictionset.rst @@ -49,6 +49,12 @@ LABEL (Least Ambiguous Set-valued Classifier) SCRIB (Set-classifier with Class-specific Risk Bounds) ------------------------------------------------------- +SCRIB's threshold search assumes the calibration set's empty-prediction +handling (``fill_max``) matches what is applied at inference time; both +``calibrate()`` and ``forward()`` resolve and use the same ``fill_max`` +value, so calibration is never optimized against behavior that inference +doesn't actually apply. + .. autoclass:: pyhealth.calib.predictionset.SCRIB :members: :undoc-members: diff --git a/examples/cxr/covid19cxr_conformal.py b/examples/cxr/covid19cxr_conformal.py index 47e2eebbc..4d71c6dd4 100644 --- a/examples/cxr/covid19cxr_conformal.py +++ b/examples/cxr/covid19cxr_conformal.py @@ -6,6 +6,11 @@ 2. Conventional conformal prediction using LABEL 3. Covariate shift adaptive conformal prediction using CovariateLabel 4. Comparison of coverage and efficiency between the two methods + +For class-specific risk control instead of the marginal/class-conditional +mis-coverage guarantees LABEL and CovariateLabel provide, see +pyhealth.calib.predictionset.SCRIB, which minimizes ambiguity subject to +per-class risk targets (see its docstring for usage). """ import numpy as np diff --git a/pyhealth/calib/predictionset/favmac/core.py b/pyhealth/calib/predictionset/favmac/core.py index 7e9b7d354..9e15e5633 100644 --- a/pyhealth/calib/predictionset/favmac/core.py +++ b/pyhealth/calib/predictionset/favmac/core.py @@ -11,6 +11,41 @@ class FavMac: + """Online value-maximizing prediction sets with conformal cost control. + + This is the internal calibration/inference engine backing the + user-facing :class:`pyhealth.calib.predictionset.FavMac`. It is + abstract: ``_greedy_sequence`` must be supplied by a subclass (see + :class:`FavMac_GreedyRatio`, the concrete class the public wrapper + actually uses). Costs, values, and their proxies must be normalized + so cost lies in ``[0, C_max]`` before being passed in here. + + Paper: + Lin, Zhen, Shubhendu Trivedi, Cao Xiao, and Jimeng Sun. "Fast + Online Value-Maximizing Prediction Sets with Conformal Cost + Control." ICML 2023. + + Examples: + >>> import numpy as np + >>> from pyhealth.calib.predictionset.favmac import AdditiveSetFunction + >>> from pyhealth.calib.predictionset.favmac.core import FavMac_GreedyRatio + >>> K = 3 + >>> C_max = float(K) + >>> cost_fn = AdditiveSetFunction(np.ones(K) / C_max, mode="cost") + >>> util_fn = AdditiveSetFunction(np.ones(K), mode="util") + >>> proxy_fn = AdditiveSetFunction(np.ones(K) / C_max, mode="proxy") + >>> fm = FavMac_GreedyRatio( + ... cost_fn, util_fn, proxy_fn, target_cost=1.0 / C_max, C_max=1.0) + >>> rng = np.random.default_rng(0) + >>> for _ in range(30): + ... logit = rng.normal(size=K) + ... y = (rng.uniform(size=K) < 0.4).astype(int) + ... _ = fm.update(logit, y) + >>> predset, _ = fm(np.array([1.0, -0.5, 0.2]), update=False) + >>> predset + array([1, 0, 0]) + """ + def __init__(self, cost_fn, util_fn, proxy_fn, target_cost, delta=None, C_max=1.) -> None: self.target_cost = target_cost self.delta = delta @@ -56,7 +91,15 @@ def _query_threshold(self): cutoff = self.target_cost * (n+1) - self.C_max return self.quantiletree.query_cumu_weight(cutoff, prev=False) else: - cutoff = self.delta * (n+1) - 1# We should assume a violation for the next point? Should we minus 1?? + # The "-1" reserves one unit of probability mass for the + # unseen (N+1)-th test point, the same finite-sample + # correction used throughout split conformal prediction (see + # e.g. base_conformal._query_quantile's ceil((1-alpha)(N+1))). + # This matches the paper's own Appendix B.3, Algorithm 5 + # (COMPUTE THRESHOLD: q = ((N+1)*delta - 1) / tree.root.sum), + # which follows directly from the Tc,delta derivation in the + # proof of Theorem 4.6 (Appendix A.4). + cutoff = self.delta * (n+1) - 1 return self.quantiletree.query_cumu_weight(cutoff, prev=False) def _greedy_sequence(self, pred:np.ndarray): diff --git a/pyhealth/calib/predictionset/scrib/__init__.py b/pyhealth/calib/predictionset/scrib/__init__.py index 3406bef3b..c0e703bb2 100644 --- a/pyhealth/calib/predictionset/scrib/__init__.py +++ b/pyhealth/calib/predictionset/scrib/__init__.py @@ -200,10 +200,15 @@ class SCRIB(SetPredictor): The higher the lk, the more penalty on risk violation (likely higher ambiguity). fill_max: Whether to fill the class with max predicted score when no class exceeds the threshold. In other words, if fill_max, - the null region will be filled with max-prediction class. - Defaults to {'lk': 1e4, 'fill_max': False} + the null region will be filled with max-prediction class. This is + applied consistently during both threshold search and inference + (``forward()``). If you pass ``loss_kwargs`` explicitly without a + ``fill_max`` key, it defaults to False here (regardless of the + ``fill_max`` argument below), matching the underlying search + routines' own default. + Defaults to {'lk': 1e4, 'fill_max': fill_max} (see the ``fill_max`` argument below). fill_max (bool, optional): Whether to fill the empty prediction set with the max-predicted class. - Defaults to True. + Only takes effect when ``loss_kwargs`` is left as None. Defaults to True. Examples: @@ -264,6 +269,10 @@ def __init__( if loss_kwargs is None: loss_kwargs = {"lk": 1e4, "fill_max": fill_max} self.loss_kwargs = loss_kwargs + # The value actually used by the threshold search (loss_kwargs may + # silently override the fill_max argument above); forward() must + # apply the same fill-max behavior it was calibrated under. + self.fill_max = loss_kwargs.get("fill_max", False) self.t = None @@ -297,7 +306,16 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: :rtype: Dict[str, torch.Tensor] """ ret = self.model(**kwargs) - ret["y_predset"] = ret["y_prob"] > self.t + y_predset = ret["y_prob"] > self.t + if self.fill_max: + # Match the calibration-time assumption: when no class clears + # its threshold, fall back to the max-predicted class instead + # of returning an empty set. + empty = y_predset.sum(dim=1) == 0 + if empty.any(): + argmax_idx = ret["y_prob"].argmax(dim=1) + y_predset[empty, argmax_idx[empty]] = True + ret["y_predset"] = y_predset return ret diff --git a/pyhealth/calib/predictionset/scrib/quicksearch.py b/pyhealth/calib/predictionset/scrib/quicksearch.py index 09ebf910f..e63b07d75 100644 --- a/pyhealth/calib/predictionset/scrib/quicksearch.py +++ b/pyhealth/calib/predictionset/scrib/quicksearch.py @@ -20,7 +20,7 @@ def _thresholding_py(ts, output): return pred def __loss_overall_helper(total_err, total_sure, alpha, N, la, lc, lcs): - ambiguity_loss = (1. - total_sure / float(N)) ** 2 + ambiguity_loss = 1. - total_sure / float(N) risk = total_err / float(max(total_sure, 1)) tempf = risk - alpha coverage_loss = np.power(max(tempf, 0), 2) diff --git a/pyhealth/calib/predictionset/scrib/quicksearch_cython.pyx b/pyhealth/calib/predictionset/scrib/quicksearch_cython.pyx index 52a3c2eda..5e8e20b10 100644 --- a/pyhealth/calib/predictionset/scrib/quicksearch_cython.pyx +++ b/pyhealth/calib/predictionset/scrib/quicksearch_cython.pyx @@ -26,7 +26,7 @@ cdef int _update_counts(int pred, int truth, int* err, int* sure, int inc): cdef double loss_overall_helper__(int total_err, int total_sure, double alpha, int N, double la, double lc, double lcs): #if total_sure == 0: return np.inf - cdef double a_loss = (1. - total_sure / N) ** 2 + cdef double a_loss = (1. - total_sure / N) cdef double tempf = total_err / max(total_sure,1) - alpha cdef double c_loss = max(tempf, 0.) ** 2 cdef double cs_loss = tempf ** 2 diff --git a/tests/core/test_scrib.py b/tests/core/test_scrib.py new file mode 100644 index 000000000..336c5a4e8 --- /dev/null +++ b/tests/core/test_scrib.py @@ -0,0 +1,166 @@ +"""Tests for SCRIB, focused on the fill_max inference behavior and the +overall-risk loss's ambiguity term. +""" + +import unittest + +import numpy as np +import torch + +from pyhealth.calib.predictionset import SCRIB +from pyhealth.calib.predictionset.scrib.quicksearch import loss_overall_py +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.models import MLP + + +class TestSCRIB(unittest.TestCase): + """Test cases for the SCRIB prediction set constructor.""" + + def setUp(self): + np.random.seed(0) + torch.manual_seed(0) + + # 5 classes, enough samples per class for stable class-specific + # thresholds during calibration. + self.samples = [ + { + "patient_id": f"p{i}", + "visit_id": f"v{i}", + "procedures": np.random.randn(6).tolist(), + "label": i % 5, + } + for i in range(60) + ] + self.dataset = create_sample_dataset( + samples=self.samples, + input_schema={"procedures": "tensor"}, + output_schema={"label": "multiclass"}, + dataset_name="test", + ) + self.model = MLP( + dataset=self.dataset, + feature_keys=["procedures"], + label_key="label", + mode="multiclass", + ) + self.model.eval() + + def test_fill_max_default_resolves_true(self): + """fill_max=True is the constructor default and should be reflected + in the resolved self.fill_max used by forward().""" + m = SCRIB(self.model, risk=0.2) + self.assertTrue(m.fill_max) + + def test_fill_max_false(self): + m = SCRIB(self.model, risk=0.2, fill_max=False) + self.assertFalse(m.fill_max) + + def test_loss_kwargs_without_fill_max_key_resolves_false(self): + """Passing loss_kwargs explicitly without a 'fill_max' key must not + silently inherit the fill_max=True constructor default -- it should + match the underlying search routines' own default of False.""" + m = SCRIB(self.model, risk=0.2, loss_kwargs={"lk": 500.0}) + self.assertFalse(m.fill_max) + + def test_loss_kwargs_with_explicit_fill_max_key(self): + m = SCRIB(self.model, risk=0.2, loss_kwargs={"lk": 500.0, "fill_max": True}) + self.assertTrue(m.fill_max) + + def test_forward_never_empty_when_fill_max_true(self): + """The core regression test: forward() must apply the same + fill_max behavior used during calibration, so no empty prediction + sets should ever be returned when fill_max=True -- even when + thresholds are (artificially) set so high that no class clears + them naturally.""" + m = SCRIB(self.model, risk=0.1, fill_max=True) + m.calibrate(cal_dataset=self.dataset) + # Force every threshold above any predicted probability, so every + # sample's natural prediction set is empty prior to the fill_max + # fallback. + m.t = torch.nn.Parameter(torch.ones_like(m.t) * 0.999) + + loader = get_dataloader(self.dataset, batch_size=len(self.samples), shuffle=False) + batch = next(iter(loader)) + with torch.no_grad(): + out = m(**batch) + + set_sizes = out["y_predset"].sum(dim=1) + self.assertTrue(torch.all(set_sizes == 1), "fill_max=True must fill every empty set with exactly the argmax class") + + def test_forward_allows_empty_when_fill_max_false(self): + """Sanity check that fill_max actually gates the behavior: with the + same forced thresholds, fill_max=False should leave sets empty.""" + m = SCRIB(self.model, risk=0.1, fill_max=False) + m.calibrate(cal_dataset=self.dataset) + m.t = torch.nn.Parameter(torch.ones_like(m.t) * 0.999) + + loader = get_dataloader(self.dataset, batch_size=len(self.samples), shuffle=False) + batch = next(iter(loader)) + with torch.no_grad(): + out = m(**batch) + + set_sizes = out["y_predset"].sum(dim=1) + self.assertTrue(torch.all(set_sizes == 0), "fill_max=False should leave sets empty when no class clears threshold") + + def test_forward_fills_with_argmax_class(self): + """The filled-in class for an empty set must be the model's own + argmax prediction, not an arbitrary class.""" + m = SCRIB(self.model, risk=0.1, fill_max=True) + m.calibrate(cal_dataset=self.dataset) + m.t = torch.nn.Parameter(torch.ones_like(m.t) * 0.999) + + loader = get_dataloader(self.dataset, batch_size=len(self.samples), shuffle=False) + batch = next(iter(loader)) + with torch.no_grad(): + base_out = self.model(**batch) + out = m(**batch) + + argmax_idx = base_out["y_prob"].argmax(dim=1) + predicted_idx = out["y_predset"].float().argmax(dim=1) + torch.testing.assert_close(predicted_idx, argmax_idx) + + def test_overall_risk_calibration_runs_and_controls_risk(self): + """Overall (float) risk mode should calibrate without error and the + resulting sure-prediction error rate should be near the target.""" + m = SCRIB(self.model, risk=0.2) + m.calibrate(cal_dataset=self.dataset) + self.assertIsNotNone(m.t) + self.assertEqual(m.t.shape[0], 5) + + def test_class_specific_risk_calibration_runs(self): + """Class-specific (array) risk mode should calibrate without error + and produce one threshold per class.""" + risk = np.array([0.2, 0.3, 0.15, 0.25, 0.2]) + m = SCRIB(self.model, risk=risk) + m.calibrate(cal_dataset=self.dataset) + self.assertEqual(m.t.shape[0], 5) + + def test_forward_before_calibration_uses_none_threshold(self): + m = SCRIB(self.model, risk=0.2) + self.assertIsNone(m.t) + + def test_ambiguity_term_is_not_squared(self): + """The overall-risk loss's chance-ambiguity term must be linear + (1 - total_sure/N), matching the paper's Eq. 2 / Algorithm 2, not + squared. Regression test for the fixed bug.""" + n, total_sure = 20, 12 + preds = np.zeros((n, 3), dtype=np.int32) + # First `total_sure` rows: exactly one class included (|H|=1). + preds[:total_sure, 0] = 1 + # Remaining rows: two classes included (|H|=2, ambiguous). + preds[total_sure:, 0] = 1 + preds[total_sure:, 1] = 1 + labels = np.zeros((n, 3)) + labels[:, 0] = 1 # every true label is class 0 + max_classes = np.zeros(n, dtype=np.int32) + + loss = loss_overall_py(preds, labels, max_classes, risk=0.5, lk=1e4, fill_max=False) + # total_err=0 (every "sure" row correctly includes class 0), so the + # risk-penalty term is 0 and the loss should equal the unsquared + # ambiguity term exactly: 1 - total_sure/N. + expected = 1.0 - total_sure / n + self.assertAlmostEqual(loss, expected, places=10) + + +if __name__ == "__main__": + unittest.main()