$ show openclimatefix/skillful_nowcasting
This repository is a PyTorch Lightning implementation of DeepMind's Skillful Nowcasting GAN (DGMR) for precipitation nowcasting. It provides a modular, well-structured library of neural network components including residual blocks (GBlock, DBlock, LBlock), conditioning stacks (ContextConditioningStack, LatentConditioningStack), a ConvGRU-based sampler, spatial and temporal discriminators, and multiple loss functions (hinge, grid cell, SSIM, etc.). The core model (DGMR) wraps these components into a PyTorch Lightning module with manual optimization for the GAN training loop. It supports HuggingFace Hub integration for pretrained weights and includes training scripts for the UK Nimrod radar dataset. The codebase is well-documented, type-annotated, and includes comprehensive unit tests for individual components, making it suitable both as a library for nowcasting research and as a reference implementation for complex GAN architectures.
The LLM identified 3 files worth highlighting for reuse:
A compact spatial attention module using 1x1 convolutions for query/key/value and custom einsum-based attention (not standard multi-head). It uses a learnable gamma parameter for residual scaling. The implementation is self-contained (only imports torch and einops) and could be extracted and reused in any 2D feature map attention scenario.
A clean ConvGRU cell and wrapper implementation. The cell uses separate spectrally normalized convolutions for read gate, update gate, and output, with sigmoid and ReLU activations. The wrapper unrolls the cell over a sequence of hidden states. It is independent of other project components and can be reused in video prediction, spatiotemporal forecasting, or other sequence modeling tasks.
A collection of well-implemented loss functions useful for image generation and GAN training: GridCellLoss (weighted per-pixel L1 with a customizable weight function), NowcastingLoss (hinge loss wrapper), loss_hinge_disc/gen functions, SSIM/MS-SSIM losses, and a get_loss factory. These are self-contained (import only torch, numpy, pytorch_msssim) and can be extracted and reused in other generative models.