Learning Differentiable Weak-Form Corrections to Accelerate Finite Element Simulations

TL;DR

The paper addresses accelerating incompressible finite element simulations on coarse grids without sacrificing stability by embedding learned corrections directly in the weak (variational) form. It introduces a weak-form correction framework that learns three coefficient fields, $c_{\mathrm{adv}}$, $\nu_t$, and $\gamma$, which modify the bilinear forms while preserving Galerkin structure, and trains them end-to-end via discrete adjoints coupled to PyTorch within the Firedrake solver. Across 1D convection–diffusion and 2D Navier–Stokes flow past a cylinder, the weak-form approach consistently yields lower rollout errors and improved long-horizon stability compared to strong-form corrections, while maintaining substantial speedups relative to high-fidelity simulations. These results demonstrate a principled, structure-preserving pathway to accurate coarse-grid simulations of incompressible flows and suggest extensions to 3D, broader geometries, and alternative neural architectures. The work highlights the practical potential of differentiable, weak-form learning to enable robust, scalable hybrid physics–ML solvers in CFD contexts.

Abstract

We present a differentiable weak-form learning approach for accelerating finite element simulations. Rather than introducing black-box source terms in the strong form of the governing equations, we augment the momentum equation directly in the variational (weak) form with parameterized bilinear operators. The coefficients of these operators are learned from high-resolution simulations so that unresolved small-scale dynamics can be represented on coarse grids. Applying the correction at the weak-form level aligns the learned model with the finite element discretization, preserving key numerical structure and better respecting the fundamental properties of incompressible flow. In the same setting, the approach yields solutions that are more accurate and more stable over long time horizons than comparable strong-form corrections. We implement the proposed method in the Firedrake finite element solver and evaluate it on benchmark problems, including the one-dimensional convection-diffusion equation and the two-dimensional incompressible Navier-Stokes equations. End-to-end differentiable training is enabled by coupling PyTorch with the Firedrake adjoint framework. Across these tests, the learned variational operators improve long-term accuracy while reducing computational cost. Overall, our results suggest that weak-form learning provides a principled, structure-preserving route to accurate and stable coarse-grid simulations of incompressible flows.

Figures (11)

• Figure 1: 1D CONVECTION-DIFFUSION EQUATION: (a) INITIAL CONDITION AND (b) SOLUTION AT $T=2$.
• Figure 2: SPACE-TIME FIELDS FOR THE 1D CONVECTION-DIFFUSION EQUATION. (a) HIGH-RESOLUTION SOLUTION $u_{H}$. (b) PROJECTED FIELD $\mathsfbi{P}u_{H}$. (c) LOW-RESOLUTION SOLUTION $u_{L}$.
• Figure 3: 2D FLOW PAST A CYLINDER: (a) COMPUTATIONAL MESH COLORED BY THE INSTANTANEOUS VELOCITY MAGNITUDE (ONLY ELEMENT EDGES ARE SHOWN). (b) HIGH-RESOLUTION GRID ($p_v,p_p = 6,5$) WITH ZOOMED-IN MESH INSETS NEAR THE CYLINDER. (c) LOW-RESOLUTION GRID ($p_v,p_p = 2,1$).
• Figure 4: ARCHITECTURE OF CORRECTIVE OPERATOR FOR 1D CONVECTION-DIFFUSION EQUATION.
• Figure 5: TRAINING LOSS VERSUS EPOCH FOR STRONG-FORM AND WEAK-FORM CORRECTIONS IN 1D CONVECTION-DIFFUSION EQUATION.