Machine-Learned Preconditioners for Linear Solvers in Geophysical Fluid Flows

Semi-implicit grid-point models for the atmosphere and the ocean require linear solvers that are working efficiently on modern supercomputers. The huge advantage of the semi-implicit time-stepping approach is that it enables large model time-steps. This however comes at the cost of having to solve a computationally demanding linear problem each model time-step to obtain an update to the model’s pressure/fluid-thickness field. In this study, we investigate whether machine learning approaches can be used to increase the efficiency of the linear solver.

Our machine learning approach aims at replacing a key component of the linear solver—the preconditioner. In the preconditioner an approximate matrix inversion is performed whose quality largely defines the linear solver’s performance. Embedding the machine-learning method within the framework of a linear solver circumvents potential robustness issues that machine learning approaches are often criticized for, as the linear solver ensures that a sufficient, pre-set level of accuracy is reached. The approach does not require prior availability of a conventional preconditioner and is highly flexible regarding complexity and machine learning design choices.

Several machine learning methods of different complexity from simple linear regression to deep feed-forward neural networks are used to learn the optimal preconditioner for a shallow-water model with semi-implicit time-stepping. The shallow-water model is specifically designed to be conceptually similar to more complex atmosphere models. The machine-learning preconditioner is competitive with a conventional preconditioner and provides good results even if it is used outside of the dynamical range of the training dataset.