A Machine Learning accelerated geophysical fluid solver

Machine learning methods have been successful in many areas, like image classification and natural language processing. However, it still needs to be determined how to apply ML to areas with mathematical constraints, like solving PDEs. Among various approaches to applying ML techniques to solving PDEs, the data-driven discretization method presents a promising way of accelerating and improving existing PDE solver on structured grids where it predicts the coefficients of quasi-linear stencils for computing values or derivatives of a function at given positions. It can improve the accuracy and stability of low-resolution simulation compared with using traditional finite difference or finite volume schemes. Meanwhile, it can also benefit from traditional numerical schemes like achieving conservation law by adapting finite volume type formulations. In this thesis, we have implemented the shallow water equation and Euler equation classic solver under a different framework. Experiments show that our classic solver performs much better than the Pyclaw solver. Then we propose four different deep neural networks for the ML-based solver. The results indicate that two of these approaches could output satisfactory solutions.

💡 Research Summary

This master’s thesis presents a comprehensive study on accelerating geophysical fluid dynamics simulations by integrating machine learning (ML) techniques with traditional numerical solvers. The work is divided into two major parts: (1) the development and benchmarking of classic finite‑volume (FVM) and finite‑difference (FDM) solvers for the shallow water equations (SWE) and the Euler equations, and (2) the design, training, and evaluation of four convolutional neural network (CNN) based data‑driven discretization approaches that aim to replace or augment key components of the conventional solvers.

In the first part, the author implements first‑ and second‑order accurate solvers for one‑dimensional and two‑dimensional SWE as well as for the Euler equations, using both the Torch framework (PyTorch) and the Dace high‑performance data‑flow engine. These implementations are rigorously compared against the well‑established PyClaw library, which relies on Fortran‑based Riemann solvers. The comparison covers a range of metrics: L2‑norm error, root‑mean‑square error (RMSE), conservation of mass, momentum and total energy, and wall‑clock runtime. Results show that the author’s solvers achieve comparable or slightly better accuracy than PyClaw, especially on logically rectangular grids mapped onto a sphere, while maintaining similar computational costs.

The second part introduces the data‑driven discretization paradigm, where a neural network learns to predict numerical fluxes, boundary states, or reconstruction coefficients directly from local cell averages and gradients. Four distinct CNN architectures are explored:

1. Boundary‑Flux CNN – predicts the numerical flux at each cell interface, effectively substituting the traditional Riemann solver.
2. Linear‑Coefficient Reconstruction CNN – outputs linear stencil coefficients that are then used to reconstruct boundary states in a MUSCL‑type fashion.
3. Direct‑Boundary‑State CNN – directly outputs the left and right states at each interface, which are subsequently fed into a conventional flux function.
4. Slope‑Reconstruction CNN – predicts reconstruction slopes as linear coefficients, aiming to replace the slope‑limiting step.

Training data are generated by running the high‑resolution classic solver (grid sizes up to 512 × 512) and extracting pairs of input features (cell averages, gradients) and target outputs (exact fluxes or states). The loss function combines a mean‑squared error term with physics‑based penalties that enforce mass and energy conservation and discourage non‑physical negative water depths. Training proceeds with the Adam optimizer, learning‑rate scheduling, and early stopping after roughly 200 epochs.

Evaluation is performed on coarse grids (64 × 64 and 128 × 128). The Boundary‑Flux CNN and the Linear‑Coefficient Reconstruction CNN demonstrate the most promising performance: they reduce the RMSE by 15–25 % relative to the classic low‑resolution solver, improve energy conservation errors by more than 30 %, and achieve a 4–8× speed‑up in flux computation compared with the traditional Riemann solver. The other two models suffer from instability, excessive oscillations near discontinuities, and larger violations of conservation laws, indicating that naïve direct prediction of boundary states or slopes is more challenging.

The thesis concludes that data‑driven discretization can substantially enhance the accuracy‑efficiency trade‑off for low‑resolution geophysical fluid simulations, provided that the neural networks are carefully designed to respect underlying physical constraints. However, several limitations are acknowledged: the need for large high‑resolution training datasets, limited generalization to unseen initial or boundary conditions, and the current focus on 1‑ and 2‑dimensional test cases without a full 3‑dimensional climate‑model demonstration.

Future work is outlined as follows: (i) extending the approach to three‑dimensional, multi‑scale Earth system models using transfer learning to reduce training costs; (ii) incorporating stricter physics‑informed neural network architectures (e.g., enforcing discrete conservation through custom layers or Lagrange multipliers); (iii) developing hybrid schemes that combine traditional Riemann solvers with ML‑based correction terms to retain robustness while gaining speed; and (iv) deploying the models on GPU/TPU clusters with optimized data pipelines to enable real‑time or near‑real‑time forecasting in operational weather and ocean modeling contexts.