| # 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: |
|
|
| ```bash |
| 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](https://huggingface.co/datasets/coarse-grained-forecasting-benchmark/euler) |
| - [SQG](https://huggingface.co/datasets/coarse-grained-forecasting-benchmark/sqg) |
| - [Heterogeneous waves](https://huggingface.co/datasets/coarse-grained-forecasting-benchmark/het-waves) |
| - [Nonlinear Klein–Gordon](https://huggingface.co/datasets/coarse-grained-forecasting-benchmark/nonlinear-kg) |
| - [Passive scalar](https://huggingface.co/datasets/coarse-grained-forecasting-benchmark/passive-scalar) |
| - [Model checkpoints](https://huggingface.co/coarse-grained-forecasting-benchmark/forecasting-models) |
|
|
| For example, download the Euler dataset and the checkpoints: |
|
|
| ```bash |
| 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: |
|
|
| ```bash |
| uv run python -m reproduce.tables |
| ``` |
|
|
| To generate 16 conditional forecasts for one Euler test condition using Cascaded SI: |
|
|
| ```bash |
| 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: |
|
|
| ```bash |
| 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: |
|
|
| ```bash |
| 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)](https://github.com/NVlabs/edm). Image evaluation uses [Clean-FID](https://github.com/GaParmar/clean-fid). |
|
|
