Add EnbPI conformal prediction intervals for time series regression

docs/api/calib.rst

@@ -27,6 +27,10 @@ confidence levels:
 - :class:`~pyhealth.calib.predictionset.ClusterLabel`: K-means cluster-based conformal prediction
 - :class:`~pyhealth.calib.predictionset.NeighborhoodLabel`: Neighborhood Conformal Prediction (NCP)
 
+**Prediction Interval Constructors** give online intervals for time series regression:
+
+- :class:`~pyhealth.calib.predictionset.EnbPI`: Ensemble batch prediction intervals, approximately valid without exchangeability
+
 Getting Started
 ---------------
 

docs/api/calib/pyhealth.calib.predictionset.rst

@@ -31,6 +31,7 @@ Available Methods
    pyhealth.calib.predictionset.RC3P
    pyhealth.calib.predictionset.ClusterLabel
    pyhealth.calib.predictionset.NeighborhoodLabel
+   pyhealth.calib.predictionset.EnbPI
 
 BaseConformal (Standard Split Conformal Prediction)
 ----------------------------------------------------
@@ -123,6 +124,37 @@ NeighborhoodLabel (Neighborhood Conformal Prediction)
    :undoc-members:
    :show-inheritance:
 
+EnbPI (Ensemble Batch Prediction Intervals)
+--------------------------------------------
+
+EnbPI builds prediction intervals for time series regression rather than
+prediction sets. It trains its own bootstrap ensemble from a ``model_fn``
+factory, so it uses no separate calibration split. Intervals are returned as
+``y_lower`` and ``y_upper`` next to the point prediction ``y_prob``. In
+deployment, call ``forward`` for intervals and ``update`` once labels are
+revealed. ``predict_sequential`` is a convenience method for evaluation: it
+runs the paper's online predict, reveal and update loop over a fully labeled
+test stream. Score the results with
+:func:`pyhealth.metrics.regression_metrics_fn`, passing ``y_lower`` and
+``y_upper`` and asking for ``miscoverage_pi``, ``interval_width`` or
+``rolling_miscoverage_pi@W``.
+
+.. autoclass:: pyhealth.calib.predictionset.EnbPI
+   :members: calibrate, forward, update, reset, predict_sequential
+   :show-inheritance:
+
+Interval Interface and Time Series Helpers
+------------------------------------------
+
+``IntervalPredictor`` fixes only the public surface of interval methods,
+``calibrate``, ``forward`` and ``update``. Each time series method implements
+its own online loop, using the helper functions below where they fit.
+
+.. autoclass:: pyhealth.calib.base_classes.IntervalPredictor
+   :members: calibrate, forward, update
+
+.. autofunction:: pyhealth.calib.predictionset.timeseries.conformal_quantile
+
 Helper Functions
 ----------------
 

docs/api/metrics.rst

@@ -25,6 +25,7 @@ reference includes the ECE definition and formula used by these helpers.
     metrics/pyhealth.metrics.calibration
     metrics/pyhealth.metrics.bootstrap
     metrics/pyhealth.metrics.prediction_set
+    metrics/pyhealth.metrics.prediction_interval
     metrics/pyhealth.metrics.fairness
     metrics/pyhealth.metrics.interpretability
     metrics/pyhealth.metrics.generative

docs/api/metrics/pyhealth.metrics.prediction_interval.rst

@@ -0,0 +1,5 @@
+[core] prediction_interval
+===================================
+
+.. automodule:: pyhealth.metrics.prediction_interval
+    :members:

examples/conformal_timeseries/enbpi_synthetic.py

