G-Adaptive mesh refinement - leveraging graph neural networks and differentiable finite element solvers

In the following we consider the solution of an abstract PDE on domain $\Omega$ of dimension $d$ in the form

Note we will treat both time-dependent and time-independent problems in the example section, and it is thus understood that $\widetilde{\Omega}=[0,T]\times{\Omega}$ and $\tilde{x}=(t,{x})$ in the time-dependent case (and that $\widetilde{\Omega}={\Omega},\tilde{x}=x$ in the time-independent case). We solve (1) with a finite element method in the weak formulation, which means we commence with a mesh $\mathcal{X}$ and a set of piecewise linear basis functions $S_{\mathcal{X}}$ on that mesh. For us trial and test functions will be the same and a finite element discretisations will typically mean that we seek $U_{\mathcal{X}}\in S_{\mathcal{X}}$ using Galerkin optimality conditions of the form

where $a$ is a bilinear form that is usually found by integration by parts of $(\mathcal{F}(U_{t},\dots),v)_{L^{2}({\Omega})}$. In the case of Poisson’s equation with zero Dirichlet boundary conditions the above simplifies to

This method of solution results in an expansion $U_{\mathcal{X}}=\sum_{i=1}^{N_{x}}\alpha_{\mathcal{X}}\phi_{\mathcal{X},i}$ where $S_{\mathcal{X}}=\{\phi_{\mathcal{X},i}\}_{i=1}^{N_{x}}$. The goal in $r$-adaptive meshing is then to minimise the solution error

with respect to the mesh locations $\mathcal{X}$ without changing the number of meshpoints $N_{x}$. Note we will on occasion suppress the subscript $\mathcal{X}$ from $U$ if the dependence is clear from context.

Relocation based r-adaptivity moves a mesh with a certain topology to another with the same topology, where the mesh points are moved, but their connectivity (and hence the associated data structures) is unaltered. Such methods typically map a fixed mesh in a computational domain (i.e. a representation of the mesh graph) to a moving mesh in the physical domain where the PDE is posed. We denote the mesh points in the physical domain by ${\mathbf{x}}^{(i)}$ and they define a triangulation $\mathcal{T}$ so that

We define the following domains and coordinates: $\Omega_{C}\subseteq\mathbb{R}^{d}$ - “computational" domain eg. the plane, sphere, mapped to $\Omega_{P}\subseteq\mathbb{R}^{d}$ - “physical" domain, with $\bm{\xi}$ - coordinates in the computational domain, $\bm{x}$ - coordinates in the physical domain. To construct an adaptive mesh we consider a differentiable time-dependent deformation map $\mathbf{F}:\Omega_{C}\rightarrow\Omega_{P}$, so that $\bm{X}=\mathbf{F}(\mathbf{\xi},t)$ and $\mathbf{F}{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}(\partial\Omega_{C})=}\partial\Omega_{P}$. If ${\bm{\xi}}^{(i)}$ are the fixed mesh points in the computational domain then $\bm{x}^{(i)}=\mathbf{F}(\bm{\xi}^{(i)},t)$ Assuming the mesh in the computational domain is regular, then determining the (properties of the) mesh in the physical domain, reduces to finding, and analysing, the function ${\mathbf{F}}$.

with ${\mathbf{v}}={\bm{0}}$ when $\bm{x}^{(i)}=\bm{\xi}^{(i)}.$ The choice of velocity function ${\mathbf{v}}$ is critical to the success of such methods, and is often motivated by natural Lagrangian structures of the underlying PDE, [22] Chapter 7. The G-adaptive method in this paper is a velocity based method in which ${\mathbf{v}}$ is derived by calculating the rate of change of the FE solution error. As we outline in section 4 and the supplementary material, G-adaptive methods avoid the mesh tangling often associated with velocity methods.

We consider a graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$, with $N_{x}:=|\mathcal{V}|$ nodes and $\mathcal{E}\subset\mathcal{V}\times\mathcal{V}$ edges. The adjacency matrix $\mathbf{A}$ is defined as $a_{ij}=1$ if $(i,j)\in\mathcal{E}$ and zero otherwise. Node features $\left\{\mathbf{x}_{i}\in\mathbb{R}^{d}_{0}:i\in\mathrm{V}\right\}$ can be represented in matrix notation as $\mathbf{X}\in\mathbb{R}^{N_{x}\times d}$.

A graph neural network (GNN) $\mathcal{M}_{\theta}:\mathbb{R}^{N_{x}\times d_{0}}\times\mathcal{E}\rightarrow\mathbb{R}^{N_{x}\times d_{N}}$ is constructed with layers: $\mathcal{L}_{\theta_{k}}:\mathbb{R}^{N_{x}\times d_{k}}\rightarrow\mathbb{R}^{N_{x}\times d_{k+1}}$ acting node wise as $\mathbf{x}_{i}^{k+1}=\mathcal{L}_{\theta_{k}}(\mathbf{x}_{i}^{k})=\phi_{\theta_{k}}\left(\mathbf{x}_{i}^{k},\sum\limits_{j\in\mathcal{N}_{i}}\varphi_{\theta_{k}}(\mathbf{x}_{j}^{k})\right)$

