DeepR: Deep Reanalysis.

Global reanalysis downscaling to regional scales by means of deep learning techniques.

Introduction

In the rapidly evolving landscape of climate science and data analysis, the need for high-resolution data has become increasingly evident. Climate researchers and professionals in various fields, from agriculture to disaster management, rely heavily on accurate and detailed data to make informed decisions. However, the existing global reanalysis data, such as ERA5, with its coarse spatial resolution, often falls short in meeting these requirements. In response to this pressing challenge, the DeepR project, led by Antonio Pérez, Mario Santa Cruz, and Javier Diez, was conceived and executed with the aim of downscaling ERA5 data to a finer resolution, thus enabling enhanced accuracy and applicability across a wide range of industries and research domains.

Project Motivation

The ERA5 Global reanalysis data, with its spatial resolution of approximately 0.25 degrees, has proven to be a valuable resource for multiple sectors. Still, its limitations in resolution can hinder precise decision-making and analysis across diverse domains. The primary motivation behind the DeepR project was to bridge this gap by downscaling ERA5 data to a finer resolution of approximately 0.05 degrees, termed as CERRA resolution. This enhancement aimed to unlock the full potential of climate data for improved decision support.

Super-Resolution in Climate Science

The project drew inspiration from the field of image processing and computer vision, specifically the concept of super-resolution. In image processing, super-resolution involves augmenting the resolution or quality of an image, typically generating a high-resolution image from one or more low-resolution iterations. DeepR adapted this concept to climate science, making it a super-resolution task tailored to atmospheric fields.

Data: The Foundation of DeepR

In any data-intensive project, data plays a pivotal role, and DeepR is no exception. The project relies on extensive datasets sourced from the publicly accessible Climate Data Store (CDS), ensuring transparency and open access to valuable climate information.

The data used in this project has been generously provided by our mentors and is used in its raw form without any processing. To download the data from the repository, you can access the european_weather_cloud.py script.

Additionally, we have developed a script to directly download data from the Climate Data Store. You can find this script at climate_data_store.py.

Focus on a specific domain

The project focuses on a specific subdomain within the original domain. In our case, this domain encompasses diverse ecosystems, including mountains, rivers, coastal areas, and more. This simplification helps reduce the dimensionality of the problem while maintaining the diversity necessary for comprehensive research.

The selected domain is shown here:

To achieve this spatial selection of the data, we utilize the data_spatial_selection.py script, which transforms the data into the desired domain.

The process of rewriting the data for the smaller domain aims to expedite data access, enhancing both memory and time efficiency for smoother and faster data handling. Furthermore, this approach empowers users to define their own specific domains and seamlessly retrain the model according to their unique research requirements.

Definition in configuration file

The data configuration section outlines how the project manages and processes the data. This section is divided into three main parts: features_configuration, label_configuration, and split_coverages.

Features Configuration

This part focuses on the configuration of features used in the project.

features_configuration:
  variables:
  - t2m
  data_name: era5
  spatial_resolution: "025deg"
  add_auxiliary:
    time: true
    lsm-low: true
    orog-low: true
    lsm-high: true
    orog-high: true
  spatial_coverage:
    longitude: [-8.35, 6.6]
    latitude: [46.45, 35.50]
  standardization:
    to_do: true
    cache_folder: /PATH/TO/.cache_reanalysis_scales
    method: domain-wise
  data_location: /PATH/TO/features/
  land_mask_location: /PATH/TO/static/land-mask_ERA5.nc
  orography_location: /PATH/TO/static/orography_ERA5.nc

Variables: The variables to be included, such as t2m (2-meter temperature data).
Data Name: The source of the feature data, which is era5.
Spatial Resolution: The spatial resolution used for feature data is 0.25 degrees.
Add Auxiliary Data: Specifies whether auxiliary data is added. In this case, time, lsm-low (low-resolution land-sea mask), orog-low (low-resolution orography), lsm-high (high-resolution land-sea mask), and orog-high (high-resolution orography) are added.
Spatial Coverage: The selected spatial coverage, defined by longitude and latitude ranges.
Standardization: Indicates whether standardization is performed. The to_do flag is set to true, and the standardization method is domain-wise. Other possible methods include pixel-wise and landmask-wise.
Data Location: The directory where feature data is stored.
Land Mask Location: The location of the land-sea mask data for ERA5.
Orography Location: The location of the orography data for ERA5.

