piDLRad (physics-informed Deep Learning based Radiation solver)

PyTorch implementation for the paper:

Revisiting Machine Learning Approaches for Short- and Longwave Radiation Inference in Weather and Climate Models

Authors:

• Salman Mohebi
• Firat Ozdemir

Installation

If you do not have conda, you can install Miniconda:

• Download the Miniconda installer for your OS from the official page.
• Run the installer and follow the prompts.

Once conda is available, it is recommended to use a conda environment to manage dependencies and avoid conflicts:

conda create -n pidlrad python=3.8
conda activate pidlrad

You can install piDLRad using pip:

pip install git+https://github.com/swissdatasciencecenter/deepcloud-pidlrad.git

Alternatively, clone the repository and install in editable mode:

git clone https://github.com/swissdatasciencecenter/deepcloud-pidlrad.git
cd deepcloud-pidlrad
pip install -e .

Usage

Here's a quick example of how to use piDLRad for inference:

import torch
from torchvision import transforms
from src.pidlrad.utils.load_data import IconDataset, IconH5Metadata, HeightCutter
from src.pidlrad.nn import get_model

# Load metadata and normalization statistics
icon_metadata = IconH5Metadata('path/to/metadata.h5')
x3d_mean, x3d_std, x2d_mean, x2d_std = icon_metadata.get_mean_std('pfph')

# Create mock args object with model configuration
class Args:
    model = 'mlp'  # or 'unet', 'lstm', 'vit', etc.
    height_in = 70
    height_out = 71
    channel_3d = 6
    channel_2d = 6
    # Add other model-specific parameters as needed

args = Args()

# Initialize and load trained model
model = get_model(x3d_mean, x3d_std, x2d_mean, x2d_std, args)
model.load_state_dict(torch.load('path/to/trained_model.pth'))
model.eval()

# Load data
transform = transforms.Compose([HeightCutter(height_in=args.height_in, height_out=args.height_out)])
dataset = IconDataset(
    data_dir='path/to/data',
    transform=transform
)

# Run inference on a sample
x3d, x2d, y_true = dataset[0]  # Get first sample
x3d, x2d = x3d.unsqueeze(0), x2d.unsqueeze(0) # Add batch dimension

with torch.no_grad():
    y_pred = model(x3d, x2d)
    print(f"Predicted radiative fluxes shape: {y_pred.shape}")

Dataset

The entire dataset is available at: https://doi.org/10.3929/ethz-b-000721647

Dataset summary:

This dataset contains HDF5 files from offline simulations of the ICON aquaplanet model, designed for radiation solver emulation. Each file represents a single time step and includes 2D and 3D atmospheric input features as well as radiative flux outputs. Additional files provide normalization statistics, grid information, and sample loading scripts. All variables follow ICON and ecrad conventions.

• Input features:
  • 2D: surface pressure, solar zenith angle, surface humidity, visible/near-IR albedo, surface temperature
  • 3D: cloud cover, temperature, pressure, cloud water/ice content, specific humidity
• Output features:
  • Longwave and shortwave radiative fluxes (upward and downward)

Data is provided under a CC BY-NC 4.0 license. For details on file structure, variable definitions, and usage, see the README and included sample scripts.

Acknowledgement

This work was funded through a grant by the Swiss Data Science Center (SDSC grant C20-03). This research was supported by computation resources provided by the EXCLAIM project funded by ETH Zurich (CSCS project number d121). The Center for Climate Systems Modeling (C2SM) at ETH Zurich is acknowledged for providing technical and scientific support. Sebastian Schemm and Stefan Rüdisühli are supported by the European Research Council, H2020 European Research Council (grant no. 848698). From November 2024 onward, Guillaume Bertoli was supported by the SNF postdoc mobility grant P500PN_225397.

Copyright

Copyright © 2025-2028 Swiss Data Science Center (SDSC), www.datascience.ch. All rights reserved. The SDSC is jointly established and legally represented by the École Polytechnique Fédérale de Lausanne (EPFL) and the Eidgenössische Technische Hochschule Zürich (ETH Zürich). This copyright encompasses all materials, software, documentation, and other content created and developed by the SDSC.