A Machine Learning accelerated geophysical fluid solver

This thesis explores the application of machine learning to geophysical fluid dynamics by implementing a data-driven discretization method that predicts stencil coefficients to accelerate and improve the accuracy and stability of shallow water and Euler equation solvers.

The Gist

Imagine you are trying to predict the path of a massive storm or a tsunami. To do this, scientists use supercomputers to run "simulations"—essentially digital versions of the ocean or the atmosphere.

However, there is a massive problem: The math is too heavy. To get a perfect, crystal-clear picture of a storm, you need a "high-resolution" grid (like a 4K ultra-HD screen). But a 4K simulation requires so much computing power that it might take weeks to finish. If you use a "low-resolution" grid (like an old, blurry 1990s TV), the simulation runs fast, but it’s inaccurate—it might miss the exact moment a wave crashes or a wind gust hits.

This Master’s thesis by Yang Bai explores a way to get the speed of the blurry TV with the accuracy of the 4K screen by using Machine Learning.

---

The Core Idea: The "Smart Sketch Artist"

Think of a traditional physics solver like a strict mathematician. Every single step, the mathematician looks at the current state of the water and uses complex, rigid formulas to calculate exactly where the water will move next. It’s incredibly accurate, but it’s slow because the mathematician has to do a mountain of arithmetic for every single drop of water.

The author proposes replacing parts of that math with a Machine Learning "Sketch Artist" (a Neural Network).

Instead of making the mathematician calculate every tiny detail from scratch, we train a neural network by showing it millions of examples of "High-Definition" simulations. The neural network learns the patterns of how water moves.

Now, when we run a "Low-Definition" simulation, the math handles the big, obvious movements, but whenever the simulation hits a tricky spot (like a sharp wave or a sudden change in depth), we call in the Sketch Artist. The AI looks at the blurry data and says, "I've seen this pattern before; based on my training, the high-def version of this wave should actually look like THIS." It "fills in the blanks" with high-resolution intelligence.

---

The Four Experiments (The "Trial and Error" Phase)

The author didn't just try one way to use AI; they tried four different "training styles" to see which one worked best:

1. The Direct Approach (The "Guess the Whole Result" Method): The AI tries to predict the entire movement of the water directly.
   • Result: Failure. It was like asking a sketch artist to draw a whole landscape from memory without looking at the actual ground; it eventually lost track of reality and the simulation "exploded."
2. The Interpolation Approach (The "Connect the Dots" Method): The AI tries to guess the values at the edges of the grid cells.
   • Result: Unstable. It was a bit too jittery. It created "glitches" in the water that made the simulation crash.
3. The Direct Reconstruction (The "Smart Detailer" Method): The AI looks at the center of a cell and predicts what the edges should look like.
   • Result: Success! This worked well. It was much better than the old-school methods, though it still had a little bit of "digital noise" (like static on a radio).
4. The Slope Method (The "Precision Architect" Method): Instead of guessing the values, the AI predicts the angle or slope of the water's surface.
   • Result: The Winner! This was the most stable and accurate. It’s like giving the sketch artist a ruler instead of just a pencil. It allowed the simulation to stay smooth, realistic, and incredibly sharp.

---

Why does this matter?

In the real world, we can't wait three weeks to find out if a hurricane is going to hit a coastline. We need answers in minutes or hours.

By "embedding" AI into the traditional math, Yang Bai has shown that we can create a hybrid solver. It’s a system that respects the unbreakable laws of physics (the math) but uses the pattern-recognition genius of AI to skip the tedious calculations. This paves the way for faster, smarter, and more accurate weather and ocean predictions.

Technical Summary

Technical Summary: A Machine Learning Accelerated Geophysical Fluid Solver

Author: Yang Bai
Date: November 30, 2022 (with arXiv update Feb 2026)
Program: MSc. Scientific Computing, Technical University of Berlin / ETH Zurich

---

1. Problem Statement

The simulation of complex physical systems via nonlinear Partial Differential Equations (PDEs)—specifically the Shallow Water Equations (SWE) and the Euler Equations—is computationally expensive. These equations are critical for modeling climate, weather, tsunamis, and gas dynamics.

