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