Label Configuration

This part focuses on the configuration of labels used in the project.

label_configuration:
  variable: t2m
  data_name: cerra
  spatial_resolution: "005deg"
  spatial_coverage:
    longitude: [-6.85, 5.1]
    latitude: [44.95, 37]
  standardization:
    to_do: true
    cache_folder: /PATH/TO/.cache_reanalysis_scales
    method: domain-wise # pixel-wise, domain-wise, landmask-wise
  data_location: /PATH/TO/labels/
  land_mask_location: /PATH/TO/static/land-mask_CERRA.nc
  orography_location: /PATH/TO/static/orography_CERRA.nc

Variable: The variable used as labels, which is t2m (2-meter temperature data).
Data Name: The source of the label data, which is cerra.
Spatial Resolution: The spatial resolution used for label data is 0.05 degrees.
Spatial Coverage: The selected spatial coverage, defined by longitude and latitude ranges.
Standardization: Indicates whether standardization is performed. The to_do flag is set to true, and the standardization method is domain-wise. Other possible methods include pixel-wise and landmask-wise.
Data Location: The directory where label data is stored.
Land Mask Location: The location of the land-sea mask data for CERRA.
Orography Location: The location of the orography data for CERRA.

Split Coverages

Splitting the data into different time periods for training, validation and test.

split_coverages:
  train:
    start: 1981-01
    end: 2013-12
    frequency: MS
  validation:
    start: 2014-01
    end: 2017-12
    frequency: MS
  test:
    start: 2018-01
    end: 2020-12
    frequency: MS

Train: Data split for training begins from 1981-01 and ends at 2013-12, with a frequency of Monthly (MS).
Validation: Data split for validation starts from 2014-01 and ends at 2017-12, with a frequency of Monthly (MS).
Test: Data split for validation starts from 2018-01 and ends at 2020-12, with a frequency of Monthly (MS).

These configuration settings are crucial for organizing, processing, and standardizing the data used in the project.

Methodology

Standardization

Why is it important?

In the context of deep learning for climatology, standardizing climatological data is a crucial step. Standardization refers to the process of transforming the data to have a mean of zero and a standard deviation of one. This process is vital for several reasons:

Preventing Dominance: Standardization prevents one variable from dominating the learning process. In climate data, variables can have vastly different scales and magnitudes. Without standardization, variables with larger scales could overshadow others, leading to biased model training.
Capturing Complex Patterns: Standardized data allows the deep learning model to effectively capture complex climate patterns across diverse geographical regions. By removing scale differences, the model can focus on extracting meaningful patterns and relationships within the data.
Facilitating Convergence: Deep neural networks benefit from standardized input data. It helps in the convergence of the network during training. When the input data has consistent scales and distributions, the optimization process becomes more stable, and the model is more likely to converge to a meaningful solution.

Application

To apply standardization to climatological data, we use the script located in scaler.py. This script automates the process of standardization, making it easy to preprocess large datasets efficiently.

In summary, standardizing climatological data is a fundamental preprocessing step that ensures the deep learning model can learn effectively, prevent variable dominance, capture intricate climate patterns, and converge efficiently during training. It plays a pivotal role in enhancing the model's performance and its ability to provide valuable insights into climatic phenomena.

Modeling

The two main modeling approaches covered are:

Diffusion model

The probabilistic generative model employed in this context is a sophisticated framework designed to denoise images. It leverages a diffusion process, which is a mathematical concept representing the gradual spread of information or change across data. In the context of image denoising, the diffusion process helps in gradually removing noise from an image while preserving the underlying structure and content.

Advantages of the Model:

