Fix SCRIB bugs

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:

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

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

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
 
 

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)

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 / <double> N) ** 2
+    cdef double a_loss = (1. - total_sure / <double> N)
     cdef double tempf = total_err / <double> max(total_sure,1) - alpha
     cdef double c_loss = max(tempf, 0.) ** 2
     cdef double cs_loss = tempf ** 2

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