Skip to main content

PyTorch Skillful Nowcasting GAN Implementation

Project description

Skillful Nowcasting with Deep Generative Model of Radar (DGMR)

Implementation of DeepMind's Skillful Nowcasting GAN Deep Generative Model of Radar (DGMR) (https://arxiv.org/abs/2104.00954) in PyTorch Lightning.

This implementation matches as much as possible the pseudocode released by DeepMind. Each of the components (Sampler, Context conditioning stack, Latent conditioning stack, Discriminator, and Generator) are normal PyTorch modules. As the model training is a bit complicated, the overall architecture is wrapped in PyTorch Lightning.

The default parameters match what is written in the paper.

Installation

Clone the repository, then run

pip install -r requirements.txt
pip install -e .

Alternatively, you can also install through pip install dgmr

Training Data

The open-sourced UK training dataset has been mirrored to HuggingFace Datasets! This should enable training the original architecture on the original data for reproducing the results from the paper. The full dataset is roughly 1TB in size, and unfortunately, streaming the data from HF Datasets doesn't seem to work, so it has to be cached locally. We have added the sample dataset as well though, which can be directly streamed from GCP without costs.

The dataset can be loaded with

from datasets import load_dataset

dataset = load_dataset("openclimatefix/nimrod-uk-1km")

For now, only the sample dataset support streaming in, as its data files are hosted on GCP, not HF, so it can be used with:

from datasets import load_dataset

dataset = load_dataset("openclimatefix/nimrod-uk-1km", "sample", streaming=True)

The authors also used MRMS US precipitation radar data as another comparison. While that dataset was not released, the MRMS data is publicly available, and we have made that data available on HuggingFace Datasets as well here. This dataset is the raw 3500x7000 contiguous US MRMS data for 2016 through May 2022, is a few hundred GBs in size, with sporadic updates to more recent data planned. This dataset is in Zarr format, and can be streamed without caching locally through

from datasets import load_dataset

dataset = load_dataset("openclimatefix/mrms", "default_sequence", streaming=True)

This steams the data with 24 timesteps per example, just like the UK DGMR dataset. To get individual MRMS frames, instead of a sequence, this can be achieved through

from datasets import load_dataset

dataset = load_dataset("openclimatefix/mrms", "default", streaming=True)

Pretrained Weights

Pretrained weights are be available through HuggingFace Hub, currently weights trained on the sample dataset. The whole DGMR model or different components can be loaded as the following:

from dgmr import DGMR, Sampler, Generator, Discriminator, LatentConditioningStack, ContextConditioningStack
model = DGMR.from_pretrained("openclimatefix/dgmr")
sampler = Sampler.from_pretrained("openclimatefix/dgmr-sampler")
discriminator = Discriminator.from_pretrained("openclimatefix/dgmr-discriminator")
latent_stack = LatentConditioningStack.from_pretrained("openclimatefix/dgmr-latent-conditioning-stack")
context_stack = ContextConditioningStack.from_pretrained("openclimatefix/dgmr-context-conditioning-stack")
generator = Generator(conditioning_stack=context_stack, latent_stack=latent_stack, sampler=sampler)

Example Usage

from dgmr import DGMR
model = DGMR(
        forecast_steps=4,
        input_channels=1,
        output_shape=128,
        latent_channels=384,
        context_channels=192,
        num_samples=3,
    )
x = torch.rand((2, 4, 1, 128, 128))
out = model(x)
y = torch.rand((2, 4, 1, 128, 128))
loss = F.mse_loss(y, out)
loss.backward()

Citation

@article{ravuris2021skillful,
  author={Suman Ravuri and Karel Lenc and Matthew Willson and Dmitry Kangin and Remi Lam and Piotr Mirowski and Megan Fitzsimons and Maria Athanassiadou and Sheleem Kashem and Sam Madge and Rachel Prudden Amol Mandhane and Aidan Clark and Andrew Brock and Karen Simonyan and Raia Hadsell and Niall Robinson Ellen Clancy and Alberto Arribas† and Shakir Mohamed},
  title={Skillful Precipitation Nowcasting using Deep Generative Models of Radar},
  journal={Nature},
  volume={597},
  pages={672--677},
  year={2021}
}

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

dgmr-1.3.4.tar.gz (20.4 kB view details)

Uploaded Source

Built Distribution

dgmr-1.3.4-py3-none-any.whl (22.3 kB view details)

Uploaded Python 3

File details

Details for the file dgmr-1.3.4.tar.gz.

File metadata

  • Download URL: dgmr-1.3.4.tar.gz
  • Upload date:
  • Size: 20.4 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.1 CPython/3.9.13

File hashes

Hashes for dgmr-1.3.4.tar.gz
Algorithm Hash digest
SHA256 dd847f3943b2c9be6697abf4c37e5c1298e4714fec767f34160345ff5c6c49bf
MD5 ec5a32d5b146430a59fd4f156c3b4185
BLAKE2b-256 e004ec73d31d21dee3618608f4567bcf55840c5b07366a62972ec54ce40429b0

See more details on using hashes here.

File details

Details for the file dgmr-1.3.4-py3-none-any.whl.

File metadata

  • Download URL: dgmr-1.3.4-py3-none-any.whl
  • Upload date:
  • Size: 22.3 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.1 CPython/3.9.13

File hashes

Hashes for dgmr-1.3.4-py3-none-any.whl
Algorithm Hash digest
SHA256 fff98942380de0a7bc8a81cc350eb02eb3ac2a595515935dc2b94ab82a696795
MD5 56539b61e3691d5c580604af705975fa
BLAKE2b-256 d1097799166cd82451ac73feab990cbd5beb1c749370973e41811ee950eb4d9d

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page