Learning Capacity: The probabilistic generative model is endowed with significant learning capacity. It has the ability to learn intricate structures and patterns from data. This means it can effectively capture complex features, textures, and nuances present in images. By learning from a diverse range of images, it becomes proficient in identifying and preserving the underlying information even in noisy or low-resolution inputs.
Extrapolation: The model exhibits a remarkable generalization capability known as extrapolation. It means that once the model has learned from a set of training data, it can extend its knowledge to new and unseen scenarios. This ability is invaluable in real-world applications where the model encounters image inputs it hasn't explicitly seen during training. Despite this, it can produce high-quality denoised outputs.
Realism: A key strength of the probabilistic generative model is its capacity to produce denoised images that maintain a high level of realism. This realism extends to preserving fine details, textures, and nuances in the upscaled images. Additionally, the model is adept at handling artifacts that may be present in the input images, resulting in outputs that closely resemble natural, artifact-free images.

The scheme of the Diffusion process:

The diffusers library

The Diffusers library provides a comprehensive set of options for working with Diffusion Models. In this documentation, we explore various options and functionalities available in the library that can be tailored to specific use cases or extended with custom implementations.

diffusers.UNet2DModel

The diffusers.UNet2DModel class closely resembles our U-net architecture. It offers flexibility in designing the down and up blocks, making it a versatile choice for various tasks.

Down Blocks

You can choose from a variety of down block types, including:

DownBlock2D
ResnetDownsampleBlock2D
AttnDownBlock2D
CrossAttnDownBlock2D
SimpleCrossAttnDownBlock2D
SkipDownBlock2D
AttnSkipDownBlock2D
DownEncoderBlock2D
AttnDownEncoderBlock2D
KDownBlock2D
KCrossAttnDownBlock2D

Up Blocks

The available up block types include:

UpBlock2D
ResnetUpsampleBlock2D
CrossAttnUpBlock2D
SimpleCrossAttnUpBlock2D
AttnUpBlock2D
SkipUpBlock2D
AttnSkipUpBlock2D
UpDecoderBlock2D
AttnUpDecoderBlock2D
KUpBlock2D
KCrossAttnUpBlock2D

Here's an example configuration for diffusers.UNet2DModel:

training_configuration:
  type: diffusion
  model_configuration:
    eps_model:
      class_name: diffusers.UNet2DModel
      kwargs:
        block_out_channels: [112, 224, 336, 448]
        down_block_types: [DownBlock2D, AttnDownBlock2D, AttnDownBlock2D, AttnDownBlock2D]
        up_block_types: [AttnUpBlock2D, AttnUpBlock2D, AttnUpBlock2D, UpBlock2D]
        layers_per_block: 2
        time_embedding_type: positional
        num_class_embeds: 24
        in_channels: 2
        norm_num_groups: 4
    scheduler:
      class_name: LMSDiscreteScheduler
      kwargs:
        num_train_timesteps: 1000
        beta_start: 0.0001
        beta_end: 0.02
        beta_schedule: linear
        prediction_type: epsilon
        timestep_spacing: trailing

The diffusers.UNet2DModel also accepts conditioning on labels through its argument class_labels. First, the embedding type must be specified in the __init__ method trough:

Passing class_embed_type (Options are 'timestep', 'identity' or None).
Passing num_class_embeds with the size of the dictionary of embeddings to use.

For example, to consider the hour of the data as covariate in this model we have two options:

Option A: Set num_class_embeds = 24 in the model creation and hour_embed_type = class in training configuration. This way the model learns an Embedding table for each hour.

Option B: Set class_embed_type = identity in the model configuration and hour_embed_type = positional in training configuration.

Option C: Set class_embed_type = timestep in the model configuration and hour_embed_type = timestep in training configuration. This configuration applies the same cos & sin transformation as in Option B maintaining the same max_duration=10000. Unlike Option B, we fit 2 nn.Linear after the embedding before feeding it to the NN.

diffusers.UNet2DConditionModel

The diffusers.UNet2DConditionModel is an extension of the previous diffusers.UNet2DModel to consider conditions during the reverse process such as time stamps, or other covariables.

One interesting parameter to tune is the activation funcion used in the time embedding which can be: Swish, Mish, SiLU or GELU.

