Skip to content

feat: run forecasters as ensembles in the inference step - #287

Open
frazane wants to merge 8 commits into
feat/output-variablesfrom
feat/ensemble-inference
Open

frazane wants to merge 8 commits into
feat/output-variablesfrom
feat/ensemble-inference

Conversation

@frazane

@frazane frazane commented Nov 5, 2026 •

Copy link
Copy Markdown
Contributor

Adds an optional ensemble block to forecaster run entries, so the inference step produces N members per initialisation. This PR stops at GRIB on disk; ensemble verification comes in a follow-up PR. Stacked on #286: ensembles rely on output_variables to keep storage manageable.

runs:
  - forecaster:
      checkpoint: https://...
      steps: 0/120/6
      ensemble:
        members: 10
        seed: 0   # member m runs with seed + m
      output_variables: [2t, 2d, 10u, 10v, sp, msl, tp, z]

anemoi-inference produces one member per process, so the prepare step writes one config_{NNN}.yaml per member, writing to grib/{NNN}/ with the file names of a deterministic run and tagged with number, numberOfForecastsInEnsemble and dataType=pf. One inference_execute job per init time runs the members on one node, inference_resources.gpu at a time (default min(N, 4), one GPU each), seeding torch with seed + m before anemoi-inference starts. Each member has its own log and .ok file under logs/inference_execute/{run_id}-{init_time}/, a retry only reruns the members that failed, and the rule's .ok means all members exist. Until ensemble verification lands, verification, plots and MEC use member 001 and log a warning. The ensemble block is part of the run-specific hash (run ids without it are unchanged) and is rejected on the forecaster of a temporal downscaler.

Tested with stormy-plateau (anemoi-inference 0.11.1), 2 inits × 4 members on 2 GPUs: the GRIB keys read back correctly with the COSMO definitions, members differ, and a member rerun with the same seed is bit-identical.

Changes: ForecasterConfig and the JSON schema, inference_prepare.py, inference_execute (new inference_execute_members.sh), run_specific_hash, a run_grib_dir helper used by the verification, plot and MEC rules, README, unit tests.

Notes

  • eps goes in grib2_keys, not encoding: the COSMO definitions have no eps key, and anemoi-inference reads it from either to pick the ensemble template for accumulated fields.
  • Existing inference_execute outputs will be rerun once after this change: Snakemake records each rule's shell template and params, and both changed (the rendered command of a deterministic run is identical). To keep existing outputs, remove their records once with snakemake --cleanup-metadata <files>. New runs are not affected. feat: select the variables written to the inference GRIB output #286 has the same effect through the new output_variables param of inference_prepare_forecaster.

@frazane
frazane added this pull request to stack #288 November 5, 2026 15:00
@frazane
frazane force-pushed the feat/ensemble-inference branch 2 times, most recently from 76bd4dc to 68fefec Compare November 6, 2026 09:06
@frazane
frazane force-pushed the feat/ensemble-inference branch from 68fefec to 1d6496d Compare November 9, 2026 17:00
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