@@ -0,0 +1,94 @@
+"""EnbPI prediction intervals on a synthetic time series.
+
+Builds an AR(1) series whose noise level doubles halfway through the test
+period, trains an EnbPI ensemble of small MLPs on the first part, and runs
+the online predict, reveal and update loop over the rest. It prints overall
+miscoverage and width, and the lowest and highest miscoverage over 100-step
+windows, for two feedback batch sizes ``s``.
+
+Expect overall miscoverage a little above the 0.1 target and a rise in
+rolling miscoverage after the shift. The residual window still holds ``T``
+residuals from before the shift and needs about ``T`` steps to catch up
+(Remark 2 of the paper). Adaptive methods such as ACI target this lag.
+
+Paper:
+    Xu, Chen, and Yao Xie. "Conformal prediction interval for dynamic
+    time-series." ICML 2021.
+
+Run:
+    python examples/conformal_timeseries/enbpi_synthetic.py
+"""
+
+import numpy as np
+import torch
+
+from pyhealth.calib.predictionset import EnbPI
+from pyhealth.datasets import create_sample_dataset
+from pyhealth.metrics import regression_metrics_fn
+from pyhealth.models import MLP
+
+N_TRAIN = 300
+N_TEST = 900
+LAGS = 3
+ALPHA = 0.1
+
+
+def make_series(seed: int = 0):
+    """AR(1) series whose noise standard deviation doubles mid-test."""
+    rng = np.random.default_rng(seed)
+    n = LAGS + N_TRAIN + N_TEST
+    shift_at = LAGS + N_TRAIN + N_TEST // 2
+    y = np.zeros(n)
+    for t in range(1, n):
+        noise_sd = 1.0 if t < shift_at else 2.0
+        y[t] = 0.8 * y[t - 1] + rng.normal(scale=noise_sd)
+    samples = [
+        {"patient_id": "p0", "x": y[t - LAGS : t].tolist(), "y": float(y[t])}
+        for t in range(LAGS, n)
+    ]
+    return create_sample_dataset(samples, {"x": "tensor"}, {"y": "regression"})
+
+
+def main():
+    torch.manual_seed(0)
+    series = make_series()
+    train = series.subset(range(N_TRAIN))
+    test = series.subset(range(N_TRAIN, N_TRAIN + N_TEST))
+
+    for feedback_size in (1, 50):
+        enbpi = EnbPI(
+            model_fn=lambda: MLP(
+                dataset=series,
+                feature_keys=["x"],
+                label_key="y",
+                mode="regression",
+                hidden_dim=16,
+            ),
+            alpha=ALPHA,
+            n_bootstrap=20,
+            train_kwargs={"epochs": 20},
+            feedback_size=feedback_size,
+            random_state=0,
+        )
+        enbpi.calibrate(train)
+        out = enbpi.predict_sequential(test)
+
+        scores = regression_metrics_fn(
+            out["y_true"],
+            out["y_pred"],
+            metrics=["miscoverage_pi", "interval_width", "rolling_miscoverage_pi@100"],
+            y_lower=out["y_lower"],
+            y_upper=out["y_upper"],
+        )
+        rolling = scores["rolling_miscoverage_pi@100"]
+        print(f"s = {feedback_size}")
+        print(f"  target miscoverage   {ALPHA:.2f}")
+        print(f"  miscoverage          {scores['miscoverage_pi']:.3f}")
+        print(f"  mean width           {scores['interval_width']:.3f}")
+        print(
+            f"  rolling miscoverage  {rolling.min():.3f} min, {rolling.max():.3f} max"
+        )
+
+
+if __name__ == "__main__":
+    main()

pyhealth/calib/README.md

@@ -172,6 +172,23 @@ Uses importance weighting (likelihood ratios w(x) = p_test(x)/p_cal(x)) to corre
 - Tibshirani, Ryan J., Rina Foygel Barber, Emmanuel Candes, and Aaditya Ramdas. "Conformal prediction under covariate shift." NeurIPS 2019. https://arxiv.org/abs/1904.06019
 - Laghuvarapu, Siddhartha, Zhen Lin, and Jimeng Sun. "Conformal Drug Property Prediction with Density Estimation under Covariate Shift." NeurIPS 2023. https://arxiv.org/abs/2310.12033
 