But the most remarkable difference is the possibility of conditioning the reverse diffusion process in the encoder hidden states (comming from images, text, or any other)

One example configuration to use diffusers.UNet2DConditionModel is included below:

training_configuration:
  type: diffusion
  model_configuration:
    eps_model:
      class_name: diffusers.UNet2DConditionModel
      kwargs:
        block_out_channels: [
          124,
          256,
          512
        ]
        down_block_types: [
          CrossAttnDownBlock2D,
          CrossAttnDownBlock2D,
          DownBlock2D
        ]
        mid_block_type: UNetMidBlock2DCrossAttn
        up_block_types: [
          UpBlock2D,
          CrossAttnUpBlock2D,
          CrossAttnUpBlock2D
        ]
        layers_per_block: 2
        time_embedding_type: positional
        in_channels: 2
        out_channels: 1
        sample_size: [20, 32]
        only_cross_attention: False
        cross_attention_dim: 256
        addition_embed_type: other

Tailored UNet

In particular, a tailored U-Net architecture with 2D convolutions, residual connections and attetion layers is used.

The parameteres of these model implemented in deepr/model/unet.py are:

image_channels: It is the number of channels of the high resolution imagen we want to generate, that matches with the number of channels of the output from the U-Net. Default value is 1, as we plan to sample one variable at a time.
n_channels: It is the number of output channels of the initial Convolution. Defaults to 16.
channel_multipliers: It is the multiplying factor over the channels applied at each down/upsampling level of the U-Net. Defaults to [1, 2, 2, 4].
is_attention: It represents the use of Attention over each down/upsampling level of the U-Net. Defaults to [False, False, True, True].
n_blocks: The number of residual blocks considered in each level. Defaults to 2.
conditioned_on_input: The number of channels of the conditions considered.

NOTE I: The length of channel_multipliers and is_attention should match as it sets the number of resolutions of our U-Net architecture.

NOTE II: Spatial tensors fed to Diffusion model must have shapes of length multiple of $2^{\text{num resolutions} - 1}$.

An example configuration for this model is specified in training_configuration > model_configuration > eps_model,

training_configuration:
  type: diffusion
  model_configuration:
    eps_model:
      class_name: UNet
      kwargs:
        block_out_channels: [32, 64, 128, 256]
        is_attention: [False, False, True, True]
        layers_per_block: 2
        time_embedding_type: positional
        in_channels: 2
        out_channels: 1
        sample_size: [20, 32]

Downsampling

The class Downsample represents a downsampling block. It uses a convolutional layer to reduce the spatial dimensions of the input tensor. Here are the key details:

Constructor: Initializes a nn.ConvTranspose2d layer with specified input and output channels, kernel size, stride, and padding.
Forward Method: Takes an input tensor x and a time tensor t (though t is not used in this case) and applies the convolution operation to downsample x.

Upsampling

The class Upsample represents an upsampling block. It uses a transposed convolutional layer to increase the spatial dimensions of the input tensor. Here are the key details:

Constructor: Initializes a nn.ConvTranspose2d layer with specified input and output channels, kernel size, stride, and padding.
Forward Method: Takes an input tensor x and a time tensor t (though t is not used in this case) and applies the transposed convolution operation to upsample x.

Down Block

The class Down block represents a block used in the first half of a U-Net architecture for encoding input features. It consists of a residual block followed by an optional attention block. Here are the key details:

Constructor: Initializes a ResidualBlock and, if has_attn is True, an AttentionBlock. These blocks are used for feature extraction during downsampling.
Forward Method: Takes an input tensor x and a time tensor t and passes x through the residual block and, if applicable, the attention block.

Middle Block

The class Middle block represents a block used in the middle section of a U-Net architecture. It contains two residual blocks with an attention block in between. Here are the key details:

Constructor: Initializes two ResidualBlock instances and an AttentionBlock. This block is typically used for processing features in the middle layers of the U-Net.
Forward Method: Takes an input tensor x and a time tensor t and passes x through the first residual block, the attention block, and then the second residual block.

Up Block

