DLWP: Deep Learning Weather Prediction
DLWP is a Python project containing data-processing and model-building tools for predicting the gridded atmosphere using deep convolutional neural networks.
If you use this code or find it useful please cite our publication!
For now, DLWP is not a package that can be installed using
pip or a
setup.py file, so it works like most research code: download (or checkout) and run.
It is assumed that the following are installed using Anaconda Python 3 (Python 2.7 is supported).
- TensorFlow (GPU capable version highly recommended). The
condapackage, while not the recommended installation method, is easy and also installs the required CUDA dependencies. For best performance, follow the instructions for installing from source.
conda install tensorflow-gpu
pip install keras
conda install netCDF4
conda install dask xarray
The following are required only for some of the DLWP features:
- PyTorch: for torch-based deep learning models. Again the GPU-ready version is recommended.
pip install torch torchvision
- scikit-learn: for machine learning pre-processing tools such as Scalers and Imputers
conda install scikit-learn
- scipy: for CFS data interpolation
- pygrib: for raw CFS data processing
pip install pygrib
- cdsapi: for retrieval of ERA5 data
pip install cdsapi
- pyspharm: spherical harmonics transforms for the barotropic model
conda install -c conda-forge pyspharm
DLWP is built as a weather forecasting model that can, should performance improve greatly, “replace” and existing global weather or climate model. Essentially, this means that DLWP uses a deep convolutional neural network to map the state of the atmosphere at one time to the entire state of the atmophere at the next available time. A continuous forecast can then be made by feeding the model’s predicted state back in as inputs, producing indefinite forecasts.
The classes in
DLWP.data provide tools for retrieving and processing raw data from the CFS reanalysis and reforecast and the ERA5 reanalysis. Meanwhile, the
DLWP.model.preprocessing module provides tools for formatting the data for ingestion into the deep learning models. The following examples retrieve and process data from the CFS reanalysis:
The resulting file of predictor data can be ingested into the data generators for the models.
DLWP.model module contains classes for building and training Keras and PyTorch models. The
DLWPNeuralNet class is essentially a wrapper for the simple Keras
Sequential model, adding optional run-time scaling and imputing of data. It implements a few key methods:
build_model: use a custom API to assemble layers in a
Sequentialmodel. Also implements models running on multiple GPUs.
fit: scale the data and fit the model
fit_generator: use the Keras
fit_generatormethod along with a custom data generator (see section below)
predict: predict with the model
predict_timeseries: predict a continuous time series forecast, where the output of one prediction iteration is used as the input for the next
An example of a model built and trained with the DLWP APIs using data generated by the DLWP processing methods, see
DLWP also implements a
DLWPFunctional class which implements the same methods as the
DLWPNeuralNet class but takes as input to
build_model a model assembled using the Keras functional API. For an example of training a functional model, see
Currently, due to a focus on TensorFlow/Keras models, the PyTorch implementation in DLWP is more limited, although still robust. Like the Keras models, it implements a convenient
build_model method to assemble a sequential-like model using the same API parameters as those for
DLWPNeuralNet. Additionally, it also implements a
fit method to automatically iterate through the data and optimizer, again, just like the Keras API.
The PyTorch example,
train_torch.py, is somewhat outdated and uses the spherical convolution library s2cnn. This method has yet to produce good results.
Custom layers and functions
DLWP.custom module contains many custom layers specifically for applying convolutional neural networks to the global weather prediction problem. For example,
PeriodicPadding2D implements periodic boundary conditions for padding data in space prior to applying convolutions. These custom layers are worth a look.
DLWP.model.generators contains several classes for generating data on-the-fly from a netCDF file produced by the DLWP preprocessing methods. These data generators can then be used in conjunction with a DWLP model instance’s
DataGeneratorclass is the simplest generator class. It merely returns batches of data from a file containing “predictors” and “targets” variables already formatted for use in the DLWP model. Due to this simplicity, this is the optimal way to generate data directly from the disk when system memory is not sufficient to load the entire dataset. However, this comes at the cost of generating very large files on disk with redundant data (since the targets are merely a different time shift of the predictors).
SeriesDataGeneratorclass is much more robust and memory efficient. It expects only a single “predictors” variable in the input file and generates predictor-target pairs on the fly for each batch of data. It also has the ability to prescribe external fields such as incoming solar radiation.
SmartDataGeneratoris deprecated in favor of
Advanced forecast tools
DLWP.model module also contains a
TimeSeriesEstimator class. This class can be used to make robust forward forecasts where the data input does not necessarily match the data output of a model. And example usage of this class is in
examples/validate.py, which performs basic routines to validate the forecast skill of DLWP models.
DLWP.util module contains useful utilities, including
load_model for saving and loading DLWP models (and correctly dealing with multi-GPU models).