code/README.md

Stabilizing Probabilistic Forecasts via Stochastic Interpolation Cascades

Code for reproducing the experiments accompanying an anonymous research submission on conditional generative forecasting.

We compare Direct Stochastic Interpolation (Direct SI), Cascaded Stochastic Interpolation (Cascaded SI), and Elucidating Diffusion Models (EDM) on five dynamical systems: 2-D Euler, surface quasi-geostrophic flow, heterogeneous waves, nonlinear Klein–Gordon, and passive scalar advection–diffusion.

The repository contains the model implementations (bifurcation/), configurations (configs/), recorded experimental results (records/), and reproduction scripts (reproduce/).

Installation

From the code directory:

uv sync
export CLEAN_FID_CACHE=$HOME/.cache/clean-fid
uv run python -c "from bifurcation.fid import backend as B; B.build_extractor('cpu')"

The dependencies are pinned in uv.lock. The original experiments used different PyTorch builds for some datasets, so minor numerical differences across environments are possible.

Data and checkpoints

The datasets and trained models are available on Hugging Face:

• Euler
• SQG
• Heterogeneous waves
• Nonlinear Klein–Gordon
• Passive scalar
• Model checkpoints

For example, download the Euler dataset and the checkpoints:

hf download coarse-grained-forecasting-benchmark/euler \
    --repo-type dataset --local-dir data/euler

hf download coarse-grained-forecasting-benchmark/forecasting-models \
    --local-dir models

export EULER_DATA="$PWD/data/euler/data"
export EULER_TRUTH_BANK="$PWD/data/euler/truth_bank"

The other systems follow the same directory structure. Their variables are SQG_DATA, WAVES_DATA (het-waves), NLKG_DATA (nonlinear-kg) and PASSIVE_SCALAR_DATA, each with the matching _TRUTH_BANK variable; every command also accepts --data and --truth-bank directly. hf is the Hugging Face CLI (pip install -U huggingface_hub, or run it without installing as uvx --from huggingface_hub hf ...).

Reproducing the results

To display the recorded results of Tables 2, 5, and 6:

uv run python -m reproduce.tables

To generate 16 conditional forecasts for one Euler test condition using Cascaded SI:

uv run python -m reproduce.table2 generate \
    --system euler --method cascaded_si \
    --checkpoint models/euler/cascaded-si-S10/model.pt \
    --out runs/example --gids 36 --n-samples 16 --device cpu

To reproduce a Table 2 entry, generating 10,000 forecasts per test condition and computing csFID-10k:

uv run python -m reproduce.table2 run \
    --system euler --method cascaded_si \
    --checkpoint models/euler/cascaded-si-S10/model.pt \
    --out runs/euler_cascaded_si --device cuda

Use --method direct_si or --method edm for the baselines. For Cascaded SI, --S selects a specific cascade depth.

The scripts in reproduce/figures/ and reproduce/appendix/ generate the paper's figures and additional analyses from the recorded results, datasets, or forecast samples, as applicable.

Training

Models can also be trained using the configurations of the selected experiments:

uv run python -m bifurcation.train \
    --config records/euler/runs/cascaded_si_S10/config.yaml \
    --data "$EULER_DATA" --out runs/euler_S10

The training configurations, model hyperparameters, and evaluation settings are provided in records/ and configs/.

License

MIT License.

The EDM implementation follows Karras et al. (2022). Image evaluation uses Clean-FID.