The class Up block represents a block used in the second half of a U-Net architecture for decoding features. It consists of a residual block followed by an optional attention block. Here are the key details:

Constructor: Initializes a ResidualBlock and, if has_attn is True, an AttentionBlock. These blocks are used for feature decoding during upsampling.
Forward Method: Takes an input tensor x and a time tensor t and passes x through
the residual block and, if applicable, the attention block.

Residual Block

The class Residual Block is a component commonly used in neural networks. It enhances feature extraction and information flow within the network. Key details include:

Constructor: Initializes the block with input and output channel specifications, time channels, group normalization settings, and optional dropout.
Components:
- Two convolutional layers with group normalization and Swish activation.
- Time embeddings for temporal information.
- Optional dropout.
Forward Method: Takes an input tensor and time tensor, applies convolution, adds time embeddings, and produces the output tensor.

Convolutional Swin2SR

The Convolutional Swin2SR is a state-of-the-art (SOTA) neural network designed for super-resolution tasks in computer vision. It stands out for several key features that make it a powerful tool for enhancing image resolution:

Efficient Scaling: The model's primary component is based on Swin v2 attention layers, which are known for their efficiency and effectiveness. These layers enable the network to efficiently process and generate high-resolution images while maintaining performance.
Easy Experiment Setting: Setting up experiments with the Convolutional Swin2SR is straightforward, making it accessible for researchers and practitioners. The model's architecture and parameters are designed for ease of use and experimentation.
Fast Training and Inference: Thanks to its efficient design, the Convolutional Swin2SR offers fast training and inference times. This efficiency is particularly valuable when dealing with large datasets or real-time applications.

Loss Terms: The model employs various loss terms to guide the training process effectively:

L1 Loss of Predictions and References: This loss term measures the difference between the model's predictions and the high-resolution reference images. It encourages the model to generate outputs that closely match the ground truth.
L1 Loss of Downsampled Predictions and References: To further refine the training process, the model also considers downsampled versions of both predictions and references. This helps in capturing details at multiple scales.
L1 Loss of Blurred Predictions and References: Blurring is introduced as an additional loss term, allowing the model to learn and recover fine details while handling different levels of image degradation.

For training the Convolutional-Swin2SR, a configuration similar to the one provided needs to be given:

training_configuration:
  type: end2end
  model_configuration:
    neural_network:
      class_name: ConvSwin2SR
      kwargs:
        embed_dim: 180
        depths: [6, 6, 6, 6, 6, 6]
        num_heads: [6, 6, 6, 6, 6, 6]
        patch_size: 1
        window_size: 5
        num_channels: 1
        img_range: 1
        resi_connection: "1conv"
        upsampler: "pixelshuffle"
        interpolation_method: "bicubic"
        hidden_dropout_prob: 0.0
        upscale: 5

Training configuration: commons

There are training parameters that are common to all the models:

num_epochs: The number of training epochs.
batch_size: The batch size for training.
gradient_accumulation_steps: The number of gradient accumulation steps.
learning_rate: The initial learning rate.
lr_warmup_steps: The number of warm-up steps for learning rate scheduling.
mixed_precision: Mixed-precision training, e.g., "fp16".
hour_embed_type: Type of hour embedding, e.g., "class".
hf_repo_name: The Hugging Face repository name for model storage.
output_dir: The directory for saving training outputs.
device: The device for training, e.g., "cuda" for GPU.
push_to_hub: Whether to push the trained model to the Hugging Face model hub.
seed: Random seed for reproducibility.
save_model_epochs: Frequency of saving the model during training.
save_image_epochs: Frequency of saving images during training.

An example of how they should be defined in the configuration file is provided:

training_parameters:
    num_epochs: 50
    batch_size: 8
    gradient_accumulation_steps: 2
    learning_rate: 0.001
    lr_warmup_steps: 500
    mixed_precision: "fp16"
    hour_embed_type: class # none, timestep, positional, cyclical, class
    output_dir: "/PATH/TO/WRITE/ARTIFACTS"
    device: cuda
    seed: 2023
    save_model_epochs: 5
    save_image_epochs: 5