Two important examples are GCN [25] in matrix form $\mathbf{X}^{k+1}=\sigma(\mathbf{A}\mathbf{X}^{k}\mathbf{W})$ and GAT [41] $\mathbf{X}^{k+1}=\sigma(\mathbf{A}(\mathbf{X}^{k})\mathbf{X}^{k}\mathbf{W})$ where $a_{ij}=a(x_{i},x_{j})$ are attentional coefficients. We can include encoder $E$ and decoder $D$ such that

GRAND [9] is a state-of-the-art GNN architecture that formulates a learnable non-homogeneous (or anisotropic) diffusion equation on a graph. Similar to neural ODEs [chen] it considers layers through the network as time and, by the properties of the row stochastic adjacency matrix, has favourable stability properties in inference and training. Given the differential equation approach to classical meshing strategies the architecture specifically aligns itself with these and brings the power of GNNs to the meshing problem. The GRAND architecture is achieved by making residual and treating propagation with an attentional graph Laplacian the GNN equation 7 becomes

$a\left(\mathbf{X}_{i},\mathbf{X}_{j}\right)$ is calculated using the scaled dot product attention [40]

where $a_{i,j}>0$ if $(i,j)\in\mathcal{E}$ and $\sum_{j\in\mathcal{N}(i)}a_{ij}=1$. To construct the feature matrix required for G-Adaptivity we concatenate coordinates of a regular FE grid $\bm{\xi}\in\mathbb{R}^{N_{x}\times d}$, a preliminary coarse FE solve $\tilde{u}\in\mathbb{R}^{N_{x}}$ and curvature estimate $\tilde{f}\in\mathbb{R}^{N_{x}}$ on the regular grid giving the features vector $\mathbf{X}=\left(\bm{\xi}\mathbin{\|}\tilde{u}\mathbin{\|}\tilde{f}\right)\in\mathbb{R}^{N_{x}\times d+2}$. Then we define our G-adaptive operator

$\mathcal{M}_{\theta}$ uses an identity encoder $E$ to lift the latent space to $\mathbb{R}^{N_{x}\times d+2+|\lambda|}$ and a decoder $D$ to return physical coordinates. GRAND layers are used as convolutional blocks that induce several desirable properties for the meshing problem. By construction the adjacency is row stochastic this guarantees stability properties on the solver step size. Feature channels don’t mix instead they interact via attention expressed with elements of the adjacency, this is important as the first two features are physical coordinates. Finally the learnable weights $Q$ and $K$ matrices allow for an anistropic diffusion which is akin to a learnable monitor function from classical methods.

As mentioned above, prior work for $r$-adaptivity both in classical adaptive meshing (cf. [7]) and in machine learning approaches to the problem (cf. [14, 34, 43, 21, 38]) has built on the assumption that direct minimisation of the FEM error, whilst highly desirable, is both expensive and highly non-convex and therefore computationally infeasible [22, p. 56]. It turns out that recent advances in computational power and software have brought exactly this direct minimisation into the realm of the possible. Indeed, we introduce an adjoint framework which allows us to introduce difFEM a finite element solver which is fully differentiable with respect to mesh point locations. Based on the chain rule we can then exploit this solver to directly compute gradients of the FEM error (3) with respect to the GNN parameters $\bm{\theta}$. In particular we have (writing $\mathcal{M}_{\bm{\theta}}=\mathcal{M}_{\bm{\theta}}(\mathbf{X},\mathcal{A})$ for simplicity)

In order to compute the derivative values of the FEM error $\mathcal{E}$ with respect to the mesh coordinates $\mathcal{X}=\{\mathbf{x}_{i}\}_{i=1}^{N_{x}}$ we begin by recalling that the expansion coefficients $\bm{\alpha}$ of $U_{\mathcal{X}}$ (cf. section 3) are given as the solution of a linear system (the constraint $\mathbf{g}$ arising from the bilinear form (2)):

for a given matrix $\mathsf{M}_{\mathcal{X}}$ and load vector $\bm{b}_{\mathcal{X}}$ which both depend on the mesh. Then we have

where $\bm{\lambda}$ is the solution to the discrete adjoint problem (cf. [39]) $\mathsf{M}_{\mathcal{X}}^{\top}\bm{\lambda}=-\partial\mathcal{E}/{\partial\bm{\alpha}}$. The advantage of this formulation is that we never have to compute $d\bm{\alpha}/d\mathbf{x}_{i}=d\left(\mathsf{M}_{\mathcal{X}}^{-1}\mathbf{b}_{\mathcal{X}}\right)/d\mathbf{x}_{i}$ directly, thus we avoid having to differentiate a matrix inversion. This significantly simplifies the computational tree in the difFEM and makes computations on moderate mesh resolutions sizes feasible.

Our first test problem is the elliptic Poisson’s equation

The favourable performance of our approach extends to time-dynamic problems. To demonstrate this we choose the classical example of the viscous Burgers’ equation which is a well-known example of a PDE that evolves solutions with sharp interfaces (which are resolved with adaptive meshes). In the following $0<\nu\ll 1$ is the viscosity parameter: