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

Abstract

This paper explores Machine Learning (ML) parameterizations for radiative transfer in the ICOsahedral Nonhydrostatic weather and climate model (ICON) and investigates the achieved ML model speed-up with ICON running on graphics processing units (GPUs). Five ML models, with varying complexity and size, are coupled to ICON; more specifically, a multilayer perceptron (MLP), a Unet model, a bidirectional recurrent neural network with long short-term memory (BiLSTM), a vision transformer (ViT), and a random forest (RF) as a baseline. The ML parameterizations are coupled to the ICON code that includes OpenACC compiler directives to enable GPU support. The coupling is done with the PyTorch-Fortran coupler developed at NVIDIA. The most accurate model is the BiLSTM with a physics-informed normalization strategy, a penalty for the heating rates during training, a Gaussian smoothing as postprocessing and a simplified computation of the fluxes at the upper levels to ensure stability of the ICON model top. The presented setup enables stable aquaplanet simulations with ICON for several weeks at a resolution of about 80 km and compares well with the physics-based default radiative transfer parameterization, ecRad. Our results indicate that the compute requirements of the ML models that can ensure the stability of ICON are comparable to GPU optimized classical physics parameterizations in terms of memory consumption and computational speed.