---
license: cc-by-4.0
tags:
  - physics
  - probabilistic-forecasting
  - generative-models
  - pytorch
---

# Conditional Forecasting Models

This repository contains the trained model checkpoints that accompany an anonymous research submission on conditional generative forecasting from coarse-grained observations.

## Models

The five benchmark datasets each have Direct SI, Cascaded SI and EDM checkpoints, plus the Cascaded SI models of the depth sweep:

| Dataset | Direct SI | Cascaded SI (selected depth) | Cascaded SI (other depths) | EDM |
|---|---|---|---|---|
| `euler` | `direct-si` | `cascaded-si-S10` | `cascaded-si-S02`, `S05`, `S20`, `S40` | `edm` |
| `sqg` | `direct-si` | `cascaded-si-S10` | `cascaded-si-S02`, `S05`, `S20`, `S40` | `edm` |
| `het-waves` | `direct-si` | `cascaded-si-S05` | `cascaded-si-S02`, `S10`, `S20`, `S40` | `edm` |
| `nonlinear-kg` | `direct-si` | `cascaded-si-S20` | `cascaded-si-S02`, `S05`, `S10`, `S40` | `edm` |
| `passive-scalar` | `direct-si` | `direct-si` (S = 1) | `cascaded-si-S02`, `S05`, `S10`, `S20`, `S40` | `edm` |

S is the cascade depth. For the passive scalar the selected depth is S = 1, so its Cascaded SI model is the Direct SI model.

Each directory contains:

- `model.pt` — PyTorch checkpoint with the model and EMA weights, the training step and the training configuration
- `config.yaml` — the run configuration used for sampling

Sampling uses the EMA weights. The checkpoints contain only tensors and plain Python values and load with `torch.load(..., weights_only=True)`.

## Datasets

The corresponding datasets are `coarse-grained-forecasting-benchmark/euler`, `sqg`, `het-waves`, `nonlinear-kg` and `passive-scalar`. Each model uses the normalization of its dataset (`data/normalization.npz`).

## Files and usage

The reproduction code is in the [`code/`](https://huggingface.co/coarse-grained-forecasting-benchmark/forecasting-models/tree/main/code) directory of this repository: the implementations of Direct SI, Cascaded SI and EDM, the configurations, the recorded results and the reproduction scripts. Its [README](https://huggingface.co/coarse-grained-forecasting-benchmark/forecasting-models/blob/main/code/README.md) gives the installation and all commands. In short, 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')"

hf download coarse-grained-forecasting-benchmark/forecasting-models --local-dir models
hf download coarse-grained-forecasting-benchmark/euler --repo-type dataset --local-dir data/euler
export EULER_DATA="$PWD/data/euler/data"
export EULER_TRUTH_BANK="$PWD/data/euler/truth_bank"

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
```

The code names the systems `euler`, `sqg`, `waves`, `nlkg` and `passive_scalar` and the methods `direct_si`, `cascaded_si` and `edm`; `--S <depth>` selects a depth-sweep model. `SHA256SUMS` lists the file checksums.

## License

Creative Commons Attribution 4.0 International (CC BY 4.0).