Traditional numerical methods, such as Finite Volume Methods (FVM), face a fundamental trade-off: increasing spatial resolution to capture fine-scale phenomena (like shock waves or turbulence) leads to a massive increase in computational cost. While Machine Learning (ML) has succeeded in unstructured data tasks (image/language), applying it to PDEs is challenging because the ML models must respect mathematical constraints, such as conservation laws, stability, and positivity (e.g., water height must remain $>0$).

2. Methodology

The thesis approaches the problem through data-driven discretization. Instead of using ML to replace the entire solver, the author embeds ML components into a classical FVM framework. This ensures the solver retains the structural benefits of traditional numerical schemes (like conservation) while using ML to improve accuracy on coarse grids.

A. Classical Solver Development:
The author implemented high-performance solvers for 1D, 2D (planar), and spherical geometries using the Torch and DACE (Data-Centric Parallel Programming) frameworks. The solvers include:

• First and Second-order schemes using various numerical fluxes (Lax-Friedrichs, Rusanov, Roe, HLL, HLLE, and HLLC).
• Slope limiters (MINMOD) to suppress numerical oscillations.
• Spherical mappings (Logically rectangular grids) to avoid polar singularities common in latitude-longitude grids.

B. Machine Learning Architectures:
The author proposed and tested four distinct Convolutional Neural Network (CNN) architectures to accelerate the SWE solver:

1. Direct Flux Prediction: The CNN predicts the boundary numerical flux directly from cell averages.
2. Linear Coefficient Prediction: The CNN predicts coefficients for linear interpolation to reconstruct boundary states.
3. Direct Boundary State Reconstruction: The CNN reconstructs the values at the cell interfaces (left/right states) directly from cell averages.
4. Reconstruction Slope Prediction: The CNN predicts the "slopes" (alpha coefficients) used in high-order reconstruction (similar to WENO schemes).

3. Key Contributions

• High-Performance Implementation: Developed a custom SWE/Euler solver that outperforms the industry-standard Pyclaw in terms of speed and flexibility (supporting GPU and arbitrary grid resolutions).
• Hybrid ML-FVM Framework: Demonstrated a method for "embedding" neural networks into classical solvers, allowing the model to benefit from the stability of traditional physics-based schemes.
• Comparative Analysis of ML Strategies: Systematically evaluated four different ways to integrate ML into the discretization process, identifying which mathematical approaches are stable and which fail.
• Spherical Geometry Integration: Successfully applied data-driven methods to a logically rectangular grid on a sphere, a complex task for standard ML models.

4. Results

• Classical Performance: The custom solvers showed superior performance over Pyclaw across 1D, 2D, and spherical tests. The Roe flux was identified as the most accurate for the Euler equations.
• ML Method Evaluation:
  • Method 1 (Direct Flux): Failed. It failed to maintain momentum and produced physically incorrect results.
  • Method 2 (Linear Coefficients): Unstable. While boundary values were close to reality, the method introduced intermittent discontinuities, leading to numerical instability (NaN values).
  • Method 3 (Direct Boundary States): Feasible. This method produced results that were sharper than low-resolution classical solvers, though it exhibited some "noise" due to training limitations.
  • Method 4 (Slope Prediction): Most Successful. This was the most stable and accurate approach. It produced sharp, non-oscillatory solutions and maintained energy conservation levels comparable to second-order classical solvers.
• Energy Conservation: The successful ML methods (3 and 4) demonstrated stable evolution of total energy and potential enstrophy, proving they respect the underlying physics better than the first two methods.

5. Significance

This work provides a roadmap for the development of Physics-Informed Machine Learning in Computational Fluid Dynamics (CFD). It proves that the most effective way to use ML in geophysical fluid solvers is not to replace the physics, but to use ML to "super-resolve" the local reconstruction (slopes/boundary states). This allows for high-fidelity simulations on coarse, computationally efficient grids, potentially revolutionizing large-scale climate and ocean modeling.