Skip to content

Streaming CP implementation to handle large datasets - #31

Draft
radka-j wants to merge 10 commits into
mainfrom
streaming
Draft

radka-j wants to merge 10 commits into
mainfrom
streaming

Conversation

@radka-j

@radka-j radka-j commented Aug 11, 2026

Copy link
Copy Markdown
Member

Contributes to #29

This pull request introduces a streaming calibration interface to the conformal calibrators, allowing calibration data to be processed incrementally in small chunks rather than requiring the entire calibration tensor in memory at once. The Ensemble conformal calibrator is extended with a member-streaming API, enabling ensemble members to be processed one at a time, including specialized handling for both quantile and standard deviation modes. Documentation is updated to describe the new streaming APIs and usage patterns.

Streaming Calibration Interface:

  • Added a streaming calibration API to the base conformal calibrator, including reset, update, and accumulate_score methods, which allow calibration data to be accumulated incrementally and concatenated only when needed. (src/autouq/calibrators/conformal/conformal.py src/autouq/calibrators/conformal/conformal.pyR41-R119)
  • Updated the documentation to explain the new streaming calibration workflow and how it differs from the previous batch-oriented API.
calibrator.reset()
for true_chunk, pred_chunk in calibration_chunks:
    calibrator.update(true_chunk, pred_chunk)
intervals = calibrator.predict(pred_test, alphas=[0.1, 0.2])

Ensemble Calibrator Streaming API:

  • Implemented a member-streaming interface for the Ensemble calibrator, with methods such as reset_stream, update_stream, stream_score, and stream_predict, enabling ensemble members to be processed one at a time and supporting both quantile and std modes efficiently. (src/autouq/calibrators/conformal/ensemble.py [1] [2]
  • Refactored internal logic to support streaming: ensemble summaries and score computations are now incremental, with Welford’s algorithm for std mode and chunked quantile computation for quantile mode. (src/autouq/calibrators/conformal/ensemble.py [1] [2]
  • Updated the ensemble calibrator documentation to describe the streaming member interface and provide usage examples. (docs/ensemble.md docs/ensemble.mdR62-R105)
# Outer loop: one iteration per calibration example (e.g. one forecast
# initialization time). Each example's raw members are only ever held one
# at a time, and are discarded once that example's score has been folded
# into the calibration bank via accumulate_score.
for true_for_this_example, members in calibration_examples:
    calibrator.reset_stream()
    for member in members:                  # one model forward pass each
        calibrator.update_stream(member, true=true_for_this_example)
    score = calibrator.stream_score(chunk_size=2048)   # chunk_size optional
    calibrator.accumulate_score(score)       # feed the base calibration layer

intervals = calibrator.predict(pred_test, alphas=[0.1, 0.2])

fastnet-anemoi (a downstream consumer) pins torch>=2.8,<2.9 for
flash-attn/torch-cluster CUDA wheel compatibility, which is incompatible
with the >=2.9.1 pin here. Verified the full test suite still passes
under torch==2.8.0.
Comment thread pyproject.toml
"jaxtyping>=0.3.3",
"numpy>=2.4.6",
"torch>=2.9.1",
"torch>=2.8",

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change is for compatibility with downstream project that uses this implementation. Revisit whether we are happy with this change before merging.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant