🌧️ DYffCast: Regional Scale Precipitation Nowcasting Using IMERG Satellite Data

DYffusion forecasts a sequence of h snapshots x_t+i1 , x_t+i2 , ... , x_t+h given the initial condition x₀. This is analogous to how standard diffusion models are used to sample from a distribution.

Summary

• PAPER: https://arxiv.org/abs/2412.02723

• This work extends the probabilistic spatio-temporal forecasting framework, DYffusion (from Rose-STL-Lab/dyffusion) to the task of precipitation nowcasting.

• The aim of this study was to forecast IMERG satellite precipitation data up to a 4-hour horizon.

• This work focuses on the South American countries, Colombia, Ecuador, and Perú — countries more susceptible to increased flooding due to climate change and lacking freely available ground-based weather radar data.

• A novel loss function, referred to as LCB, is introduced, combining MSE, MAE, and the LPIPS perceptual score.

• DYffusion trained with LCB (DYffusion_LCB), outperformed five competitor models: two ConvLSTM baselines trained with LCB and binary cross-entropy (BCE) loss respecitvely, a DYffusion model trained with its native L1 loss and the STEPS nowcast method.

• The models were also evaluated on the heavy rain event, Cyclone Yaku. DYffusion_LCB was able to capture the small and large scale features, generating sharp and stable forecasts up to a 2-hour horizon. See below for Cyclone Yaku results:

The evolution of Cyclone Yaku over 4 hours, showing undersampled intervals of 1-hour, beginning at 03:00 UTC on March 9, 2023. From top to bottom: Ground truth IMERG data (1^st row), DYffusion_LCB (2^nd row), DYffusion_L1 (3^rd row), ConvLSTM_LCB (4^th row), ConvLSTM_BCE (5^th row) and STEPS nowcast method implemented in PySteps.

Usage

The DYffusion code is stored in rainnow/src/dyffusion, for simplicity, the structure is inherited from the original DYffusion repo Rose-STL-Lab/dyffusion.

rainnow
├── train.py
├── train_convlstm.py
└── src
    ├── dyffusion
    │   ├── configs
    │   ├── datamodules
    │   ├── diffusion
    │   ├── experiment_types
    │   └── utilities
    ├── models
    │   ├── __pycache__
    │   ├── modules
    │   ├── _base_model.py
    │   ├── conv_lstm.py
    │   ├── pysteps_steps.py
    │   ├── unet_resnet.py
    │   └── unet.py
    ├── utilities
    │   ├── cmaps
    │   ├── instantiators.py
    │   ├── loading.py
    │   └── utils.py
    ├── convlstm_trainer.py
    ├── datasets.py
    ├── interpolator_evaluation.py
    ├── loss.py
    ├── normalise.py
    └── plotting.py

DYffusion Training

This project is entirely config-driven. You can either train an Interpolator (I(φ)) or a Forecaster (F(θ)) network using the train.py script. train.py instantiates the main_config.yaml file (in dyffusion/configs) to set up the training parameters. The main_config.yaml acknowledges the following (override) .yaml files: trainer.yaml, model.yaml, diffusion.yaml, datamodule.yaml, module.yaml, callbacks.yaml and logger.yaml, that override the main_config.yaml attributes on instantiation. These are all declared at the top of the main_config.yaml config:

defaults:
  - _self_
  - trainer: trainer.yaml
  - model: unet_resnet.yaml                           # unet.yaml
  - diffusion: dyffusion.yaml                         # null
  - datamodule: imerg_precipitation.yaml
  - module: forecasting_multi_horizon_dyffusion.yaml  # interpolation.yaml
  - callbacks: callbacks.yaml
  - logger: csv_logger.yaml

.yaml files

• logger.yaml: Specifies the logging used. Currently, it's set to csv_logger but you can add other logging i.e. weights & biases.
• trainer.yaml: Training-related parameters i.e. number of epochs, number of ensemble predictions, gradient accumulation batches and gradient clipping.
• diffusion.yaml: DYffusion related parameters i.e. trained interpolator checkpoint path and forward condition etc.
• datamodule.yaml: Data related parameters i.e. batch size, data splits and normalisation parameters.
• module.yaml: The experiment type. This will either load the InterpolationExperiment() (Interpolator I(φ)) or the MultiHorizonForecastingDYffusion() (Forecastor, F(θ)).
• callbacks.yaml: Sets the learning rate scheduler, checkpointing, and early stopping. Additional callbacks can be added.

Note: The optimizer.yaml is always included, regardless of experiment type. It sets the optimizer i.e. Adam, AdamW or SGD and the associated learning rate etc.

Example Configurations

(1) Interpolation Experiment:

To run an Interpolator I(φ) (experiment) training, set the defaults to:

defaults:
  - _self_
  - trainer: trainer.yaml
  - model: unet_resnet.yaml
  - diffusion: null
  - datamodule: imerg_precipitation.yaml
  - module: interpolation.yaml
  - callbacks: callbacks.yaml
  - logger: csv_logger.yaml

(2) DYffusion Experiment:

To run a Forecastor F(θ) (experiment) training, set the defaults to:

defaults:
  - _self_
  - trainer: trainer.yaml
  - model: unet_resnet.yaml
  - diffusion: dyffusion.yaml 
  - datamodule: imerg_precipitation.yaml
  - module: forecasting_multi_horizon_dyffusion.yaml
  - callbacks: callbacks.yaml
  - logger: csv_logger.yaml

Additional Information

Please see rainnow/notebooks/README.md and rainnow/src/data_prep/README.md for more detail on the notebooks in this repo and the custom IMERG data pipeline respectively.

Installation

1. To clone the GitHub repository, run the following code in your terminal at a chosen directory:

   git clone https://github.com/Dseal95/DYffcast.git

2. Navigate to the cloned project directory and run the env_setup.sh shell script to set up the conda environment:

   cd <git repository DIR>
   source env_setup.sh

This will set up the rainnow conda environment.

Testing

All unit tests can be located in the /tests directory. To run the tests, make sure you are in the rainnnow environment and run the following command:

python -m pytest tests