A Neural-Stabilised Framework for Accurate Simulation of Atmospheric Flows in General Circulation Models

Abstract

Overview This project implements a Hybrid Discontinuous Galerkin Finite Element to solve the Shallow Water Equations (SWE) on a rotating sphere (an Icosphere grid). It utilizes a Neural Network (MLP) to predict the optimal artificial viscosity required to stabilize the simulation. By offloading the stabilization logic to a neural network, this framework aims to run simulations at coarser resolutions (Level 3) while maintaining stability and accuracy comparable to higher-resolution traditional solvers. Key Features Domain: Geodesic icosahedral grid (Icosphere). Physics: Shallow Water Equations (SWE) with Coriolis force (Rossby-Haurwitz Wave test case). Solver: Explicit Runge-Kutta 2 (RK2) time integration with Rusanov Flux. ML Component: A PyTorch-based Multilayer Perceptron (MLP) trained to predict "Teacher Viscosity" based on local flow features (Gradients, Divergence, Vorticity). 📂 File Structure To run this workflow, organize your code into the following three scripts: 01_generate_data.py: Generates the "Ground Truth" and harvests training data using a high-resolution traditional solver. 02_train_mlp.py: Trains the Neural Network using the harvested data. 03_hybrid_solver.py: Runs the fast, coarse-grid simulation using the trained NN for stabilization and compares it against the Ground Truth. ⚙️ Prerequisites This code is designed to run in Google Colab (due to Google Drive path dependencies) or a local environment with the following libraries: