Physics-based machine learning for mantle convection simulations

TL;DR

A physics-based machine learning approach that predicts creeping flow velocities as a function of temperature while conserving mass, thereby bypassing the numerical solution of the Stokes problem and enabling autoregressive rollout at inference.

Abstract

Mantle convection simulations are an essential tool for understanding how rocky planets evolve. However, the poorly known input parameters to these simulations, the non-linear dependence of transport properties on pressure and temperature, and the long integration times in excess of several billion years all pose a computational challenge for numerical solvers. We propose a physics-based machine learning approach that predicts creeping flow velocities as a function of temperature while conserving mass, thereby bypassing the numerical solution of the Stokes problem. A finite-volume solver then uses the predicted velocities to advect and diffuse the temperature field to the next time-step, enabling autoregressive rollout at inference. For training, our model requires temperature-velocity snapshots from a handful of simulations (94). We consider mantle convection in a two-dimensional rectangular box with basal and internal heating, pressure- and temperature-dependent viscosity. Overall, our model is up to 89 times faster than the numerical solver. We also show the importance of different components in our convolutional neural network architecture such as mass conservation, learned paddings on the boundaries, and loss scaling for the overall rollout performance. Finally, we test our approach on unseen scenarios to demonstrate some of its strengths and weaknesses.

Figures (9)

• Figure 1: The ranges of different fields in the simulations (a-e). Scaling the velocities brings the ranges across different simulation parameters closer to each other (f,g). For ease of visualization, we normalize the $u$ and $v$ range for the scaled and unscaled cases by dividing them by the maximum (f,g,h).
• Figure 2: Our hybrid physics-based model for time-stepping. (a) We use a data-driven approach to model the primary computational bottleneck in mantle convection simulations: solving the Stokes equation. The velocities predicted by the convolutional neural network (CNN) can be fed to the numerical solver for time-stepping. (b) The CNN architecture is inspired by tompson2017 for modeling solutions of linear systems as a function of the right-hand-side of the equation. (c) Instead of standard padding methods (such as zero paddings), we use boundary-learned convolutions. A unique set of filters is responsible for predicting each corner and each edge as well as the interior domain. (d) Mass conservation is enforced by predicting a field whose curl yields divergence-free velocity components $u$ and $v$.
• Figure 3: Examples of velocity predictions compared to ground truth from the direct solver for two different unseen simulations (a,b) and (c,d) characterized by widely different parameters leading to relatively low (a,b) and high velocities (c,d). A single model is able to learn across four orders of magnitude. The prediction error is divided by the absolute maximum value and displayed as a percentage error in the right column.
• Figure 4: (a) U-net schematic for predicting from state at time $t$ to state at $t+1$. (b) Four different U-nets are trained to evaluate if different parameter counts and components like boundary-learned convolution improve the prediction accuracy. These are plotted alongside our best Stokes surrogate model that predicts velocities (solid grey lines). (c) When rolling out autoregressively with the best performing network (dashed black lines), the predictions diverge within $16$ steps.
• Figure 5: Evaluation of rollout performance on three qualitatively different simulations (a,b), (c,d) and (e,f). For each simulation, we plot the ground truth in the first column, followed by the prediction (either baseline of $100$ momentum skips, or the ML model) in the second, the horizontally-average temperature profile in the third, and the mean temperature as a function of time in the fourth.