README.md

metadata

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/ directory of this repository: the implementations of Direct SI, Cascaded SI and EDM, the configurations, the recorded results and the reproduction scripts. Its README gives the installation and all commands. In short, 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')"

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