Training a model

To train your model, follow these steps:

Prepare your dataset: Ensure that your dataset is properly formatted with all the different netCDF files inside the same folder structure.
Configure your training parameters: Create a configuration file (usually in YAML format) that specifies various training hyperparameters, such as learning rate, batch size, number of epochs, etc. You can use the provided configuration examples as a starting point.
Start training: Run the training script, specifying the path to your configuration file. The training script is located at train_model.py

Generating model predictions

To make predictions using the model you've defined, you can use the provided script: generate_model_predictions.py

This script is designed to generate predictions using your trained model. You can run it with the appropriate input data to obtain model predictions for your specific task or application.

Appendix I: Positional Embeddings

When working with sequential data, the order of the elements is important, and we must pay attention to how we pass this information to our models.

In our particular case, the timesteps $t$ is encoded with positional embeddings as proposed in the Denoising Diffusion Probabilistic Models paper.

Besides, we may encode other important features as the hour of the day or the day of the year, which are cyclical. This is different from positional encodings because we want the encoding from hour 23 to be more similar to the one from 0 than from hour 18.

References

Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33, 6840-6851.
Song, J., Meng, C., & Ermon, S. (2020). Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502.
Conde, M. V., Choi, U. J., Burchi, M., & Timofte, R. (2022). Swin2SR: Swinv2 transformer for compressed image super-resolution and restoration. arXiv preprint arXiv:2209.11345.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., & Ommer, B. (2022). High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 10684-10695).
Annotated Deep Learning Paper implementations

License

Copyright 2023, European Union.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Workflow for developers/contributors

For best experience create a new conda environment (e.g. DEVELOP) with Python 3.10:

mamba create -n deepr-cuda -c conda-forge python=3.10
mamba activate deepr-cuda
make mamba-cuda_env-update

A data directory for the testing data must be created:

cd tests
mkdir data
cd data
mkdir features
mkdir labels

Once the directories have been created, testing data can be downloaded:

cd tests
wget -O data.zip https://cloud.predictia.es/s/zen8PGwJbi7mTCB/download
unzip data.zip
rm data.zip

Before pushing to GitHub, run the following commands:

Update conda environment: make conda-env-update
Install this package: pip install -e .
Sync with the latest template (optional): make template-update
Run quality assurance checks: make qa
Run tests: make unit-tests
Run the static type checker: make type-check
Build the documentation (see Sphinx tutorial): make docs-build

Name		Name	Last commit message	Last commit date
Latest commit History 335 Commits
.github/workflows		.github/workflows
ci		ci
deepr		deepr
docs		docs
resources		resources
scripts		scripts
tests		tests
.cruft.json		.cruft.json
.gitignore		.gitignore
.pre-commit-config-cruft.yaml		.pre-commit-config-cruft.yaml
.pre-commit-config.yaml		.pre-commit-config.yaml
Dockerfile		Dockerfile
LICENSE		LICENSE
Makefile		Makefile
README.md		README.md
environment.yml		environment.yml
environment_CUDA.yml		environment_CUDA.yml
pyproject.toml		pyproject.toml
setup.cfg		setup.cfg

License

ECMWFCode4Earth/DeepR

Folders and files

Latest commit

History

Repository files navigation

DeepR: Deep Reanalysis.

Introduction

Project Motivation

Super-Resolution in Climate Science

Data: The Foundation of DeepR

Focus on a specific domain

Definition in configuration file

Features Configuration

Label Configuration

Split Coverages

Methodology

Standardization

Why is it important?

Application

Modeling

Diffusion model

The diffusers library

diffusers.UNet2DModel

Down Blocks

Up Blocks

diffusers.UNet2DConditionModel

Tailored UNet

Downsampling

Upsampling

Down Block

Middle Block

Up Block

Residual Block

Convolutional Swin2SR

Training configuration: commons

Training a model

Generating model predictions

Appendix I: Positional Embeddings

References

License

Workflow for developers/contributors

About

Resources

License

Stars

Watchers

Forks

Releases

Packages 0

Contributors 3

Languages

Packages