Generalizing machine-learned discretization for climate simulations: addressing ‘out-of-sample’ challenges for 2D data-driven advection discretization

The effective spatial and temporal scales resolved by Earth System Models (ESMs) remain a key limitation in reducing uncertainties in climate projections. While increasing model resolution is computationally prohibitive, machine learning (ML)-based parameterizations offer a promising alternative. However, these approaches often face generalization challenges in ‘out-of-sample’ scenarios, leading to numerical instabilities when integrated into ESMs. In this study, we aim to tackle these challenges by developing a data-driven discretization neural network for multidirectional advection in ocean models. The canonical 1D advection problem is revisited by using neural networks to predict the coefficients of a three-node stencil trained on high-resolution solutions projected onto coarser spatial and temporal grids. Conventional discretizations generalize to all scalar fields, while the data-driven approach is, by construction, tied to the training data. First, it is shown that we can normalize inputs with min-max scaling to achieve generalization, while training on coarsened high-resolution data across multiple grid configurations reduces sensitivity to time steps and mesh resolution. We find that coarsening based on triangular test functions, instead of averaging, enables unique mapping of the fine-scale variations of high-resolution solutions, leading to monotonicity of the neural network. Hybrid ML discretizations that predict advective fluxes are investigated, with a focus on enforcing desirable numerical properties—such as monotonicity, accuracy, and stability. Finally, we aim to test the numerical and generalization properties of the new data-driven discretization on 2D geostrophic flows. These results provide guidance for the development of better end-to-end data-driven parameterizations and discretizations in ESMs.