# 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).