Rollout-calibrated conformal prediction for autoregressive single-cell foundation models.
Project description
rollout-cp-scfm
Rollout-calibrated conformal prediction for autoregressive single-cell foundation models.
When should we trust a cell-state prediction?
rollout-cp-scfmis a distribution-free uncertainty wrapper for autoregressive scFMs (CellTempo, PerturbGen, autoregressive ImmuneWorld), with finite-sample coverage guarantees that explicitly handle the deployment-time rollout regime.
PyPI distribution: rollout-cp-scfm. Python import path: acp
(from acp.calibration import RolloutSplitConformalRegressor); the
distribution / import split mirrors scikit-learn / sklearn.
What this package does
The library implements split conformal prediction with four score functions, all sharing a common API:
| Module | Use case | Reference |
|---|---|---|
split_cp.py |
Marginal CP for one-shot scFM outputs | Vovk 2005, Lei et al. 2018 |
weighted_cp.py |
CP under covariate shift across cohorts | Tibshirani et al. 2019 |
adaptive_cp.py |
Online ACI for streaming time series | Gibbs & Candès 2021 |
rollout_cp.py |
Per-step / cumulative CP under autoregressive rollouts | This work |
The rollout_cp wrapper is what's needed when the scFM is autoregressive
(CellTempo-style next-state generation), because teacher-forced
calibration provably produces narrower bands than the deployment-time
rollout regime requires (see docs/THEOREM_DRAFT.md Corollary 1).
Rollout calibration restores nominal coverage at every horizon t.
Why this exists
Existing CP wrappers for scFMs (e.g. Cilantro-SL for Geneformer) target one-shot prediction, not multi-step rollouts. Existing rollout-CP work (HopCast, CP-Traj, ConformalDAgger, Conformal Language Modeling) targets robotics / forecasting / language, not single-cell foundation models. This package is the bridge: a clean, tested, pip-installable wrapper for autoregressive scFMs.
Quick start
# Once published to PyPI, just:
# pip install rollout-cp-scfm
# Or, to develop / reproduce §4 figures from source:
git clone https://github.com/shenfuqili/acp-scfm
cd acp-scfm
python3.11 -m venv .venv && source .venv/bin/activate
pip install -e ".[dev,repro]" # core + repro + dev tooling
pytest tests/ # 18 tests should pass
python scripts/simulate_autoregressive.py # generates E3 figure
Expected output of the simulation:
linear ε_M ≈ 0.000 cov_TF[20] = 0.736 cov_RO[20] = 0.901
misspec ε_M ≈ 0.085 cov_TF[20] = 0.272 cov_RO[20] = 0.901
The simulation produces figures/E3/E3_sim_coverage_vs_horizon.png,
which is the paper's main experimental figure: rollout calibration
holds nominal coverage 0.90 across 20 horizons in both well-specified
and misspecified linear-Gaussian regimes; teacher-forced calibration
collapses to 0.27 in the misspec regime.
API at a glance
import numpy as np
from acp.calibration.rollout_cp import (
RolloutSplitConformalRegressor,
per_step_coverage,
)
# cal_rollouts: shape (n_cal, T, d) -- recursive M̂^t(x_0) on cal trajectories.
# cal_truths: shape (n_cal, T, d) -- ground-truth states at t=1..T.
cp = RolloutSplitConformalRegressor(alpha=0.1, mode="per_step").fit(
cal_rollouts, cal_truths
)
# At deployment we only have rollouts; CP gives per-coord prediction bands.
lower, upper = cp.predict_interval(test_rollouts)
curve = per_step_coverage(lower, upper, test_truths)
print(f"per-step coverage: {curve.round(3)}") # ~ [0.9, 0.9, 0.9, ...]
mode="cumulative" aggregates the score across the entire trajectory
and produces a single qhat — useful when joint trajectory coverage
P[∀ t: y_t ∈ C_t] is the quantity of interest.
Repository layout
src/acp/
calibration/
split_cp.py marginal CP (classifier + regressor)
weighted_cp.py Tibshirani 2019 weighted CP, weight clipping
adaptive_cp.py Gibbs & Candès 2021 ACI
rollout_cp.py this work — rollout-calibrated split CP
evaluation/
coverage.py empirical-coverage / efficiency metrics
selective_pred.py risk-coverage curves
backbones/
base.py abstract BackboneWrapper interface
scripts/
simulate_autoregressive.py E3 simulation, two regimes
run_E1_sklearn.py E1 sklearn-LR baseline on Norman 2019
tests/ pytest suite for split + rollout CP
docs/
PIVOT_PLAN.md strategic plan (post v1/v2/v3 adversarial review)
THEOREM_DRAFT.md Proposition 1 + Corollary 1 with proofs
RELATED_WORK.md paper related-work section
E1/ sklearn-LR Norman 2019 baseline writeup
E3/ synthetic AR experiment writeup
Experiments
| Goal | Backbone | Status | |
|---|---|---|---|
| E1 marginal coverage | Norman 2019 (290 perturbations) | sklearn LR baseline | done (±0.002 hit at 5 alphas) |
| E1-Geneformer | same, with Geneformer | Geneformer | planned |
| E3-sim | rollout vs teacher-forced | linear M̂ on linear truth | done (figures/E3/) |
| E3-misspec | activate ε_M term | linear M̂ on tanh-perturbed truth | done (figures/E3/) |
| E3-real | autoregressive scFM | CellTempo on scBaseTraj | planned |
Citing
A bibtex entry will be added once the arXiv preprint is up. For now:
@misc{wang2026rolloutcpscfm,
author = {Zhibin Wang},
title = {rollout-cp-scfm: Rollout-calibrated conformal prediction for
autoregressive single-cell foundation models},
year = {2026},
url = {https://github.com/shenfuqili/acp-scfm}
}
License
MIT.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file rollout_cp_scfm-0.1.0.tar.gz.
File metadata
- Download URL: rollout_cp_scfm-0.1.0.tar.gz
- Upload date:
- Size: 20.1 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.12.7
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
56eca33a2c7291a8439582d9ccce1adcf7bdf9208c091e184b9b0d473f73db6e
|
|
| MD5 |
bab6aa599bc54109f70aa7a761bc4a90
|
|
| BLAKE2b-256 |
6796375a0cf941471258b5c506e1e68837e9a72ca08c7ea3ff39fc7540b72a5c
|
File details
Details for the file rollout_cp_scfm-0.1.0-py3-none-any.whl.
File metadata
- Download URL: rollout_cp_scfm-0.1.0-py3-none-any.whl
- Upload date:
- Size: 24.0 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.12.7
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
2a1f4e8ded659fd138a5b6bfadd14a1c7b0a820d580ddc751a877c12c46274b9
|
|
| MD5 |
93653d71895d441eb1aaab7a91c22ed1
|
|
| BLAKE2b-256 |
b00a9e625023907d641843899d571036fa1d5c23f0881e3c9b781ccd1ac684f2
|