rwiegels/seasonal_afnocast

AFNOCast — Seasonal Precipitation Post-Processing (Blue Nile)

Seasonal AFNOCast (Wiegels et al., 2026)is a deep-learning post-processing model that corrects bias and downscales daily ensemble precipitation forecasts from ECMWF SEAS5 (~0.28°) to a high-resolution 0.1° grid over the Blue Nile Basin (7.1°N–17.6°N, 32.0°E–39.7°E).

A separate model state dict is provided for each of the 12 calendar months of valid time, enabling month-specific bias correction across the full 215-day SEAS5 forecast horizon.

---

Model Architecture

Seasonal AFNOCast is a 3-D convolutional encoder–decoder augmented with an Adaptive Fourier Neural Operator (AFNO) block that operates on the spectral representation of the spatial field.

Input (25/51 ens × 11 time-steps × 38 lat × 28 lon)
    │
    ├── Encoder (2× ReflectionPad3d + Conv3d + ReLU)
    ├── Encoder2 (2× ReflectionPad3d + Conv3d + ReLU)
    │     └── + ensemble std conditioning (same encoder path)
    ├── Bicubic upsampling ×2.79  →  (~106 lat × 78 lon)
    ├── Positional embedding (learnable)
    ├── AFNO block (Adaptive Fourier, 2 blocks, sparsity threshold 0.001)
    ├── Dropout3d (p=0.1, applied after AFNO — dropout_state=1)
    ├── Encoder2 (second pass)
    └── Decoder (2× stride Conv3d + ReLU → n_ens output members)

Output (n_ens × 106 lat × 78 lon)

The forward pass uses torch.vmap over the ensemble dimension so each member is processed independently through process_chain, with the per-pixel ensemble standard deviation injected as an additive conditioning signal.

Hyperparameter	Value
n_feature	10
dropout_state	1 (late dropout, after AFNO)
n_ens (output members)	1 per forward pass
AFNO blocks	2
Upsampling factor	×2.79 (bicubic, antialias)

---

Input / Output Specification

	Input	Output
Variable	Total precipitation	Total precipitation
Source	ECMWF SEAS5	post-processed forecast towards ERA5-Land
Resolution	0.281° (38 × 28 grid points)	0.1° (106 × 78 grid points)
Ensemble	25 members (hindcast), 51 members (forecast)	1 member per model call, after concatenation 25/51 members
Temporal context	11 days (target day ± 5 days)	1 day (target day)
Domain	lat 7.1–17.6°N, lon 32.0–39.7°E	same

---

Per-Month Checkpoints

One model is trained independently for each calendar month of valid time (i.e. the month in which the forecast day falls, not the initialisation month).

Month	File
January	afnocast_best_01.safetensors
February	afnocast_best_02.safetensors
March	afnocast_best_03.safetensors
April	afnocast_best_04.safetensors
May	afnocast_best_05.safetensors
June	afnocast_best_06.safetensors
July	afnocast_best_07.safetensors
August	afnocast_best_08.safetensors
September	afnocast_best_09.safetensors
October	afnocast_best_10.safetensors
November	afnocast_best_11.safetensors
December	afnocast_best_12.safetensors

---

Training Details

Setting	Value
Optimizer	Adam
Learning rate	3 × 10⁻⁴
Epochs	100 (early stopping patience = 10)
Batch size	4

Data splits:

Split	Period
Training	1981–2010
Validation (during training)	2011–2016
Hold-out test	2017–2023

Loss function:

Component	Weight	Description
CRPS	1.0	Continuous Ranked Probability Score across ensemble

---

Usage

Installation

This project is managed with uv. Clone the repository and run:

uv sync

This creates a .venv/ virtual environment and installs all dependencies. Activate it before running anything:

source .venv/bin/activate   # Linux / macOS
# on Windows: .venv\Scripts\activate

Quick inference (Python API)

from seasonal_afnocast.inference.inference import InferenceOrchestrator

orchestrator = InferenceOrchestrator(
    init_date="202505",                  # SEAS5 initialisation date (yyyymm)
    trained_model_dir="weights/",        # directory with afnocast_best_MM.safetensors
    seas_forecast_dir="/path/to/SEAS5_tp_yyyymm.nc/directory",
    output_dir="./output",
    device="cuda",                       # or "cpu"
    batch_size=4,
)

ds = orchestrator.run()
# ds is an xarray.Dataset with dims (timestep, ens_out, latitude, longitude)
# Saved automatically to output_dir/inference_202505.nc

Command-line interface

python src/seasonal_afnocast/inference/inference.py \
    --init_date yyyymm \
    --seas-forecast-dir /path/to/SEAS5/dir \
    --output-dir ./results/inference \

The --trained-model-dir and --config arguments default to the weights/ and configs/ folders at the project root and can be omitted when running from the cloned repository.

What happens internally

1. compute_target_months("2025", "05") derives the 7–8 calendar months covered by the 215-day forecast.
2. For each target month the orchestrator loads afnocast_best_<MM>.safetensors.
3. InferenceDataset reads the raw SEAS5 NetCDF, clips to the Blue Nile domain, filters to that calendar month, and builds 11-step ±5-day context windows (with boundary padding) — mirroring the training sample construction.
4. Outputs are stitched into a single xarray.Dataset with proper latitude/longitude coordinates and saved as NetCDF.

Expected SEAS5 input file

/your/seas5/dir/SEAS5_tp_<yyyymm>.nc

Dimensions: (ens, time, lat, lon) Variable: tp (total precipitation, float32).

---

Repository Structure

├── configs/
│   └── config_afnocast.yaml               # Model and data configuration
├── weights/
│   ├── afnocast_best_01.safetensors       # State dicts, one per calendar month
│   └── ...
├── src/seasonal_afnocast/
│   ├── models/AFNOCast.py                 # Model architecture
│   ├── inference/inference.py             # InferenceOrchestrator
│   ├── dataloader/
│   │   └── dataloader_inference.py        # InferenceDataset + compute_target_months
│   └── utils/config.py                   # Config loading (YAML → OmegaConf)
├── pyproject.toml
└── README.md

---

Intended Use & Limitations

• Intended use: Post-processing and downscaling of ECMWF SEAS5 ensemble precipitation forecasts over the Blue Nile Basin for seasonal hydro-meteorological applications.
• Geographic scope: Blue Nile Basin only (7.1°N–17.6°N, 32.0°E–39.7°E). The model has not been evaluated outside this domain.
• Variable: Precipitation only. The architecture would need retraining for other variables.
• Forecast system: Trained on SEAS5 reforecasts. Performance for other forecast systems is untested.

---

Citation

If you use AFNOCast in your work, please cite this repository:

@misc{afnocast2025,
  author       = {Wiegels, R., Chwala, C., Polz, J., Glawion, L., Lorenz, C., Schober, T.C., Kunstmann, H. },
  title        = {AFNOCast: Seasonal Precipitation AI Post-Processing for the Blue Nile Basin},
  year         = {2025},
  howpublished = {\url{https://huggingface.co/rwiegels/seasonal_afnocast}},
}

---

Acknowledgements

The AFNO block implementation is based on the AFNO-transformer by NVIDIA Research. SEAS5 data: © ECMWF. ERA5-Land data: Copernicus Climate Change Service.