From f2b84b54f9f4ef3b8eef2c8cfb52d4e2c18731e6 Mon Sep 17 00:00:00 2001 From: Felipe Bonchristiano Date: Thu, 29 Oct 2026 21:55:35 -0500 Subject: [PATCH 1/5] add EnbPI conformal prediction intervals for time series regression --- docs/api/calib.rst | 4 + .../calib/pyhealth.calib.predictionset.rst | 32 ++ docs/api/metrics.rst | 1 + .../pyhealth.metrics.prediction_interval.rst | 5 + .../conformal_timeseries/enbpi_synthetic.py | 94 ++++ pyhealth/calib/README.md | 17 + pyhealth/calib/base_classes.py | 25 +- pyhealth/calib/predictionset/__init__.py | 2 + .../predictionset/timeseries/__init__.py | 6 + .../calib/predictionset/timeseries/enbpi.py | 426 ++++++++++++++++++ .../calib/predictionset/timeseries/utils.py | 66 +++ pyhealth/metrics/prediction_interval.py | 88 ++++ pyhealth/metrics/regression.py | 46 +- tests/core/test_enbpi.py | 272 +++++++++++ .../core/test_prediction_interval_metrics.py | 110 +++++ tests/core/test_timeseries_cp_utils.py | 52 +++ tests/core/timeseries_cp_test_utils.py | 56 +++ 17 files changed, 1300 insertions(+), 2 deletions(-) create mode 100644 docs/api/metrics/pyhealth.metrics.prediction_interval.rst create mode 100644 examples/conformal_timeseries/enbpi_synthetic.py create mode 100644 pyhealth/calib/predictionset/timeseries/__init__.py create mode 100644 pyhealth/calib/predictionset/timeseries/enbpi.py create mode 100644 pyhealth/calib/predictionset/timeseries/utils.py create mode 100644 pyhealth/metrics/prediction_interval.py create mode 100644 tests/core/test_enbpi.py create mode 100644 tests/core/test_prediction_interval_metrics.py create mode 100644 tests/core/test_timeseries_cp_utils.py create mode 100644 tests/core/timeseries_cp_test_utils.py diff --git a/docs/api/calib.rst b/docs/api/calib.rst index acaac1e35..a123c17a6 100644 --- a/docs/api/calib.rst +++ b/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 --------------- diff --git a/docs/api/calib/pyhealth.calib.predictionset.rst b/docs/api/calib/pyhealth.calib.predictionset.rst index 89b2757ca..95c968a8a 100644 --- a/docs/api/calib/pyhealth.calib.predictionset.rst +++ b/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, 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 ---------------- diff --git a/docs/api/metrics.rst b/docs/api/metrics.rst index f03533b76..a98543f70 100644 --- a/docs/api/metrics.rst +++ b/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 diff --git a/docs/api/metrics/pyhealth.metrics.prediction_interval.rst b/docs/api/metrics/pyhealth.metrics.prediction_interval.rst new file mode 100644 index 000000000..06b40a58a --- /dev/null +++ b/docs/api/metrics/pyhealth.metrics.prediction_interval.rst @@ -0,0 +1,5 @@ +[core] prediction_interval +=================================== + +.. automodule:: pyhealth.metrics.prediction_interval + :members: diff --git a/examples/conformal_timeseries/enbpi_synthetic.py b/examples/conformal_timeseries/enbpi_synthetic.py new file mode 100644 index 000000000..f4710e9c0 --- /dev/null +++ b/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() diff --git a/pyhealth/calib/README.md b/pyhealth/calib/README.md index 5bc86751e..d96ba17c4 100644 --- a/pyhealth/calib/README.md +++ b/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. \ No newline at end of file diff --git a/pyhealth/calib/base_classes.py b/pyhealth/calib/base_classes.py index b68897b9e..d3aee4922 100644 --- a/pyhealth/calib/base_classes.py +++ b/pyhealth/calib/base_classes.py @@ -36,4 +36,27 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: def to(self, device): super().to(device) self.device = device - return self \ No newline at end of file + 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``. + """ + + 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 diff --git a/pyhealth/calib/predictionset/__init__.py b/pyhealth/calib/predictionset/__init__.py index 667830cf7..c1e444bd8 100644 --- a/pyhealth/calib/predictionset/__init__.py +++ b/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", ] diff --git a/pyhealth/calib/predictionset/timeseries/__init__.py b/pyhealth/calib/predictionset/timeseries/__init__.py new file mode 100644 index 000000000..11cd4a66c --- /dev/null +++ b/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"] diff --git a/pyhealth/calib/predictionset/timeseries/enbpi.py b/pyhealth/calib/predictionset/timeseries/enbpi.py new file mode 100644 index 000000000..23a2b54f3 --- /dev/null +++ b/pyhealth/calib/predictionset/timeseries/enbpi.py @@ -0,0 +1,426 @@ +""" +Ensemble Batch Prediction Intervals (EnbPI) + +Online conformal prediction intervals for time series regression. A bootstrap +ensemble gives every training point an out-of-bag (OOB) prediction from the +models that never saw it, and the absolute residuals of those predictions fill +a sliding window. The interval half-width is a quantile of that window, which +slides forward as test labels are revealed, without refitting any model. + +Paper: + Xu, Chen, and Yao Xie. "Conformal prediction interval for dynamic + time-series." ICML 2021. +""" + +import bisect +import logging +from collections import deque +from typing import Callable, Dict, Optional, Tuple + +import numpy as np +import torch +import torch.nn.functional as F + +from pyhealth.calib.base_classes import IntervalPredictor +from pyhealth.datasets import get_dataloader +from pyhealth.models import BaseModel +from pyhealth.trainer import Trainer + +from .utils import as_column, as_numpy, conformal_quantile + +logger = logging.getLogger(__name__) + +__all__ = ["EnbPI"] + +SUPPORTED_AGGREGATIONS = ("mean", "median") +SUPPORTED_CENTERS = ("aggregate", "quantile") + +_AGGREGATE_FNS = {"mean": np.nanmean, "median": np.nanmedian} + +# Cap on the elements of the 3-D array the median aggregation builds per chunk, +# about 32 MB of float64. +_MAX_CHUNK_ELEMENTS = 2**22 + + +def _oob_aggregate( + oob: np.ndarray, + preds: np.ndarray, + aggregation: str, + max_chunk_elements: int = _MAX_CHUNK_ELEMENTS, +) -> np.ndarray: + """Aggregates each OOB ensemble's predictions at a batch of inputs. + + Args: + oob: Boolean mask of shape (B, T'). Column ``i`` marks the models that + left training point ``i`` out, and has at least one True. + preds: Predictions of every model, shape (B, n). + aggregation: ``"mean"`` or ``"median"``. + max_chunk_elements: Cap on the size of the 3-D array built for the + median. + + Returns: + Array of shape (T', n), the prediction ``f_{-i}(x)`` of each OOB + ensemble ``i`` at each input ``x``. + """ + if aggregation == "mean": + # A masked mean is a matrix product, with no (T', B, n) array. einsum + # avoids the spurious floating-point warnings some BLAS builds raise + # from matmul. + sums = np.einsum("bt,bn->tn", oob, preds.astype(float)) + return sums / oob.sum(axis=0)[:, None] + # The median has no such shortcut, so bound memory by chunking over T'. + n_models, n_points = oob.shape + chunk = max(1, max_chunk_elements // (n_models * preds.shape[1])) + out = np.empty((n_points, preds.shape[1])) + for start in range(0, n_points, chunk): + mask = oob[:, start : start + chunk].T[:, :, None] + out[start : start + chunk] = np.nanmedian( + np.where(mask, preds[None, :, :], np.nan), axis=1 + ) + return out + + +def _check_regression(model: BaseModel) -> None: + if model.mode != "regression": + raise NotImplementedError( + f"EnbPI only supports regression models, got mode {model.mode!r}." + ) + + +class EnbPI(IntervalPredictor): + """Ensemble Batch Prediction Intervals for time series regression. + + EnbPI trains ``n_bootstrap`` models on bootstrap resamples of the + training data. Each training point ``i`` gets an out-of-bag prediction + ``f_{-i}(x_i)``, the aggregate of the models whose resample left ``i`` + out (the paper's leave-one-out predictor), and its absolute residual + enters a sliding window of length ``T``, the number of training points. + At test time the interval is + + ``f_{-t}(x_t) +/- (1 - alpha) quantile of the window`` + + where ``f_{-t}(x_t)`` aggregates the OOB ensembles' predictions at + ``x_t``. Each time ``feedback_size`` test labels are revealed, their + residuals enter the window and the oldest leave it. No model is refit. + + Coverage is approximate rather than finite-sample: + ``P(Y_t in C_t) >= 1 - alpha - o(1)`` when the errors are stationary and + strongly mixing and the OOB predictors approximate the true regression + function. Exchangeability is not required. + + Paper: + Xu, Chen, and Yao Xie. "Conformal prediction interval for dynamic + time-series." ICML 2021. + + Args: + model_fn: Returns a fresh, untrained regression ``BaseModel``. Called + once per bootstrap model, so each starts from its own random + initialization. + alpha: Target miscoverage level in (0, 1). + n_bootstrap: Number of bootstrap models ``B``. The paper finds 20 to + 30 enough. Default 20. + aggregation: How predictions are aggregated, ``"mean"`` (default, as + in the paper's experiments) or ``"median"``. + center: ``"aggregate"`` (default) uses ``aggregation`` over the OOB + ensembles' predictions at ``x_t``. ``"quantile"`` uses their + ``(1 - alpha)`` quantile, as written in Algorithm 1, line 12. + batch_size: Batch size for training the bootstrap models and for + running them over a dataset. Default 32. + train_kwargs: Keyword arguments for ``Trainer.train``, such as + ``epochs``. Default is None, which uses the Trainer defaults. + device: Device for training. Default is None, which uses the Trainer + default. + feedback_size: Points predicted before their labels are revealed in + ``predict_sequential``, the ``s`` of the paper. Default 1. + random_state: Seed for drawing the bootstrap resamples. Training also + uses torch's global RNG, so seed it with ``torch.manual_seed`` for + fully reproducible models. + + Examples: + >>> from pyhealth.calib.predictionset import EnbPI + >>> from pyhealth.models import MLP + >>> # train_ds and test_ds hold one series in time order, with a + >>> # "regression" label "y" and lagged values as the feature "x". + >>> enbpi = EnbPI( + ... model_fn=lambda: MLP(dataset=train_ds, feature_keys=["x"], + ... label_key="y", mode="regression"), + ... alpha=0.1, n_bootstrap=20, train_kwargs={"epochs": 10}, + ... random_state=0, + ... ) + >>> enbpi.calibrate(train_ds) # trains 20 models + >>> out = enbpi.predict_sequential(test_ds) + >>> from pyhealth.metrics import regression_metrics_fn + >>> # A dict with the MAE of the centers, the share of test points + >>> # outside their interval (near alpha) and the mean width. + >>> scores = regression_metrics_fn( + ... out["y_true"], out["y_pred"], + ... metrics=["mae", "miscoverage_pi", "interval_width"], + ... y_lower=out["y_lower"], y_upper=out["y_upper"], + ... ) + """ + + def __init__( + self, + model_fn: Callable[[], BaseModel], + alpha: float, + n_bootstrap: int = 20, + aggregation: str = "mean", + center: str = "aggregate", + batch_size: int = 32, + train_kwargs: Optional[Dict] = None, + device: Optional[str] = None, + feedback_size: int = 1, + random_state: Optional[int] = None, + ) -> None: + if not 0 < alpha < 1: + raise ValueError(f"alpha must be in (0, 1), got {alpha}.") + if n_bootstrap < 1: + raise ValueError(f"n_bootstrap must be at least 1, got {n_bootstrap}.") + if aggregation not in SUPPORTED_AGGREGATIONS: + raise ValueError( + f"Unknown aggregation: {aggregation!r}. Supported: " + f"{SUPPORTED_AGGREGATIONS}." + ) + if center not in SUPPORTED_CENTERS: + raise ValueError( + f"Unknown center: {center!r}. Supported: {SUPPORTED_CENTERS}." + ) + if feedback_size < 1: + raise ValueError(f"feedback_size must be at least 1, got {feedback_size}.") + super().__init__() + self.mode = "regression" + self.model_fn = model_fn + self.alpha = alpha + self.n_bootstrap = n_bootstrap + self.aggregation = aggregation + self.aggregate = _AGGREGATE_FNS[aggregation] + self.center = center + self.batch_size = batch_size + self.train_kwargs = train_kwargs or {} + self.train_device = device + self.feedback_size = feedback_size + self.rng = np.random.default_rng(random_state) + + self.models: Optional[torch.nn.ModuleList] = None + self.bootstrap_indices: Optional[np.ndarray] = None + # oob[b, i] is True when model b never saw training point i. + self.oob: Optional[np.ndarray] = None + self.residuals: Optional[deque] = None + # The same residuals in ascending order, so the width is a lookup. + self.sorted_residuals: Optional[list] = None + self.label_key: Optional[str] = None + self.calibrated = False + + def _train_models(self, train_dataset) -> None: + """Algorithm 1, lines 1-4: fit one model per bootstrap resample.""" + n_samples = len(train_dataset) + self.bootstrap_indices = self.rng.integers( + 0, n_samples, size=(self.n_bootstrap, n_samples) + ) + models = [] + for b, indices in enumerate(self.bootstrap_indices): + logger.info(f"EnbPI: training bootstrap model {b + 1}/{self.n_bootstrap}") + model = self.model_fn() + # Fail before training the first model, not after all of them. + _check_regression(model) + # The loader shuffles, so sorting keeps the same bag while letting + # disk-backed datasets read contiguous runs. + loader = get_dataloader( + train_dataset.subset(np.sort(indices).tolist()), + self.batch_size, + shuffle=True, + ) + Trainer(model=model, enable_logging=False, device=self.train_device).train( + train_dataloader=loader, **self.train_kwargs + ) + models.append(model) + self.models = torch.nn.ModuleList(models) + + def _model_predictions(self, **batch) -> np.ndarray: + """Predictions of every bootstrap model, shape (n_bootstrap, n).""" + preds = [] + with torch.no_grad(): + for model in self.models: + model.eval() + y_prob = model(**batch)["y_prob"] + preds.append(y_prob.detach().cpu().numpy().reshape(-1)) + return np.stack(preds) + + def _push_residuals(self, y_pred: np.ndarray, y_true: np.ndarray) -> None: + """Appends absolute residuals, skipping points with no label or OOB + prediction (NaN), so they neither enter nor evict anything.""" + for residual in np.abs(np.asarray(y_true) - np.asarray(y_pred)): + if np.isnan(residual): + continue + if len(self.residuals) == self.residuals.maxlen: + oldest = self.residuals[0] + del self.sorted_residuals[ + bisect.bisect_left(self.sorted_residuals, oldest) + ] + self.residuals.append(float(residual)) + bisect.insort(self.sorted_residuals, float(residual)) + + def calibrate(self, cal_dataset) -> None: + """Trains the ensemble and fills the window with OOB residuals. + + Args: + cal_dataset: Training samples. EnbPI uses no separate calibration + split, so this is the data its models are trained on. Must + be in time order. + + The ensemble is trained on the first call only. Calling + ``calibrate`` again on the same data reuses it and resets the window + to the OOB residuals, for example before evaluating another + ``alpha``. To train on different data, create a new ``EnbPI``. + """ + if self.models is None: + self._train_models(cal_dataset) + for model in self.models: + _check_regression(model) + self.label_key = self.models[0].label_keys[0] + n_samples = len(cal_dataset) + if self.bootstrap_indices.shape[1] != n_samples: + raise ValueError( + f"The ensemble was trained on {self.bootstrap_indices.shape[1]} " + f"samples, but cal_dataset has {n_samples}. Create a new EnbPI " + "to train on different data." + ) + in_bag = np.zeros((len(self.models), n_samples), dtype=bool) + np.put_along_axis(in_bag, self.bootstrap_indices, True, axis=1) + self.oob = ~in_bag + + preds, y_true = [], [] + loader = get_dataloader(cal_dataset, self.batch_size, shuffle=False) + for batch in loader: + preds.append(self._model_predictions(**batch)) + y_true.append(as_numpy(batch[self.label_key])) + preds = np.concatenate(preds, axis=1) + y_true = np.concatenate(y_true) + + # Algorithm 1, lines 6-10: OOB prediction from the models that left + # each point out. + has_oob = self.oob.any(axis=0) + if not has_oob.all(): + logger.warning( + f"EnbPI: {int((~has_oob).sum())} training points are in every " + "bootstrap sample and get no residual. Increase n_bootstrap." + ) + y_oob = np.full(n_samples, np.nan) + y_oob[has_oob] = self.aggregate( + np.where(self.oob, preds, np.nan)[:, has_oob], axis=0 + ) + self.residuals = deque(maxlen=n_samples) + self.sorted_residuals = [] + self._push_residuals(y_oob, y_true) + self.calibrated = True + + def _check_calibrated(self) -> None: + if not self.calibrated: + raise RuntimeError( + "Model must be calibrated before inference. Call calibrate() first." + ) + + def _center(self, **batch) -> np.ndarray: + """Algorithm 1, line 12: center from the OOB ensembles at ``x_t``.""" + preds = self._model_predictions(**batch) + # OOB prediction of every training point's ensemble at each test + # input, shape (training points with an OOB model, batch size). + oob = self.oob[:, self.oob.any(axis=0)] + oob_preds = _oob_aggregate(oob, preds, self.aggregation) + if self.center == "quantile": + return np.quantile(oob_preds, 1 - self.alpha, axis=0) + return self.aggregate(oob_preds, axis=0) + + def _width(self) -> float: + """Algorithm 1, line 13: the ``(1 - alpha)`` quantile of the window.""" + return conformal_quantile(self.sorted_residuals, self.alpha, presorted=True) + + def _predict_dataset(self, dataset) -> Tuple[np.ndarray, np.ndarray]: + """Centers and labels for a dataset, in dataset order.""" + loader = get_dataloader(dataset, self.batch_size, shuffle=False) + y_pred, y_true = [], [] + for batch in loader: + y_pred.append(self._center(**batch)) + y_true.append(as_numpy(batch[self.label_key])) + return np.concatenate(y_pred), np.concatenate(y_true) + + def forward(self, **kwargs) -> Dict[str, torch.Tensor]: + """Predicts intervals for one batch without changing the window. + + Works inside ``Trainer.inference``. Labels in the batch are not used + to update the window; call :meth:`update` once they are revealed. + Running ``Trainer.inference`` over a whole test set therefore gives + every point the same width, which is not the paper's online method. + Use :meth:`predict_sequential` to evaluate that. + + Returns: + A dictionary with ``y_prob`` (the center), ``y_lower`` and + ``y_upper``, each of shape [batch size, 1]. When the batch carries + labels it also has ``y_true`` and the squared error ``loss`` of + ``y_prob``, as ``Trainer.inference`` expects. + """ + self._check_calibrated() + y_pred = self._center(**kwargs) + width = self._width() + out = { + "y_prob": as_column(y_pred), + "y_lower": as_column(y_pred - width), + "y_upper": as_column(y_pred + width), + } + if self.label_key in kwargs: + out["y_true"] = as_column(as_numpy(kwargs[self.label_key])) + out["loss"] = F.mse_loss(out["y_prob"], out["y_true"]) + return out + + def update(self, revealed_dataset) -> None: + """Algorithm 1, lines 15-19: slides the window by new residuals. + + Args: + revealed_dataset: Points whose labels were just revealed, in time + order. Their residuals enter the window and as many of the + oldest leave it. + """ + self._check_calibrated() + self._push_residuals(*self._predict_dataset(revealed_dataset)) + + def predict_sequential(self, test_dataset) -> Dict[str, np.ndarray]: + """Evaluates EnbPI on a labeled test stream, as in the paper. + + A convenience method for evaluation, not part of the + :class:`~pyhealth.calib.base_classes.IntervalPredictor` interface. It + runs Algorithm 1, lines 11-19: predict ``feedback_size`` points from + the current window, then reveal their labels and slide the window, + until the stream ends. This gives the same intervals as calling + :meth:`forward` on each chunk of ``feedback_size`` points followed by + :meth:`update` on that chunk, but runs the ensemble over the stream + only once, since the centers do not depend on the window. + + All labels must be known in advance, so this is for measuring + coverage and width on held-out data. In deployment, where labels + arrive later, call :meth:`forward` and :meth:`update` directly. + + The window is changed in place: afterwards it holds the latest test + residuals, so a second call continues from there rather than starting + over. + + Args: + test_dataset: Test samples following the training data, in time + order. + + Returns: + A dictionary of arrays of shape (n,) in time order: ``y_true``, + ``y_pred`` (the center), ``y_lower`` and ``y_upper``. + """ + self._check_calibrated() + y_pred, y_true = self._predict_dataset(test_dataset) + width = np.empty(len(y_true)) + for start in range(0, len(y_true), self.feedback_size): + step = slice(start, start + self.feedback_size) + width[step] = self._width() + self._push_residuals(y_pred[step], y_true[step]) + return { + "y_true": y_true, + "y_pred": y_pred, + "y_lower": y_pred - width, + "y_upper": y_pred + width, + } diff --git a/pyhealth/calib/predictionset/timeseries/utils.py b/pyhealth/calib/predictionset/timeseries/utils.py new file mode 100644 index 000000000..342cd14d9 --- /dev/null +++ b/pyhealth/calib/predictionset/timeseries/utils.py @@ -0,0 +1,66 @@ +"""Helpers shared by the time series conformal interval methods. + +Plain functions, like ``prepare_numpy_dataset`` for the set predictors. Each +method uses the ones it needs and implements its own online loop. +""" + +import math + +import numpy as np +import torch + +__all__ = ["conformal_quantile"] + + +def conformal_quantile(scores, alpha: float, presorted: bool = False) -> float: + """Conformal ``(1 - alpha)`` quantile of a window of scores. + + Returns the ``ceil((1 - alpha) * (n + 1))``-th smallest score, the same + convention as split conformal prediction. Adaptive methods move ``alpha`` + outside ``[0, 1]``, so the edge cases are part of the contract. + + Args: + scores: Scores of shape (n,). Higher means less conforming. + alpha: Miscoverage level. Any real value is accepted. + presorted: Whether ``scores`` is already sorted in ascending order. + It is then indexed directly, without copying, which lets online + methods keep a sorted window and read its quantile in constant + time. Any sequence works, such as a list. Default False. + + Returns: + The quantile. ``+inf`` when the rank exceeds ``n``, which covers + ``alpha <= 0`` and windows too small for ``alpha``. ``-inf`` when the + rank is below 1, which covers ``alpha >= 1``. + + Examples: + >>> import numpy as np + >>> conformal_quantile(np.arange(1.0, 10.0), alpha=0.2) + 8.0 + >>> conformal_quantile(np.arange(1.0, 10.0), alpha=0.0) + inf + """ + if not presorted: + scores = np.asarray(scores, dtype=float).ravel() + n = len(scores) + # The small tolerance stops float error from bumping an exact rank up by + # one, e.g. (1 - 0.1) * 10 evaluates to 9.000000000000002. + rank = math.ceil((1 - alpha) * (n + 1) - 1e-9) + if rank > n: + return math.inf + if rank < 1: + return -math.inf + if presorted: + return float(scores[rank - 1]) + return float(np.partition(scores, rank - 1)[rank - 1]) + + +def as_numpy(values) -> np.ndarray: + """Flattens a label tensor or array to a float array of shape (n,).""" + if isinstance(values, torch.Tensor): + values = values.detach().cpu().numpy() + return np.asarray(values, dtype=float).reshape(-1) + + +def as_column(values: np.ndarray) -> torch.Tensor: + """Converts an array of shape (n,) to a float32 tensor of shape (n, 1).""" + return torch.as_tensor(np.asarray(values), dtype=torch.float32).reshape(-1, 1) diff --git a/pyhealth/metrics/prediction_interval.py b/pyhealth/metrics/prediction_interval.py new file mode 100644 index 000000000..2b2b1cd9b --- /dev/null +++ b/pyhealth/metrics/prediction_interval.py @@ -0,0 +1,88 @@ +"""Metrics for regression prediction intervals. + +These helpers back the prediction interval metrics of +:func:`pyhealth.metrics.regression_metrics_fn`, the way +:mod:`pyhealth.metrics.prediction_set` backs the prediction set metrics of +the classification metrics functions. +""" + +import numpy as np + + +def _flatten(**arrays: np.ndarray) -> list: + """Flattens each array to float and checks they all have the same length.""" + flat = {name: np.asarray(a, dtype=float).ravel() for name, a in arrays.items()} + if len({a.shape for a in flat.values()}) > 1: + shapes = ", ".join(f"{name} {a.shape}" for name, a in flat.items()) + raise ValueError(f"Inputs must have the same number of values, got {shapes}.") + return list(flat.values()) + + +def _covered( + y_lower: np.ndarray, y_upper: np.ndarray, y_true: np.ndarray +) -> np.ndarray: + """Whether each target lies in its interval, bounds included.""" + y_lower, y_upper, y_true = _flatten(y_lower=y_lower, y_upper=y_upper, y_true=y_true) + return (y_lower <= y_true) & (y_true <= y_upper) + + +def width(y_lower: np.ndarray, y_upper: np.ndarray) -> float: + """Average width of the prediction intervals. + + Infinite if any interval is unbounded. + + Example: + >>> width(np.asarray([0.0, 1.0]), np.asarray([2.0, 2.0])) + 1.5 + """ + y_lower, y_upper = _flatten(y_lower=y_lower, y_upper=y_upper) + return float(np.mean(y_upper - y_lower)) + + +def miscoverage_pi( + y_lower: np.ndarray, y_upper: np.ndarray, y_true: np.ndarray +) -> float: + """Miscoverage rate, Prob(Y not in [y_lower, y_upper]). + + Example: + >>> y_lower = np.asarray([0.0, 2.0, 0.0, 1.0]) + >>> y_upper = np.asarray([2.0, 3.0, 4.0, 3.0]) + >>> y_true = np.asarray([1.0, 2.0, 5.0, 3.0]) + >>> miscoverage_pi(y_lower, y_upper, y_true) + 0.25 + + Explanation: + 1 is in [0, 2], and 2 and 3 sit on a bound of their intervals, which + counts as covered. 5 is above [0, 4]. Thus the miscoverage rate is 1/4. + """ + return float(1 - np.mean(_covered(y_lower, y_upper, y_true))) + + +def rolling_miscoverage_pi( + y_lower: np.ndarray, + y_upper: np.ndarray, + y_true: np.ndarray, + window: int, +) -> np.ndarray: + """Miscoverage rate over each run of ``window`` consecutive time steps. + + Entry ``k`` covers steps ``k`` to ``k + window - 1``, so the result has + ``len(y_true) - window + 1`` entries. The samples must be in time order. + Overall miscoverage averages over the whole stream, so it hides periods of + undercoverage after a distribution shift that show up here. + + Example: + >>> y_lower = np.zeros(4) + >>> y_upper = np.full(4, 2.0) + >>> y_true = np.asarray([1.0, 9.0, 1.0, 1.0]) + >>> rolling_miscoverage_pi(y_lower, y_upper, y_true, window=2) + array([0.5, 0.5, 0. ]) + + Explanation: + Entry k covers steps k and k + 1. Only step 1 (9 is above [0, 2]) is + missed, so the windows starting at steps 0 and 1 have miscoverage 1/2. + """ + missed = (~_covered(y_lower, y_upper, y_true)).astype(float) + if not 1 <= window <= len(missed): + raise ValueError(f"window must be in [1, {len(missed)}], got {window}.") + return np.convolve(missed, np.ones(window) / window, mode="valid") diff --git a/pyhealth/metrics/regression.py b/pyhealth/metrics/regression.py index 2d0d02ed5..382a37c34 100644 --- a/pyhealth/metrics/regression.py +++ b/pyhealth/metrics/regression.py @@ -4,11 +4,15 @@ import numpy as np import sklearn.metrics as sklearn_metrics +import pyhealth.metrics.prediction_interval as pint + def regression_metrics_fn( x: np.ndarray, x_rec: np.ndarray, metrics: Optional[List[str]] = None, + y_lower: Optional[np.ndarray] = None, + y_upper: Optional[np.ndarray] = None, ) -> Dict[str, float]: """Computes metrics for regression. @@ -17,6 +21,12 @@ def regression_metrics_fn( - kl_divergence: KL divergence - mse: mean squared error - mae: mean absolute error + + The following metrics related to the prediction intervals are accepted as well, but will be ignored if y_lower or y_upper is None: + - miscoverage_pi: Prob(Y not in prediction interval). Check :func:`pyhealth.metrics.prediction_interval.miscoverage_pi`. + - interval_width: Average width of the prediction intervals. Check :func:`pyhealth.metrics.prediction_interval.width`. + - rolling_miscoverage_pi@W: miscoverage_pi over each run of W consecutive time steps, for any positive integer W (e.g. "rolling_miscoverage_pi@100"). Returns an array and needs the samples in time order. Check :func:`pyhealth.metrics.prediction_interval.rolling_miscoverage_pi`. + If no metrics are specified, kd_div, mse, mae are computed by default. This function calls sklearn.metrics functions to compute the metrics. For @@ -27,6 +37,12 @@ def regression_metrics_fn( x: True target data sample. x_rec: reconstructed data sample. metrics: List of metrics to compute. Default is ["kl_divergence", "mse", "mae"]. + y_lower: Optional lower bounds of the prediction intervals, one per + sample. Only used by the prediction interval metrics listed above. + Default is None. + y_upper: Optional upper bounds of the prediction intervals, one per + sample. Only used by the prediction interval metrics listed above. + Default is None. Returns: Dictionary of metrics whose keys are the metric names and values are @@ -48,9 +64,26 @@ def regression_metrics_fn( if x.shape != x_rec.shape: raise ValueError("x and x_rec must have the same shape.") + prediction_interval_metrics = ["miscoverage_pi", "interval_width"] + has_intervals = y_lower is not None and y_upper is not None + output = {} for metric in metrics: - if metric == "kl_divergence": + if metric in prediction_interval_metrics or metric.startswith( + "rolling_miscoverage_pi@" + ): + if not has_intervals: + continue + if metric == "miscoverage_pi": + output[metric] = pint.miscoverage_pi(y_lower, y_upper, x) + elif metric == "interval_width": + output[metric] = pint.width(y_lower, y_upper) + else: + window = _parse_window(metric) + output[metric] = pint.rolling_miscoverage_pi( + y_lower, y_upper, x, window + ) + elif metric == "kl_divergence": # Work on copies to avoid mutating x/x_rec for subsequent metrics x_kl = x.copy() x_rec_kl = x_rec.copy() @@ -71,6 +104,17 @@ def regression_metrics_fn( return output +def _parse_window(metric: str) -> int: + """Parses ``W`` from a metric name such as ``"rolling_miscoverage_pi@100"``.""" + window = metric.split("@", 1)[1] + if not window.isdigit() or int(window) < 1: + raise ValueError( + f"Invalid metric {metric!r}: expected 'rolling_miscoverage_pi@W' with " + "W a positive integer (e.g. 'rolling_miscoverage_pi@100')." + ) + return int(window) + + def timeseries_regression_metrics_fn( y_true: Sequence[np.ndarray], diff --git a/tests/core/test_enbpi.py b/tests/core/test_enbpi.py new file mode 100644 index 000000000..85540372a --- /dev/null +++ b/tests/core/test_enbpi.py @@ -0,0 +1,272 @@ +"""Tests for EnbPI (Xu and Xie 2021). + +The deterministic tests use hand-computed values from the paper's +Algorithm 1 with constant fake models and fixed bootstrap samples. + +Fixture: three models predict the constants 1, 2 and 4, trained on + + b0: [0, 0, 1, 1] -> out of bag for points {2, 3} + b1: [2, 2, 2, 2] -> out of bag for points {0, 1, 3} + b2: [0, 1, 2, 2] -> out of bag for point {3} + +so the OOB predictions are 2, 2, 1 and aggregate(1, 2, 4) for point 3, which +is 7/3 under the mean and 2 under the median. With targets [0, 3, 1, 5] the +mean residuals are [2, 1, 0, 8/3]. +""" + +import unittest + +import numpy as np +import torch + +from pyhealth.calib.predictionset import EnbPI +from pyhealth.calib.predictionset.timeseries import conformal_quantile +from pyhealth.calib.predictionset.timeseries.enbpi import _oob_aggregate +from pyhealth.datasets import get_dataloader +from pyhealth.metrics.prediction_interval import miscoverage_pi +from pyhealth.models import MLP +from tests.core.timeseries_cp_test_utils import ( + ConstantModel, + make_ar_stream, + make_stream, +) + +# Model outputs and forward() results are float32, so compare at float32 precision. +RTOL = 1e-6 + +BAGS = np.array([[0, 0, 1, 1], [2, 2, 2, 2], [0, 1, 2, 2]]) +Y_TRAIN = [0.0, 3.0, 1.0, 5.0] + + +def _fitted_enbpi(bags=BAGS, models=None, **kwargs): + """EnbPI with fixed fake models and bags, so calibrate skips training.""" + enbpi = EnbPI(model_fn=lambda: None, **kwargs) + if models is None: + models = [ConstantModel(1.0), ConstantModel(2.0), ConstantModel(4.0)] + enbpi.models = torch.nn.ModuleList(models) + enbpi.bootstrap_indices = np.asarray(bags) + return enbpi + + +class TestOutOfBag(unittest.TestCase): + def test_mean_residuals(self): + enbpi = _fitted_enbpi(alpha=0.5) + enbpi.calibrate(make_stream(Y_TRAIN)) + np.testing.assert_allclose( + np.asarray(enbpi.residuals), [2.0, 1.0, 0.0, 8 / 3], rtol=RTOL + ) + + def test_median_residuals(self): + enbpi = _fitted_enbpi(alpha=0.5, aggregation="median") + enbpi.calibrate(make_stream(Y_TRAIN)) + np.testing.assert_allclose( + np.asarray(enbpi.residuals), [2.0, 1.0, 0.0, 3.0], rtol=RTOL + ) + + def test_point_in_every_bag_is_dropped(self): + bags = np.array([[0, 0, 3, 1], [2, 2, 3, 3], [0, 1, 2, 3]]) + enbpi = _fitted_enbpi(bags=bags, alpha=0.5) + with self.assertLogs( + "pyhealth.calib.predictionset.timeseries.enbpi", "WARNING" + ): + enbpi.calibrate(make_stream(Y_TRAIN)) + # Point 3 has no OOB model; points 0, 1 use b1 (2) and point 2 uses b0 (1). + np.testing.assert_allclose( + np.asarray(enbpi.residuals), [2.0, 1.0, 0.0], rtol=RTOL + ) + self.assertEqual(enbpi.residuals.maxlen, 4) + + +class TestInterval(unittest.TestCase): + def setUp(self): + self.test = make_stream([1.0, 9.0, 2.0]) + self.batch = next(iter(get_dataloader(self.test, 3, shuffle=False))) + + def test_aggregate_center_and_width(self): + # OOB ensembles predict [2, 2, 1, 7/3] everywhere, mean 11/6. + # Window [2, 1, 0, 8/3], alpha 0.5: rank ceil(0.5 * 5) = 3 -> 2. + enbpi = _fitted_enbpi(alpha=0.5) + enbpi.calibrate(make_stream(Y_TRAIN)) + out = enbpi(**self.batch) + np.testing.assert_allclose( + out["y_prob"].numpy().ravel(), [11 / 6] * 3, rtol=RTOL + ) + np.testing.assert_allclose( + out["y_lower"].numpy().ravel(), [11 / 6 - 2] * 3, rtol=RTOL + ) + np.testing.assert_allclose( + out["y_upper"].numpy().ravel(), [11 / 6 + 2] * 3, rtol=RTOL + ) + + def test_quantile_center(self): + # Algorithm 1 line 12: np.quantile([2, 2, 1, 7/3], 0.5) = 2. + enbpi = _fitted_enbpi(alpha=0.5, center="quantile") + enbpi.calibrate(make_stream(Y_TRAIN)) + np.testing.assert_allclose( + enbpi(**self.batch)["y_prob"].numpy().ravel(), [2.0] * 3, rtol=RTOL + ) + + def test_window_slides_after_each_feedback_batch(self): + # Center 11/6; test residuals |1 - 11/6| = 5/6, |9 - 11/6| = 43/6. + enbpi = _fitted_enbpi(alpha=0.5, feedback_size=2) + enbpi.calibrate(make_stream(Y_TRAIN)) + out = enbpi.predict_sequential(self.test) + # The third residual is |2 - 11/6| = 1/6; the window keeps T = 4. + np.testing.assert_allclose( + np.asarray(enbpi.residuals), [8 / 3, 5 / 6, 43 / 6, 1 / 6], rtol=RTOL + ) + # The first two points share the calibration window, quantile 2. + # The third uses [0, 8/3, 5/6, 43/6]: rank 3 of 4 -> 8/3. + np.testing.assert_allclose( + out["y_upper"] - out["y_pred"], [2.0, 2.0, 8 / 3], rtol=RTOL + ) + + def test_forward_does_not_change_window(self): + enbpi = _fitted_enbpi(alpha=0.5) + enbpi.calibrate(make_stream(Y_TRAIN)) + enbpi(**self.batch) + np.testing.assert_allclose( + np.asarray(enbpi.residuals), [2.0, 1.0, 0.0, 8 / 3], rtol=RTOL + ) + + def test_forward_returns_columns_and_loss(self): + enbpi = _fitted_enbpi(alpha=0.5) + enbpi.calibrate(make_stream(Y_TRAIN)) + out = enbpi(**self.batch) + for key in ("y_prob", "y_lower", "y_upper", "y_true"): + self.assertEqual(out[key].shape, (3, 1)) + self.assertIn("loss", out) + + def test_update_slides_window_in_order(self): + # Residuals |1 - 11/6| = 5/6 then |9 - 11/6| = 43/6 evict 2 and 1. + enbpi = _fitted_enbpi(alpha=0.5) + enbpi.calibrate(make_stream(Y_TRAIN)) + enbpi.update(make_stream([1.0, 9.0])) + np.testing.assert_allclose( + np.asarray(enbpi.residuals), [0.0, 8 / 3, 5 / 6, 43 / 6], rtol=RTOL + ) + + def test_feedback_size_one_updates_every_step(self): + # Widths: 2 from [2, 1, 0, 8/3]; then [1, 0, 8/3, 5/6] -> 1; + # then [0, 8/3, 5/6, 43/6] -> 8/3. + enbpi = _fitted_enbpi(alpha=0.5, feedback_size=1) + enbpi.calibrate(make_stream(Y_TRAIN)) + out = enbpi.predict_sequential(self.test) + np.testing.assert_allclose( + out["y_upper"] - out["y_pred"], [2.0, 1.0, 8 / 3], rtol=RTOL + ) + + +class TestTimeOrder(unittest.TestCase): + def test_dataset_order_is_time_order(self): + enbpi = _fitted_enbpi(alpha=0.5) + enbpi.calibrate(make_stream(Y_TRAIN)) + out = enbpi.predict_sequential(make_stream([3.0, 1.0, 2.0])) + np.testing.assert_array_equal(out["y_true"], [3.0, 1.0, 2.0]) + + +class TestScaling(unittest.TestCase): + """The memory- and time-saving paths match the direct computation.""" + + def test_oob_aggregate_matches_masked_reference(self): + rng = np.random.default_rng(0) + oob = rng.random((6, 40)) < 0.4 + oob[0, ~oob.any(axis=0)] = True # every column needs an OOB model + preds = rng.normal(size=(6, 5)) + masked = np.where(oob.T[:, :, None], preds[None, :, :], np.nan) + for aggregation, reference in (("mean", np.nanmean), ("median", np.nanmedian)): + with self.subTest(aggregation=aggregation): + # A tiny chunk cap forces the median through several chunks. + np.testing.assert_allclose( + _oob_aggregate(oob, preds, aggregation, max_chunk_elements=50), + reference(masked, axis=1), + ) + + def test_sorted_window_matches_conformal_quantile(self): + enbpi = _fitted_enbpi(alpha=0.2) + enbpi.calibrate(make_stream(Y_TRAIN)) + rng = np.random.default_rng(0) + # Integer residuals give ties, which removal from the sorted copy + # must handle. + for residual in rng.integers(0, 5, size=30).astype(float): + enbpi._push_residuals(np.zeros(1), np.array([residual])) + window = np.asarray(enbpi.residuals) + self.assertEqual(enbpi.sorted_residuals, sorted(window.tolist())) + self.assertEqual(enbpi._width(), conformal_quantile(window, 0.2)) + + +class TestInterface(unittest.TestCase): + def test_forward_before_calibrate_raises(self): + batch = next(iter(get_dataloader(make_stream(Y_TRAIN), 4, shuffle=False))) + with self.assertRaises(RuntimeError): + _fitted_enbpi(alpha=0.5)(**batch) + + def test_non_regression_model_fn_raises_before_training(self): + enbpi = EnbPI(model_fn=lambda: ConstantModel(1.0, mode="binary"), alpha=0.1) + with self.assertRaises(NotImplementedError): + enbpi.calibrate(make_stream(Y_TRAIN)) + + def test_recalibrate_on_different_length_raises(self): + enbpi = _fitted_enbpi(alpha=0.5) + enbpi.calibrate(make_stream(Y_TRAIN)) + with self.assertRaises(ValueError): + enbpi.calibrate(make_stream(Y_TRAIN + [2.0])) + + def test_recalibrate_resets_window(self): + enbpi = _fitted_enbpi(alpha=0.5) + enbpi.calibrate(make_stream(Y_TRAIN)) + enbpi.update(make_stream([1.0, 9.0])) + enbpi.calibrate(make_stream(Y_TRAIN)) + np.testing.assert_allclose( + np.asarray(enbpi.residuals), [2.0, 1.0, 0.0, 8 / 3], rtol=RTOL + ) + + def test_non_regression_model_raises(self): + enbpi = _fitted_enbpi( + bags=BAGS[:1], models=[ConstantModel(1.0, mode="binary")], alpha=0.1 + ) + with self.assertRaises(NotImplementedError): + enbpi.calibrate(make_stream(Y_TRAIN)) + + def test_invalid_arguments_raise(self): + for kwargs in ( + {"alpha": 0.0}, + {"alpha": 0.1, "aggregation": "trimmed"}, + {"alpha": 0.1, "center": "mode"}, + {"alpha": 0.1, "n_bootstrap": 0}, + {"alpha": 0.1, "feedback_size": 0}, + ): + with self.subTest(**kwargs), self.assertRaises(ValueError): + EnbPI(model_fn=lambda: None, **kwargs) + + +class TestEndToEnd(unittest.TestCase): + def test_coverage_on_ar_series(self): + torch.manual_seed(42) + stream = make_ar_stream(n=300, seed=42) + train, test = stream.subset(range(100)), stream.subset(range(100, 300)) + enbpi = EnbPI( + model_fn=lambda: MLP( + dataset=stream, + feature_keys=["x"], + label_key="y", + mode="regression", + hidden_dim=4, + ), + alpha=0.1, + n_bootstrap=5, + train_kwargs={"epochs": 5}, + device="cpu", + random_state=42, + ) + enbpi.calibrate(train) + out = enbpi.predict_sequential(test) + self.assertTrue(np.all(out["y_lower"] <= out["y_pred"])) + self.assertTrue(np.all(out["y_pred"] <= out["y_upper"])) + miscoverage = miscoverage_pi(out["y_lower"], out["y_upper"], out["y_true"]) + self.assertGreater(miscoverage, 0.02) + self.assertLess(miscoverage, 0.2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_prediction_interval_metrics.py b/tests/core/test_prediction_interval_metrics.py new file mode 100644 index 000000000..66828b305 --- /dev/null +++ b/tests/core/test_prediction_interval_metrics.py @@ -0,0 +1,110 @@ +"""Tests for the prediction interval metrics and their regression_metrics_fn hook.""" + +import unittest + +import numpy as np + +import pyhealth.metrics.prediction_interval as pint +from pyhealth.metrics import regression_metrics_fn + + +class TestPredictionIntervalHelpers(unittest.TestCase): + def setUp(self): + self.y_true = np.array([1.0, 2.0, 5.0, 0.0]) + self.y_lower = np.array([0.0, 2.0, 0.0, 1.0]) + self.y_upper = np.array([2.0, 3.0, 4.0, 3.0]) + + def test_miscoverage_counts_bounds_as_covered(self): + # Covered: 1 in [0, 2], 2 in [2, 3]. Missed: 5 > 4, 0 < 1. + self.assertEqual( + pint.miscoverage_pi(self.y_lower, self.y_upper, self.y_true), 0.5 + ) + + def test_width(self): + # Widths [2, 1, 4, 2]. + self.assertEqual(pint.width(self.y_lower, self.y_upper), 2.25) + + def test_infinite_width(self): + self.assertEqual(pint.width(np.array([-np.inf]), np.array([1.0])), np.inf) + + def test_rolling_miscoverage(self): + np.testing.assert_array_equal( + pint.rolling_miscoverage_pi( + self.y_lower, self.y_upper, self.y_true, window=2 + ), + [0.0, 0.5, 1.0], + ) + + def test_column_inputs_are_flattened(self): + self.assertEqual( + pint.miscoverage_pi( + self.y_lower[:, None], self.y_upper[:, None], self.y_true[:, None] + ), + 0.5, + ) + + def test_shape_mismatch_raises(self): + with self.assertRaises(ValueError): + pint.miscoverage_pi(self.y_lower[:3], self.y_upper, self.y_true) + with self.assertRaises(ValueError): + pint.width(self.y_lower[:1], self.y_upper) + + def test_invalid_window_raises(self): + with self.assertRaises(ValueError): + pint.rolling_miscoverage_pi( + self.y_lower, self.y_upper, self.y_true, window=5 + ) + + +class TestRegressionMetricsFn(unittest.TestCase): + def setUp(self): + self.y_true = np.array([1.0, 2.0, 5.0, 0.0]) + self.y_pred = np.array([1.0, 2.5, 2.0, 2.0]) + self.y_lower = np.array([0.0, 2.0, 0.0, 1.0]) + self.y_upper = np.array([2.0, 3.0, 4.0, 3.0]) + + def test_interval_metrics(self): + out = regression_metrics_fn( + self.y_true, + self.y_pred, + metrics=["mae", "miscoverage_pi", "interval_width"], + y_lower=self.y_lower, + y_upper=self.y_upper, + ) + self.assertEqual(set(out), {"mae", "miscoverage_pi", "interval_width"}) + self.assertEqual(out["miscoverage_pi"], 0.5) + self.assertEqual(out["interval_width"], 2.25) + + def test_rolling_metric_uses_window_from_name(self): + out = regression_metrics_fn( + self.y_true, + self.y_pred, + metrics=["rolling_miscoverage_pi@3"], + y_lower=self.y_lower, + y_upper=self.y_upper, + ) + np.testing.assert_allclose(out["rolling_miscoverage_pi@3"], [1 / 3, 2 / 3]) + + def test_interval_metrics_ignored_without_bounds(self): + out = regression_metrics_fn( + self.y_true, + self.y_pred, + metrics=["mae", "miscoverage_pi", "rolling_miscoverage_pi@2"], + y_lower=self.y_lower, + ) + self.assertEqual(set(out), {"mae"}) + + def test_invalid_window_name_raises(self): + for metric in ("rolling_miscoverage_pi@0", "rolling_miscoverage_pi@x"): + with self.subTest(metric=metric), self.assertRaises(ValueError): + regression_metrics_fn( + self.y_true, + self.y_pred, + metrics=[metric], + y_lower=self.y_lower, + y_upper=self.y_upper, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_timeseries_cp_utils.py b/tests/core/test_timeseries_cp_utils.py new file mode 100644 index 000000000..d8614220b --- /dev/null +++ b/tests/core/test_timeseries_cp_utils.py @@ -0,0 +1,52 @@ +"""Tests for the shared time series conformal helpers. + +The quantile tests use hand-computed values. +""" + +import math +import unittest + +import numpy as np + +from pyhealth.calib.predictionset.timeseries import conformal_quantile + + +class TestConformalQuantile(unittest.TestCase): + def test_rank_convention(self): + # n = 9, alpha = 0.2: rank ceil(0.8 * 10) = 8, so the 8th smallest. + self.assertEqual(conformal_quantile(np.arange(1.0, 10.0), 0.2), 8.0) + + def test_exact_rank_is_not_bumped_by_float_error(self): + # (1 - 0.1) * 10 is 9.000000000000002 in floating point; rank is 9. + self.assertEqual(conformal_quantile(np.arange(1.0, 10.0), 0.1), 9.0) + + def test_presorted_list_matches_unsorted(self): + scores = [5.0, 1.0, 4.0, 2.0, 3.0] + for alpha in (-0.5, 0.0, 0.3, 0.6, 1.0): + with self.subTest(alpha=alpha): + self.assertEqual( + conformal_quantile(sorted(scores), alpha, presorted=True), + conformal_quantile(np.array(scores), alpha), + ) + + def test_unordered_input(self): + self.assertEqual(conformal_quantile(np.array([3.0, 1.0, 2.0]), 0.5), 2.0) + + def test_window_too_small_gives_inf(self): + # n = 4, alpha = 0.1: rank ceil(0.9 * 5) = 5 > 4. + self.assertEqual(conformal_quantile(np.ones(4), 0.1), math.inf) + + def test_empty_window_gives_inf(self): + self.assertEqual(conformal_quantile(np.array([]), 0.1), math.inf) + + def test_alpha_at_or_below_zero_gives_inf(self): + self.assertEqual(conformal_quantile(np.arange(100.0), 0.0), math.inf) + self.assertEqual(conformal_quantile(np.arange(100.0), -0.3), math.inf) + + def test_alpha_at_or_above_one_gives_minus_inf(self): + self.assertEqual(conformal_quantile(np.arange(100.0), 1.0), -math.inf) + self.assertEqual(conformal_quantile(np.arange(100.0), 1.4), -math.inf) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/timeseries_cp_test_utils.py b/tests/core/timeseries_cp_test_utils.py new file mode 100644 index 000000000..cf90b1529 --- /dev/null +++ b/tests/core/timeseries_cp_test_utils.py @@ -0,0 +1,56 @@ +"""Shared fixtures for the time series conformal prediction tests. + +Holds tiny regression streams and fake regression models with known outputs, +for the EnbPI tests now and the ACI and AgACI tests later. It defines no test +cases, so unittest discovery does not collect it. +""" + +from typing import List, Optional + +import numpy as np +import torch + +from pyhealth.datasets import create_sample_dataset + + +def make_stream( + y: List[float], + x: Optional[List[float]] = None, +): + """Regression dataset with one sample per time step. + + Each sample has a 1-d tensor feature ``x`` and a ``regression`` label + ``y``, in time order. + """ + x = list(y) if x is None else x + samples = [ + {"patient_id": "p0", "x": [float(xi)], "y": float(yi)} for xi, yi in zip(x, y) + ] + return create_sample_dataset(samples, {"x": "tensor"}, {"y": "regression"}) + + +def make_ar_stream(n: int, phi: float = 0.8, lags: int = 2, seed: int = 0): + """AR(1) series with unit Gaussian noise, ``lags`` past values as features.""" + rng = np.random.default_rng(seed) + y = np.zeros(n + lags) + for t in range(1, len(y)): + y[t] = phi * y[t - 1] + rng.normal() + samples = [ + {"patient_id": "p0", "x": y[t - lags : t].tolist(), "y": float(y[t])} + for t in range(lags, len(y)) + ] + return create_sample_dataset(samples, {"x": "tensor"}, {"y": "regression"}) + + +class ConstantModel(torch.nn.Module): + """Fake regression model that predicts ``value`` for every input.""" + + def __init__(self, value: float, mode: str = "regression") -> None: + super().__init__() + self.value = value + self.mode = mode + self.label_keys = ["y"] + + def forward(self, **batch): + n = len(batch["y"]) + return {"y_prob": torch.full((n, 1), float(self.value))} From 72624488279db5a2bd98e881c59b59967c2b9ee3 Mon Sep 17 00:00:00 2001 From: Felipe Bonchristiano Date: Thu, 29 Oct 2026 22:05:34 -0500 Subject: [PATCH 2/5] fix PR rule violations: modern type hints and docstring examples --- pyhealth/calib/base_classes.py | 39 +++++++++++-------- .../calib/predictionset/timeseries/enbpi.py | 26 ++++++------- .../calib/predictionset/timeseries/utils.py | 14 ++++++- pyhealth/metrics/regression.py | 4 +- 4 files changed, 50 insertions(+), 33 deletions(-) diff --git a/pyhealth/calib/base_classes.py b/pyhealth/calib/base_classes.py index d3aee4922..497ca69da 100644 --- a/pyhealth/calib/base_classes.py +++ b/pyhealth/calib/base_classes.py @@ -22,15 +22,28 @@ def to(self, device): return self -class SetPredictor(ABC, torch.nn.Module): - def __init__(self, model, **kwargs) -> None: - super().__init__() - self.model = model +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 forward(self, **kwargs) -> dict[str, torch.Tensor]: + ... + + def update(self, revealed_dataset): ... def to(self, device): @@ -39,13 +52,10 @@ 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``. - """ +class SetPredictor(ABC, torch.nn.Module): + def __init__(self, model, **kwargs) -> None: + super().__init__() + self.model = model def calibrate(self, cal_dataset): ... @@ -53,10 +63,7 @@ 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 + return self \ No newline at end of file diff --git a/pyhealth/calib/predictionset/timeseries/enbpi.py b/pyhealth/calib/predictionset/timeseries/enbpi.py index 23a2b54f3..ad6fd641e 100644 --- a/pyhealth/calib/predictionset/timeseries/enbpi.py +++ b/pyhealth/calib/predictionset/timeseries/enbpi.py @@ -15,7 +15,7 @@ import bisect import logging from collections import deque -from typing import Callable, Dict, Optional, Tuple +from collections.abc import Callable import numpy as np import torch @@ -167,10 +167,10 @@ def __init__( aggregation: str = "mean", center: str = "aggregate", batch_size: int = 32, - train_kwargs: Optional[Dict] = None, - device: Optional[str] = None, + train_kwargs: dict | None = None, + device: str | None = None, feedback_size: int = 1, - random_state: Optional[int] = None, + random_state: int | None = None, ) -> None: if not 0 < alpha < 1: raise ValueError(f"alpha must be in (0, 1), got {alpha}.") @@ -201,14 +201,14 @@ def __init__( self.feedback_size = feedback_size self.rng = np.random.default_rng(random_state) - self.models: Optional[torch.nn.ModuleList] = None - self.bootstrap_indices: Optional[np.ndarray] = None + self.models: torch.nn.ModuleList | None = None + self.bootstrap_indices: np.ndarray | None = None # oob[b, i] is True when model b never saw training point i. - self.oob: Optional[np.ndarray] = None - self.residuals: Optional[deque] = None + self.oob: np.ndarray | None = None + self.residuals: deque | None = None # The same residuals in ascending order, so the width is a lookup. - self.sorted_residuals: Optional[list] = None - self.label_key: Optional[str] = None + self.sorted_residuals: list | None = None + self.label_key: str | None = None self.calibrated = False def _train_models(self, train_dataset) -> None: @@ -335,7 +335,7 @@ def _width(self) -> float: """Algorithm 1, line 13: the ``(1 - alpha)`` quantile of the window.""" return conformal_quantile(self.sorted_residuals, self.alpha, presorted=True) - def _predict_dataset(self, dataset) -> Tuple[np.ndarray, np.ndarray]: + def _predict_dataset(self, dataset) -> tuple[np.ndarray, np.ndarray]: """Centers and labels for a dataset, in dataset order.""" loader = get_dataloader(dataset, self.batch_size, shuffle=False) y_pred, y_true = [], [] @@ -344,7 +344,7 @@ def _predict_dataset(self, dataset) -> Tuple[np.ndarray, np.ndarray]: y_true.append(as_numpy(batch[self.label_key])) return np.concatenate(y_pred), np.concatenate(y_true) - def forward(self, **kwargs) -> Dict[str, torch.Tensor]: + def forward(self, **kwargs) -> dict[str, torch.Tensor]: """Predicts intervals for one batch without changing the window. Works inside ``Trainer.inference``. Labels in the batch are not used @@ -383,7 +383,7 @@ def update(self, revealed_dataset) -> None: self._check_calibrated() self._push_residuals(*self._predict_dataset(revealed_dataset)) - def predict_sequential(self, test_dataset) -> Dict[str, np.ndarray]: + def predict_sequential(self, test_dataset) -> dict[str, np.ndarray]: """Evaluates EnbPI on a labeled test stream, as in the paper. A convenience method for evaluation, not part of the diff --git a/pyhealth/calib/predictionset/timeseries/utils.py b/pyhealth/calib/predictionset/timeseries/utils.py index 342cd14d9..e3ed9518f 100644 --- a/pyhealth/calib/predictionset/timeseries/utils.py +++ b/pyhealth/calib/predictionset/timeseries/utils.py @@ -55,12 +55,22 @@ def conformal_quantile(scores, alpha: float, presorted: bool = False) -> float: def as_numpy(values) -> np.ndarray: - """Flattens a label tensor or array to a float array of shape (n,).""" + """Flattens a label tensor or array to a float array of shape (n,). + + Examples: + >>> as_numpy(torch.tensor([[1.0], [2.0]])) + array([1., 2.]) + """ if isinstance(values, torch.Tensor): values = values.detach().cpu().numpy() return np.asarray(values, dtype=float).reshape(-1) def as_column(values: np.ndarray) -> torch.Tensor: - """Converts an array of shape (n,) to a float32 tensor of shape (n, 1).""" + """Converts an array of shape (n,) to a float32 tensor of shape (n, 1). + + Examples: + >>> as_column(np.array([1.0, 2.0])).shape + torch.Size([2, 1]) + """ return torch.as_tensor(np.asarray(values), dtype=torch.float32).reshape(-1, 1) diff --git a/pyhealth/metrics/regression.py b/pyhealth/metrics/regression.py index 382a37c34..7221cedfe 100644 --- a/pyhealth/metrics/regression.py +++ b/pyhealth/metrics/regression.py @@ -11,8 +11,8 @@ def regression_metrics_fn( x: np.ndarray, x_rec: np.ndarray, metrics: Optional[List[str]] = None, - y_lower: Optional[np.ndarray] = None, - y_upper: Optional[np.ndarray] = None, + y_lower: np.ndarray | None = None, + y_upper: np.ndarray | None = None, ) -> Dict[str, float]: """Computes metrics for regression. From 77c357abd3c8d80240d46a6fa6045c93ad9d6d23 Mon Sep 17 00:00:00 2001 From: Felipe Bonchristiano Date: Thu, 29 Oct 2026 22:18:43 -0500 Subject: [PATCH 3/5] reject calibration when no out-of-bag models exist --- .../calib/predictionset/timeseries/enbpi.py | 8 +++++++ tests/core/test_enbpi.py | 22 +++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/pyhealth/calib/predictionset/timeseries/enbpi.py b/pyhealth/calib/predictionset/timeseries/enbpi.py index ad6fd641e..2e184ea6f 100644 --- a/pyhealth/calib/predictionset/timeseries/enbpi.py +++ b/pyhealth/calib/predictionset/timeseries/enbpi.py @@ -300,6 +300,14 @@ def calibrate(self, cal_dataset) -> None: # Algorithm 1, lines 6-10: OOB prediction from the models that left # each point out. has_oob = self.oob.any(axis=0) + if not has_oob.any(): + # Every center would be an empty aggregate, so every interval + # would be NaN, and updates could never fill the window. + raise ValueError( + "EnbPI: every training point is in every bootstrap sample, so " + "no out-of-bag predictions exist. Increase n_bootstrap or use " + "more training data." + ) if not has_oob.all(): logger.warning( f"EnbPI: {int((~has_oob).sum())} training points are in every " diff --git a/tests/core/test_enbpi.py b/tests/core/test_enbpi.py index 85540372a..451abec88 100644 --- a/tests/core/test_enbpi.py +++ b/tests/core/test_enbpi.py @@ -76,6 +76,28 @@ def test_point_in_every_bag_is_dropped(self): ) self.assertEqual(enbpi.residuals.maxlen, 4) + def test_no_out_of_bag_model_raises(self): + # Every bag contains every point, so no OOB prediction exists. + bags = np.array([[0, 1, 2, 3], [3, 2, 1, 0], [1, 0, 3, 2]]) + enbpi = _fitted_enbpi(bags=bags, alpha=0.5) + with self.assertRaisesRegex(ValueError, "no out-of-bag"): + enbpi.calibrate(make_stream(Y_TRAIN)) + self.assertFalse(enbpi.calibrated) + + def test_no_out_of_bag_model_from_training_raises(self): + # The reported case: 2 samples, 1 model, and random_state=1 draws + # the bag [0, 1], which contains both points. + enbpi = EnbPI( + model_fn=lambda: ConstantModel(1.0), + alpha=0.5, + n_bootstrap=1, + random_state=1, + train_kwargs={"epochs": 0}, + ) + with self.assertRaisesRegex(ValueError, "no out-of-bag"): + enbpi.calibrate(make_stream([0.0, 1.0])) + self.assertFalse(enbpi.calibrated) + class TestInterval(unittest.TestCase): def setUp(self): From 6f4940b9f78b803d5cd2efe8ad75a84bb8de1896 Mon Sep 17 00:00:00 2001 From: Felipe Bonchristiano Date: Thu, 29 Oct 2026 22:26:21 -0500 Subject: [PATCH 4/5] add reset() and retrain on every calibrate() call --- .../calib/pyhealth.calib.predictionset.rst | 2 +- .../calib/predictionset/timeseries/enbpi.py | 40 ++++++++------- tests/core/test_enbpi.py | 50 +++++++++++++------ 3 files changed, 60 insertions(+), 32 deletions(-) diff --git a/docs/api/calib/pyhealth.calib.predictionset.rst b/docs/api/calib/pyhealth.calib.predictionset.rst index 95c968a8a..cad347fc2 100644 --- a/docs/api/calib/pyhealth.calib.predictionset.rst +++ b/docs/api/calib/pyhealth.calib.predictionset.rst @@ -140,7 +140,7 @@ test stream. Score the results with ``rolling_miscoverage_pi@W``. .. autoclass:: pyhealth.calib.predictionset.EnbPI - :members: calibrate, forward, update, predict_sequential + :members: calibrate, forward, update, reset, predict_sequential :show-inheritance: Interval Interface and Time Series Helpers diff --git a/pyhealth/calib/predictionset/timeseries/enbpi.py b/pyhealth/calib/predictionset/timeseries/enbpi.py index 2e184ea6f..1685a5338 100644 --- a/pyhealth/calib/predictionset/timeseries/enbpi.py +++ b/pyhealth/calib/predictionset/timeseries/enbpi.py @@ -208,6 +208,8 @@ def __init__( self.residuals: deque | None = None # The same residuals in ascending order, so the width is a lookup. self.sorted_residuals: list | None = None + # The window as calibrate left it, restored by reset. + self.initial_residuals: list | None = None self.label_key: str | None = None self.calibrated = False @@ -263,28 +265,20 @@ def _push_residuals(self, y_pred: np.ndarray, y_true: np.ndarray) -> None: def calibrate(self, cal_dataset) -> None: """Trains the ensemble and fills the window with OOB residuals. + Every call trains a new ensemble on ``cal_dataset``, drawing new + bootstrap samples, so a second call gives different models than the + first. To rerun an experiment from the same starting window without + retraining, use :meth:`reset`. + Args: cal_dataset: Training samples. EnbPI uses no separate calibration split, so this is the data its models are trained on. Must be in time order. - - The ensemble is trained on the first call only. Calling - ``calibrate`` again on the same data reuses it and resets the window - to the OOB residuals, for example before evaluating another - ``alpha``. To train on different data, create a new ``EnbPI``. """ - if self.models is None: - self._train_models(cal_dataset) - for model in self.models: - _check_regression(model) + self.calibrated = False + self._train_models(cal_dataset) self.label_key = self.models[0].label_keys[0] n_samples = len(cal_dataset) - if self.bootstrap_indices.shape[1] != n_samples: - raise ValueError( - f"The ensemble was trained on {self.bootstrap_indices.shape[1]} " - f"samples, but cal_dataset has {n_samples}. Create a new EnbPI " - "to train on different data." - ) in_bag = np.zeros((len(self.models), n_samples), dtype=bool) np.put_along_axis(in_bag, self.bootstrap_indices, True, axis=1) self.oob = ~in_bag @@ -320,8 +314,20 @@ def calibrate(self, cal_dataset) -> None: self.residuals = deque(maxlen=n_samples) self.sorted_residuals = [] self._push_residuals(y_oob, y_true) + self.initial_residuals = list(self.residuals) self.calibrated = True + def reset(self) -> None: + """Restores the window to the OOB residuals from :meth:`calibrate`. + + Undoes every :meth:`update` without retraining or needing the + training data, so an evaluation can be rerun from the same start, for + example with another ``feedback_size`` or ``alpha``. + """ + self._check_calibrated() + self.residuals = deque(self.initial_residuals, maxlen=self.residuals.maxlen) + self.sorted_residuals = sorted(self.initial_residuals) + def _check_calibrated(self) -> None: if not self.calibrated: raise RuntimeError( @@ -408,8 +414,8 @@ def predict_sequential(self, test_dataset) -> dict[str, np.ndarray]: arrive later, call :meth:`forward` and :meth:`update` directly. The window is changed in place: afterwards it holds the latest test - residuals, so a second call continues from there rather than starting - over. + residuals, so a second call continues from there. Call :meth:`reset` + first to start over from the calibration window. Args: test_dataset: Test samples following the training data, in time diff --git a/tests/core/test_enbpi.py b/tests/core/test_enbpi.py index 451abec88..585d29609 100644 --- a/tests/core/test_enbpi.py +++ b/tests/core/test_enbpi.py @@ -39,12 +39,21 @@ def _fitted_enbpi(bags=BAGS, models=None, **kwargs): - """EnbPI with fixed fake models and bags, so calibrate skips training.""" + """EnbPI whose training installs fixed fake models and bags. + + ``enbpi.train_calls`` counts how often calibrate trained. + """ enbpi = EnbPI(model_fn=lambda: None, **kwargs) if models is None: models = [ConstantModel(1.0), ConstantModel(2.0), ConstantModel(4.0)] - enbpi.models = torch.nn.ModuleList(models) - enbpi.bootstrap_indices = np.asarray(bags) + enbpi.train_calls = 0 + + def install(train_dataset): + enbpi.models = torch.nn.ModuleList(models) + enbpi.bootstrap_indices = np.asarray(bags) + enbpi.train_calls += 1 + + enbpi._train_models = install return enbpi @@ -228,27 +237,40 @@ def test_non_regression_model_fn_raises_before_training(self): with self.assertRaises(NotImplementedError): enbpi.calibrate(make_stream(Y_TRAIN)) - def test_recalibrate_on_different_length_raises(self): + def test_recalibrate_retrains(self): enbpi = _fitted_enbpi(alpha=0.5) enbpi.calibrate(make_stream(Y_TRAIN)) - with self.assertRaises(ValueError): - enbpi.calibrate(make_stream(Y_TRAIN + [2.0])) + enbpi.update(make_stream([1.0, 9.0])) + enbpi.calibrate(make_stream(Y_TRAIN)) + self.assertEqual(enbpi.train_calls, 2) + np.testing.assert_allclose( + np.asarray(enbpi.residuals), [2.0, 1.0, 0.0, 8 / 3], rtol=RTOL + ) - def test_recalibrate_resets_window(self): + def test_reset_restores_initial_window(self): enbpi = _fitted_enbpi(alpha=0.5) enbpi.calibrate(make_stream(Y_TRAIN)) enbpi.update(make_stream([1.0, 9.0])) - enbpi.calibrate(make_stream(Y_TRAIN)) + enbpi.reset() + self.assertEqual(enbpi.train_calls, 1) np.testing.assert_allclose( np.asarray(enbpi.residuals), [2.0, 1.0, 0.0, 8 / 3], rtol=RTOL ) + self.assertEqual(enbpi.sorted_residuals, sorted(enbpi.residuals)) + self.assertEqual(enbpi.residuals.maxlen, 4) - def test_non_regression_model_raises(self): - enbpi = _fitted_enbpi( - bags=BAGS[:1], models=[ConstantModel(1.0, mode="binary")], alpha=0.1 - ) - with self.assertRaises(NotImplementedError): - enbpi.calibrate(make_stream(Y_TRAIN)) + def test_reset_reruns_give_same_intervals(self): + enbpi = _fitted_enbpi(alpha=0.5) + enbpi.calibrate(make_stream(Y_TRAIN)) + test = make_stream([1.0, 9.0, 2.0]) + first = enbpi.predict_sequential(test) + enbpi.reset() + second = enbpi.predict_sequential(test) + np.testing.assert_array_equal(first["y_upper"], second["y_upper"]) + + def test_reset_before_calibrate_raises(self): + with self.assertRaises(RuntimeError): + _fitted_enbpi(alpha=0.5).reset() def test_invalid_arguments_raise(self): for kwargs in ( From fce756d92f3893912ee63625bec65947710c813b Mon Sep 17 00:00:00 2001 From: Felipe Bonchristiano Date: Thu, 29 Oct 2026 22:33:32 -0500 Subject: [PATCH 5/5] save calibration state in state_dict for checkpointing --- .../calib/predictionset/timeseries/enbpi.py | 70 ++++++++++++++++++- tests/core/test_enbpi.py | 69 ++++++++++++++++++ 2 files changed, 136 insertions(+), 3 deletions(-) diff --git a/pyhealth/calib/predictionset/timeseries/enbpi.py b/pyhealth/calib/predictionset/timeseries/enbpi.py index 1685a5338..ef468ff79 100644 --- a/pyhealth/calib/predictionset/timeseries/enbpi.py +++ b/pyhealth/calib/predictionset/timeseries/enbpi.py @@ -80,6 +80,13 @@ def _oob_aggregate( return out +def _oob_mask(bootstrap_indices: np.ndarray) -> np.ndarray: + """Mask of shape (B, T), True where model ``b`` never saw point ``i``.""" + in_bag = np.zeros(bootstrap_indices.shape, dtype=bool) + np.put_along_axis(in_bag, bootstrap_indices, True, axis=1) + return ~in_bag + + def _check_regression(model: BaseModel) -> None: if model.mode != "regression": raise NotImplementedError( @@ -108,6 +115,12 @@ class EnbPI(IntervalPredictor): strongly mixing and the OOB predictors approximate the true regression function. Exchangeability is not required. + ``state_dict`` saves the whole predictor: the bootstrap models' weights + and the calibration state, including the current residual window. To + restore it, create an ``EnbPI`` with the same ``model_fn`` and + ``n_bootstrap`` and call ``load_state_dict``, or use + ``Trainer.save_ckpt`` and ``Trainer.load_ckpt``. + Paper: Xu, Chen, and Yao Xie. "Conformal prediction interval for dynamic time-series." ICML 2021. @@ -279,9 +292,7 @@ def calibrate(self, cal_dataset) -> None: self._train_models(cal_dataset) self.label_key = self.models[0].label_keys[0] n_samples = len(cal_dataset) - in_bag = np.zeros((len(self.models), n_samples), dtype=bool) - np.put_along_axis(in_bag, self.bootstrap_indices, True, axis=1) - self.oob = ~in_bag + self.oob = _oob_mask(self.bootstrap_indices) preds, y_true = [], [] loader = get_dataloader(cal_dataset, self.batch_size, shuffle=False) @@ -438,3 +449,56 @@ def predict_sequential(self, test_dataset) -> dict[str, np.ndarray]: "y_lower": y_pred - width, "y_upper": y_pred + width, } + + # ------------------------------------------------------------------ + # Checkpointing, not part of Algorithm 1. PyTorch's extra state hooks + # let state_dict and load_state_dict carry the calibration state. + # ------------------------------------------------------------------ + def get_extra_state(self) -> dict: + """Calibration state that ``state_dict`` saves next to the weights. + + Holds the bootstrap indices, the residual window, the initial window + and the label key, as tensors and plain values, so checkpoints load + with ``torch.load(..., weights_only=True)`` as ``Trainer.load_ckpt`` + does. + """ + if not self.calibrated: + return {"calibrated": False} + return { + "calibrated": True, + "label_key": self.label_key, + "bootstrap_indices": torch.as_tensor(self.bootstrap_indices), + "residuals": torch.tensor(list(self.residuals), dtype=torch.float64), + "initial_residuals": torch.tensor( + self.initial_residuals, dtype=torch.float64 + ), + } + + def set_extra_state(self, state: dict) -> None: + """Restores the state saved by :meth:`get_extra_state`. + + ``load_state_dict`` calls this before loading the bootstrap models' + weights, so it first builds ``n_bootstrap`` models with ``model_fn`` + for those weights to load into. + """ + if not state["calibrated"]: + self.calibrated = False + return + bootstrap_indices = state["bootstrap_indices"].numpy() + if len(bootstrap_indices) != self.n_bootstrap: + raise ValueError( + f"The checkpoint has {len(bootstrap_indices)} bootstrap models, " + f"but this EnbPI has n_bootstrap={self.n_bootstrap}." + ) + if self.models is None or len(self.models) != self.n_bootstrap: + self.models = torch.nn.ModuleList( + self.model_fn() for _ in range(self.n_bootstrap) + ) + self.bootstrap_indices = bootstrap_indices + self.oob = _oob_mask(bootstrap_indices) + self.label_key = state["label_key"] + residuals = state["residuals"].tolist() + self.residuals = deque(residuals, maxlen=bootstrap_indices.shape[1]) + self.sorted_residuals = sorted(residuals) + self.initial_residuals = state["initial_residuals"].tolist() + self.calibrated = True diff --git a/tests/core/test_enbpi.py b/tests/core/test_enbpi.py index 585d29609..61bc59ffa 100644 --- a/tests/core/test_enbpi.py +++ b/tests/core/test_enbpi.py @@ -14,6 +14,7 @@ mean residuals are [2, 1, 0, 8/3]. """ +import io import unittest import numpy as np @@ -284,6 +285,74 @@ def test_invalid_arguments_raise(self): EnbPI(model_fn=lambda: None, **kwargs) +class TestCheckpoint(unittest.TestCase): + """state_dict and load_state_dict round-trip the whole predictor.""" + + def setUp(self): + torch.manual_seed(0) + self.stream = make_ar_stream(n=60, seed=0) + self.train = self.stream.subset(range(30)) + self.revealed = self.stream.subset(range(30, 40)) + self.test = self.stream.subset(range(40, 60)) + + def _enbpi(self): + return EnbPI( + model_fn=lambda: MLP( + dataset=self.stream, + feature_keys=["x"], + label_key="y", + mode="regression", + hidden_dim=4, + ), + alpha=0.2, + n_bootstrap=3, + train_kwargs={"epochs": 1}, + device="cpu", + random_state=0, + ) + + def _round_trip(self, enbpi): + buffer = io.BytesIO() + torch.save(enbpi.state_dict(), buffer) + buffer.seek(0) + restored = self._enbpi() + # weights_only=True, as Trainer.load_ckpt loads checkpoints. + restored.load_state_dict(torch.load(buffer, weights_only=True)) + return restored + + def test_round_trip_resumes_online_updates(self): + enbpi = self._enbpi() + enbpi.calibrate(self.train) + enbpi.update(self.revealed) + restored = self._round_trip(enbpi) + self.assertTrue(restored.calibrated) + self.assertEqual(list(restored.residuals), list(enbpi.residuals)) + expected = enbpi.predict_sequential(self.test) + actual = restored.predict_sequential(self.test) + for key in ("y_pred", "y_lower", "y_upper"): + np.testing.assert_array_equal(actual[key], expected[key]) + + def test_round_trip_keeps_initial_window(self): + enbpi = self._enbpi() + enbpi.calibrate(self.train) + initial = list(enbpi.residuals) + enbpi.update(self.revealed) + restored = self._round_trip(enbpi) + restored.reset() + self.assertEqual(list(restored.residuals), initial) + + def test_uncalibrated_round_trip(self): + restored = self._round_trip(self._enbpi()) + self.assertFalse(restored.calibrated) + + def test_n_bootstrap_mismatch_raises(self): + enbpi = self._enbpi() + enbpi.calibrate(self.train) + other = EnbPI(model_fn=enbpi.model_fn, alpha=0.2, n_bootstrap=2) + with self.assertRaises(ValueError): + other.load_state_dict(enbpi.state_dict()) + + class TestEndToEnd(unittest.TestCase): def test_coverage_on_ar_series(self): torch.manual_seed(42)