+## Prediction Interval Methods (Time Series Conformal Prediction)
+
+These methods build prediction intervals `[y_lower, y_upper]` for regression on time series. Data arriving in time order is not exchangeable, so the split conformal guarantee above does not apply. These methods instead update online: as labels are revealed, recent residuals replace old ones, so the interval width tracks the current error level without refitting the model.
+
+### EnbPI (Ensemble Batch Prediction Intervals)
+
+**Modes**: `regression`
+
+**Class**: `pyhealth.calib.predictionset.EnbPI`
+
+EnbPI trains a bootstrap ensemble on the training series. Every training point gets an out-of-bag prediction from the models whose bootstrap sample left it out, and the absolute residuals of these predictions form a sliding window. The interval is the ensemble prediction plus or minus the (1-α) quantile of the window. Each time a batch of `s` test labels is revealed, their residuals enter the window and the oldest leave, with no model refit. No separate calibration split is needed.
+
+**Guarantee**: Approximate marginal coverage, P(Y_t ∈ C(X_t)) ≥ 1-α - o(1) as the training length grows, without exchangeability. It assumes a stationary, strongly mixing error process and out-of-bag predictors that estimate the regression function well. This is not a finite-sample guarantee.
+
+**Reference**:
+- Xu, Chen, and Yao Xie. "Conformal prediction interval for dynamic time-series." ICML 2021. https://proceedings.mlr.press/v139/xu21h.html
+
 ## Usage
 
 See the [full documentation](https://pyhealth.readthedocs.io/en/latest/api/calib.html) for detailed API references and examples.

pyhealth/calib/base_classes.py

@@ -22,6 +22,36 @@ def to(self, device):
         return self
 
 
+class IntervalPredictor(ABC, torch.nn.Module):
+    """Interface for regression methods that output prediction intervals.
+
+    ``forward`` returns the point prediction as ``y_prob`` together with
+    ``y_lower`` and ``y_upper``. Online methods also learn from revealed
+    labels through ``update``.
+
+    Examples:
+        >>> from pyhealth.calib.predictionset import EnbPI
+        >>> predictor = EnbPI(model_fn=model_fn, alpha=0.1)
+        >>> predictor.calibrate(train_ds)
+        >>> out = predictor(**batch)  # y_prob, y_lower and y_upper
+        >>> predictor.update(revealed_ds)  # learn from newly revealed labels
+    """
+
+    def calibrate(self, cal_dataset):
+        ...
+
+    def forward(self, **kwargs) -> dict[str, torch.Tensor]:
+        ...
+
+    def update(self, revealed_dataset):
+        ...
+
+    def to(self, device):
+        super().to(device)
+        self.device = device
+        return self
+
+
 class SetPredictor(ABC, torch.nn.Module):
     def __init__(self, model, **kwargs) -> None:
         super().__init__()

pyhealth/calib/predictionset/__init__.py

@@ -8,6 +8,7 @@
 from pyhealth.calib.predictionset.label_shift import LabelShift, bbse, em
 from pyhealth.calib.predictionset.rc3p import RC3P
 from pyhealth.calib.predictionset.scrib import SCRIB
+from pyhealth.calib.predictionset.timeseries import EnbPI
 
 __all__ = [
     "BaseConformal",
@@ -21,4 +22,5 @@
     "em",
     "ClusterLabel",
     "NeighborhoodLabel",
+    "EnbPI",
 ]

pyhealth/calib/predictionset/timeseries/__init__.py

@@ -0,0 +1,6 @@
+"""Online conformal prediction intervals for time series regression."""
+
+from pyhealth.calib.predictionset.timeseries.enbpi import EnbPI
+from pyhealth.calib.predictionset.timeseries.utils import conformal_quantile
+
+__all__ = ["EnbPI", "conformal_quantile"]