Accelerating the discovery of steady-states of planetary interior dynamics with machine learning

TL;DR

The benefit of this method lies in requiring very few simulations to train on, providing a solution with no prediction error as the authors initialize a numerical method, and posing minimal computational overhead at inference time.

Abstract

Simulating mantle convection often requires reaching a computationally expensive steady-state, crucial for deriving scaling laws for thermal and dynamical flow properties and benchmarking numerical solutions. The strong temperature dependence of the rheology of mantle rocks causes viscosity variations of several orders of magnitude, leading to a slow-evolving stagnant lid where heat conduction dominates, overlying a rapidly-evolving and strongly convecting region. Time-stepping methods, while effective for fluids with constant viscosity, are hindered by the Courant criterion, which restricts the time step based on the system's maximum velocity and grid size. Consequently, achieving steady-state requires a large number of time steps due to the disparate time scales governing the stagnant and convecting regions. We present a concept for accelerating mantle convection simulations using machine learning. We generate a dataset of 128 two-dimensional simulations with mixed basal and internal heating, and pressure- and temperature-dependent viscosity. We train a feedforward neural network on 97 simulations to predict steady-state temperature profiles. These can then be used to initialize numerical time stepping methods for different simulation parameters. Compared to typical initializations, the number of time steps required to reach steady-state is reduced by a median factor of 3.75. The benefit of this method lies in requiring very few simulations to train on, providing a solution with no prediction error as we initialize a numerical method, and posing minimal computational overhead at inference time. We demonstrate the effectiveness of our approach and discuss the potential implications for accelerated simulations for advancing mantle convection research.

Figures (10)

• Figure 1: A visualization of the normalized simulation parameters and temperature profiles for train (blue), validation (cv, orange) and test (green) sets. The test set is deliberately picked in this manner to check for NNs ability to extrapolate slightly.
• Figure 2: We use a feedforward neural network to predict the temperature at a given height $y$ as a function of the simulation parameters: internal heating ($Q$), viscosity contrast due to temperature ($\gamma$), and viscosity contrast due to pressure ($\beta$). This pointwise prediction, instead of providing the complete temperature profile at once, helps prevent oscillations in the output of the network. The network has skip connections to regularize the loss landscape. We also condition the last hidden layer with the input vector by concatenation, as this tends to improve the accuracy of the predictions. Although not shown in the figure, we apply $SELU()$ activation to each hidden layer, i.e., each layer except the input and the output.
• Figure 3: The "tangent method" is used for determining the thickness of the lid. a) A tangent is drawn from the point of maximum gradient of the horizontally averaged profile of the RMS velocity to where it equals zero. This is taken as the lid thickness from top or the height of the lid in terms of the y-coordinate. b) The temperature contrast between this point and the maximum temperature (i.e. $T=1$ at the bottom of the domain) is used to determine an effective Rayleigh number of the system with respect to which scaling laws are derived.
• Figure 4: Comparison of neural network predictions and those of the baseline algorithms against ground truth temperature profiles for the validation (cv) set.
• Figure 5: Comparison of neural network predictions and those of the baseline algorithms against ground truth temperature profiles for the test set.