Interpretable Neural Networks to Estimate Momentum Fluxes of Orographic Gravity Waves

State-of-the-art Earth system models (ESMs) cannot explicitly resolve many small-scale atmospheric processes such as atmospheric gravity waves, and thus must represent, or parameterize, them based on the resolved state. Machine learning (ML) has the potential to address this. In our study, we train neural networks on ERA5 reanalysis data to predict momentum fluxes of orographic gravity waves as function of the lower resolution state variables as would be represented by a coarse ESM. Employing a full year of ERA5 data, we filter inertia-gravity waves by normal-mode function decomposition using the software MODES, and train ML models, more precisely: U-Nets, on data coarse-grained to the ESM's target resolution. We consider four different cases: the full spectrum of resolved inertia-gravity waves or just its subgrid-scale part, both over all land or just over mountainous terrain. Our neural networks successfully predict momentum fluxes, with a global coefficient of determination (R²) ranging from 0.72 to 0.56, depending on the case, when evaluated offline with unseen data. An analysis of our models using SHAP values, an explainable AI technique, shows that the networks are learning physically meaningful relationships. In addition, we give a comparison with the physics-based parameterization scheme by Lott and Miller. These results offer the opportunity for the development of operational ML-based parameterizations to improve the representation of gravity waves and their effects in climate models.