Project No: 16301322

Title: Large Eddy Simulation Code in JAX: An Accelerated and Differentiable Atmospheric Model for Turbulence Parameterization Development

Principal Investigator: Prof. Xiaoming SHI

Abstract:

Large-eddy simulations (LESs) are an essential tool for studies on atmospheric turbulence and clouds. LESs play critical roles in the development of turbulence and convection parameterizations. Current global models are approaching kilometer-scale resolution as supercomputing facilities advance. However, this resolution range is in the so-called gray zone, where subgrid-scale (SGS) turbulence actively interacts with resolved motion and significantly influences the large-scale characteristics of simulated weather systems. Thus, developing SGS turbulence models for the gray zone requires new LES models, which must run sufficiently fast and energy-efficiently when simulating large domains. Meanwhile, the development of next-generation Earth system models (ESMs) also demands new LES code because future ESMs are expected to be capable of learning from targeted LESs in some ESM columns for updating uncertain parameters on demand. The recent development of JAX, a Python library for high-performance numerical computing, enables the development of new LES code that can utilize GPU acceleration efficiently. The first goal of this project is to develop a new LES code in JAX with adequate parallelism for turbulence parameterization development in both gray-zone and next-generation ESM scenarios. We will use the generalized pseudo-incompressible equations as the dynamical core of the new LES model. This set of equations filters sound waves on a physical approximation basis, making the computational workflow simpler than that of other approaches. Moreover, flexibility to adopt other governing equations will also be ensured in the development. We will evaluate the accuracy and performance of the new LES code with classic testing cases and ensure sufficient hardware acceleration is achieved for targeting large-domain applications. Our second goal is to establish frameworks for training machine learning (ML)-based SGS turbulence models with the new LES code. In addition to hardware acceleration, JAX enables automatic differentiation (AD) through the LES model. AD tools have been used extensively in the atmospheric science community for data assimilation and sensitivity studies. Here, with the AD capacity, the new LES code can be used for coupled-online training of ML-based SGS turbulence models. The loss (cost) function for training can be constructed based on observations of model states, and backpropagation through the dynamical core enables us to update the ML model parameters for optimization. We will use dry and moist boundary layer cases to evaluate the efficiency and efficacy of different ML training strategies. The LES and ML codes from this project will be provided to the community on an open-source basis.