File size: 4,132 Bytes
41dd7e0 4a6a13e 41dd7e0 4a6a13e 41dd7e0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | # 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).
|