# Learning Interface Conditions in Domain Decomposition Solvers

Ali Taghibakhshi
Mechanical Science and Engineering
University of Illinois Urbana-Champaign
Urbana, IL 61801, USA
<EMAIL_ADDRESS>
&Nicolas Nytko
Computer Science
University of Illinois Urbana-Champaign
Urbana, IL 61801, USA
<EMAIL_ADDRESS>
&Tareq Zaman
Scientific Computing Program
Memorial University of Newfoundland
and Labrador
St. John’s, NL, Canada
<EMAIL_ADDRESS>
&Scott MacLachlan
Mathematics and Statistics
Memorial University of Newfoundland
and Labrador
St. John’s, NL, Canada
<EMAIL_ADDRESS>
&Luke Olson
Computer Science
University of Illinois Urbana-Champaign
Urbana, IL 61801, USA
<EMAIL_ADDRESS>
&Matthew West
Mechanical Science and Engineering
University of Illinois Urbana-Champaign
Urbana, IL 61801, USA
<EMAIL_ADDRESS>

###### Abstract

Domain decomposition methods are widely used and effective in the
approximation of solutions to partial differential equations. Yet the optimal
construction of these methods requires tedious analysis and is often available
only in simplified, structured-grid settings, limiting their use for more
complex problems. In this work, we generalize optimized Schwarz domain
decomposition methods to unstructured-grid problems, using Graph Convolutional
Neural Networks (GCNNs) and unsupervised learning to learn optimal
modifications at subdomain interfaces. A key ingredient in our approach is an
improved loss function, enabling effective training on relatively small
problems, but robust performance on arbitrarily large problems, with
computational cost linear in problem size. The performance of the learned
linear solvers is compared with both classical and optimized domain
decomposition algorithms, for both structured- and unstructured-grid problems.

##  1 Introduction

Domain decomposition methods (DDMs) Toselli and Widlund (2005); Quarteroni and
Valli (1999); Dolean et al. (2015) are highly effective in solving the linear
and nonlinear systems of equations that arise from the numerical approximation
of solutions to partial differential equations (PDEs). While most effective on
elliptic boundary-value problems, DDMs can also be applied to nonlinear
problems, either using their nonlinear variants, or successively solving
linearizations. Time-dependent problems are normally solved by using a time
stepping algorithm in the time domain for implicit methods, which require the
solution of a spatial problem for each time step. Of these methods, Schwarz
methods are particularly popular given their relative simplicity and ease of
parallelization. The common theme is to break the global problem into
subproblems, derived either by projection or by discretizing the same PDE over
a physical subdomain, and to use solutions on the subdomains as a
preconditioner for the global discretization. Classical Schwarz methods
generally make use of Dirichlet or Neumann boundary conditions for these
subdomain problems, while Optimized Schwarz Methods (OSMs) aim to improve the
convergence of the algorithm by using more general interface conditions Gander
et al. (2000). Notably, Gander and Kwok (2011) demonstrates that optimal, but
non-local, interface conditions exist for more general decompositions.

Much of the OSM literature considers only one-level additive Schwarz methods,
although multilevel extensions do exist. For one-level methods (i.e., domain
decomposition approaches without a “coarse grid”), restricted additive Schwarz
(RAS) approaches Cai and Sarkis (1999) are arguably the most common; optimized
restricted additive Schwarz (ORAS) methods are considered in St-Cyr et al.
(2007). The OSM idea has also been extended to asynchronous Schwarz methods
Magoulès et al. (2017), where the computations on each subdomain are done
using the newest information available in a parallel computing environment
without synchronizing the solves on each subdomain.

With a recent focus on machine learning (ML) techniques for solving PDE
systems Raissi et al. (2019); Li et al. (2020a), there is also effort to apply
learning-based methods to improve the performance of iterative solvers for
PDEs, including DDM and algebraic multigrid (AMG) methods. Within AMG methods,
ML techniques have been applied to learning interpolation operators Greenfeld
et al. (2019); Luz et al. (2020) and to coarse-grid selection in reduction-
based AMG Taghibakhshi et al. (2021). Of particular note here is the loss
function employed in Greenfeld et al. (2019); Luz et al. (2020), where they
use unsupervised learning to train a graph neural network to minimize the
Frobenius norm of the error-propagation operator of their iterative method.
Within DDM, significant effort has been invested in combining ML techniques
with DDM, as in Heinlein et al. (2021), where two main families of approaches
are given: 1) using ML within classical DDM methods to improve convergence,
and 2) using deep neural networks, such as Physics Informed Neural Networks
(PINNs) Raissi et al. (2019), as a discretization module and solver for DDM
problems. In Heinlein et al. (2019), a fully connected neural network is used
to predict the geometric locations of constraints for coarse space
enhancements in an adaptive Finite Element Tearing and Interconnecting-Dual
Primal (FETI-DP) method. Using the continuous formulation of DDM, the so-
called D3M Li et al. (2019) uses a variational deep learning solver,
implementing local neural networks on physical subdomains in a parallel
fashion. Likewise, Deep-DDM Li et al. (2020b) utilizes PINNs to discretize and
solve DDM problems, with coarse space corrections Mercier et al. (2021) being
used to improve scalability.

In this paper, we advance DDM-based solvers by developing a framework for
learning optimized Schwarz preconditioners. A key aspect of this is
reconsidering the loss function to use a more effective relaxation of the
ideal objective function than Frobenius norm minimization Greenfeld et al.
(2019); Luz et al. (2020). Moreover, the approach introduced here offers an
opportunity to reconsider existing limitations of optimized Schwarz methods,
where optimal parameter choice is based on Fourier analysis and requires a
highly regular subdomain structure, such as in the classical cases of square
domains split into two equal subdomains or into regular grids of square
subdomains. Our framework learns the optimized Schwarz parameters via training
on small problem sizes, in a way that generalizes effectively to large
problems, and in a way that allows us to consider both irregular domains and
unstructured grids, with no extraordinary limitations on subdomain shape.
Furthermore, the evaluation time of our algorithm scales linearly with problem
size. This allows significant freedom in defining optimized Schwarz methods,
in comparison to classical approaches, allowing us to explore the potential
benefits of optimized Schwarz over classical (restricted) additive methods on
unstructured grids for the first time.

##  2 Background

Let $\Omega\subset\mathbb{R}^{2}$ be an open set, and consider the positive-
definite Helmholtz problem

$Lu=(\eta-\Delta)u=f\quad\text{in $\Omega$},$ (1)

with inhomogeneous Dirichlet conditions imposed on the boundary
$\partial\Omega$. In (1), the parameter $\eta>0$ represents a shift in the
Helmholtz problem. In the numerical results below, we consider both finite-
difference discretizations of (1) on regular grids, as well as piecewise
linear finite-element (FE) discretizations on arbitrary triangulations. In
both cases, we denote the set of degrees of freedom as $D$, and note that
these are in a one-to-one correspondence with the nodes of the underlying
mesh. Consider a decomposition of $D$ into non-overlapping subdomains
$D_{i}^{0},i\in\\{1,2,\ldots,S\\}$ such that each node is contained within
exactly one subdomain $D_{i}^{0}$, yielding $\cup D_{i}^{0}=D$. In this
subdomain notation, the superscript denotes the amount of overlap in the
subdomains, which is zero for the non-overlapping subdomains that we first
consider. Let $R_{i}^{0}$ be the restriction operator onto the set of degrees
of freedom (DoFs) in $D_{i}^{0}$, and let $\left(R_{i}^{0}\right)^{T}$ be the
corresponding extension operator from $D_{i}^{0}$ into set $D$. Then, an FE
discretization of the Helmholtz problem leads to a linear system of the form
$Ax=b$, where $A$ is the global stiffness matrix and
$A_{i}^{0}=R_{i}^{0}A\left(R_{i}^{0}\right)^{T}$ is the subdomain stiffness
matrix for $D_{i}^{0}$. We note that alternate definitions to the Galerkin
projection for $A_{i}^{0}$ are possible, and are commonly considered in
optimized Schwarz settings (as noted below).

In the case of restricted additive Schwarz (RAS) Cai and Sarkis (1999), the
subdomains are extended to allow for overlap: nodes near the “boundary” of
their subdomain are potentially included in two or more subdomains. We denote
the amount of overlap by $\delta\in\mathbb{N}$, defining subdomains
$D_{i}^{\delta}$ recursively, by
$D_{i}^{\delta}=D_{i}^{\delta-1}\cup\left\\{j\mid a_{kj}\neq 0\text{ and }k\in
D_{i}^{\delta-1}\right\\}$ for $\delta>0$. The conventional RAS preconditioner
is defined as

$M_{\text{RAS}}=\sum\limits_{i=1}^{S}(\tilde{R}^{\delta}_{i})^{T}(A_{i}^{\delta})^{-1}{R_{i}^{\delta}},$
(2)

where $R_{i}^{\delta}$ is the standard restriction operator to subdomain
$\Omega_{i}^{\delta}$, $\tilde{R}_{i}^{\delta}$ is a modified restriction
operator from $D$ to the DoFs in $D_{i}^{\delta}$ that takes nonzero values
only for DoFs in $D_{i}^{0}$, and
$A_{i}^{\delta}={(R_{i}^{\delta})}^{T}AR_{i}^{\delta}$. Figure 1 shows an
example unstructured grid with two subdomains and overlap $\delta=1$.

In an optimized Schwarz setting, we modify the subdomain systems,
$A_{i}^{\delta}$. Rather than using a Galerkin projection onto
$D_{i}^{\delta}$, we rediscretize (1) over the subdomain of $\Omega$
corresponding to these DoFs, imposing a Robin-type boundary condition on the
boundary of the subdomain. We define this matrix to be
$\tilde{A}_{i}^{\delta}=A_{i}+L_{i}$, where $A_{i}$ is the term resulting from
discretization of (1) with Neumann boundary conditions, and $L_{i}$ is
additional the term resulting from the Robin-type condition, as in:

$\displaystyle\text{Dirichlet: }u$ $\displaystyle=g_{\text{D}}(x),$
$\displaystyle\text{Neumann: }\vec{n}\cdot\nabla u$
$\displaystyle=g_{\text{N}}(x),$ $\displaystyle\text{Robin: }\alpha
u+\vec{n}\cdot\nabla u$ $\displaystyle=g_{\text{R}}(x),$ (3)

where $\vec{n}$ is the outward normal to the edge on the boundary and $g$
denotes inhomogeneous data.

The matrix $L_{i}$ has the same dimensions as $A_{i}$, so that
$\tilde{A}_{i}^{\delta}$ is well-defined. However, it has a significantly
different sparsity pattern, with nonzero entries only in rows/columns
corresponding to nodes on the boundary of subdomain $D_{i}^{\delta}$. In
practice, we identify a cycle or path in the graph corresponding to $A_{i}$
with the property that every node in the cycle is on the boundary of
$D_{i}^{\delta}$ but not the boundary of $D$ (the discretized domain), and
then restrict the nonzeros in $L_{i}$ to the entries corresponding to the
edges in this cycle/path (including self-edges, corresponding to entries on
the diagonal of $L_{i}$).

Figure 1: Two subdomains with overlap $\delta=1$ on a 58-node unstructured
grid. The blue and orange shading denotes the original non-overlapping
partitions ($D_{1}^{0}$ and $D_{2}^{0}$), while the blue and orange dashed
outlines show the overlapping subdomains ($D_{1}^{1}$ and $D_{2}^{1}$). Nodes
belonging to only one subdomain are marked with a solid circle, while those in
white belong to both subdomains. The connections in the boundary matrices
$L_{1}$ and $L_{2}$ are denoted by the edges shaded in blue (for $L_{1}$) and
orange (for $L_{2}$).

Using this notation, the ORAS preconditioner can be written as

$M_{\text{ORAS}}=\sum\limits_{i=1}^{S}\left(\tilde{R}_{i}^{\delta}\right)^{T}\left(\tilde{A}_{i}^{\delta}\right)^{-1}R_{i}^{\delta}.$
(4)

Our aim is to learn the values in the matrices $L_{i}$, aiming to outperform
the classical choices of these values predicted by Fourier analysis, in the
cases where those values are known, and to learn suitable values for cases,
such as finite-element discretizations on unstructured grids, where no known
optimized Schwarz parameters exist. We optimize the values for the case of
stationary (Richardson) iteration, but evaluate the performance of the
resulting methods both as stationary iterations and as preconditioners for
FGMRES.

#### Graph Neural Networks (GNNs):

Applying learning techniques to graph structured data necessitates stepping
beyond multilayer perceptron (MLP) and conventional convolutional neural
networks (CNN) to a type of network that leverages the underlying graph nature
of the problem, namely graph convolutional neural networks (GCNNs). GCNNs are
typically divided into two categories: spectral and spatial Wu et al. (2020).
Spectral GCNNs, first introduced by Bruna et al. Bruna et al. (2013), consider
a graph Laplacian eigenbasis and define convolution as diagonal operators. As
such, spectral GCNN methods suffer from time complexity problems due to the
necessity for the eigendecomposition of the Laplacian matrix. Nevertheless, in
follow-up work Defferrard et al. (2016); Kipf and Welling (2016), remedies
have been proposed to mitigate this. Unlike spectral methods, spatial GCNNs
consider local propagation of information in graph domains as a convolution
graph. One popular framework is the message passing neural network (MPNN)
Gilmer et al. (2017), which is based on sharing information among neighbor
nodes in each round of a convolution pass. This has been generalized Battaglia
et al. (2018) by introducing a configurable graph network block structure
consisting of node and edge convolution modules and a global attribute. In an
effort to alleviate computational complexity of GCNNs Du et al. (2017),
topology adaptive graph convolution networks (TAGCN) can be constructed by
defining learnable filters. This is not only computationally simpler, but also
allows for adapting to the topology of the graphs when scanning them for
convolution.

##  3 Method

###  3.1 Optimization problem and the loss function

Throughout this paper, we use $\|\cdot\|$ to denote the $\ell^{2}$ norm of a
matrix or vector, $\|A\|_{F}$ for the Frobenius norm of $A$, and $\rho(A)$ as
the spectral radius of $A$. The optimization problem that we seek to solve is
to find optimal values for the entries in the matrices $L_{i}$, constrained by
given sparsity patterns, to minimize $\rho(T)$, where $T=I-M_{\text{ORAS}}A$
is the error-propagation operator for the stationary iteration corresponding
to $M_{\text{ORAS}}$ defined in (4). The spectral radius $\rho(T)$ corresponds
to the asymptotic convergence factor of the stationary iteration, giving a
bound on the asymptotic convergence of the method. Formally defined in terms
of the extremal eigenvalue of $T^{T}T$ (since $T$ is not symmetric), direct
minimization of $\rho(T)$ is difficult since backpropagation of an
eigendecomposition is numerically unstable Wang et al. (2019). To overcome
this, Luz et al. Luz et al. (2020) propose to relax the minimization of
$\rho(T)$ (for a similar AMG preconditioner) to minimizing the Frobenius norm
of $T$, $\|T\|_{F}$. In our experiments, however, we find that this is
insufficient, leading to preconditioners that do not scale. One reason is that
while the Frobenius norm is an upper bound on the spectral radius, it does not
appear to be a suitably “tight” bound for use in this context (see Section 4.2
and Figure 6). Instead, we use a relaxation inspired by Gelfand’s formula,
that $\rho(T)=\displaystyle\lim_{K\rightarrow\infty}\|T^{K}\|^{\frac{1}{K}}$,
and the common bound that

$\rho(T)\leq\|T^{K}\|^{\frac{1}{K}}=\sup_{x\neq
0}\left(\frac{\|T^{K}x\|}{\|x\|}\right)^{\frac{1}{K}}=\sup_{x:\|x\|=1}(\|T^{K}x\|)^{\frac{1}{K}}$
(5)

for some finite $K\in\mathbb{N}$. This results in the optimization problem

$\min_{\begin{subarray}{c}L_{i},i=1,2,\ldots,S\\\ \text{sparsity of
$L_{i}$}\end{subarray}}\sup_{x:\|x\|=1}\|T^{K}x\|.$ (6)

###  3.2 Numerical evaluation of the loss function

We denote the action of evaluating the GNN by $f^{(\theta)}$ (where $\theta$
represents the network parameters), and consider a discretized problem with
DoF set $D$, of size $n$. The set $D$ can be decomposed into subdomains either
by using fixed geometric choices of the subdomain (e.g., for finite-difference
discretizations), using the METIS graph partitioner Karypis and Kumar (1998),
or a $k$-means-based clustering algorithm (best known as Lloyd’s algorithm
which has $O(n)$ time complexity) Bell (2008); Lloyd (1982). For unstructured
problems, we use a $k$-means-based algorithm, decomposing $D$ to subdomains
$D_{i}$ for $i=1,2,\ldots,S$, with overlap $\delta$; see Supplementary
Materials for details. The GNN then takes $D$ and its decomposition as inputs,
as well as sparsity constraints on the matrices $L_{i}$ for $i=1,2,\ldots,S$,
and outputs values for these matrices:

${L_{1}^{(\theta)},L_{2}^{(\theta)},\ldots,L_{S}^{(\theta)}}\leftarrow
f^{(\theta)}(D).$ (7)

Using the learned subdomain interface matrices, we then obtain the modified
MLORAS (Machine Learning Optimized Restricted Additive Schwarz) operator,
$M^{(\theta)}_{\text{ORAS}}$, simply using
$\tilde{A}_{i}^{\delta}=A_{i}+L_{i}^{(\theta)}$ in (4). We denote the
associated error propagation operator by
$T^{(\theta)}=I-M^{(\theta)}_{\text{ORAS}}A$.

While Gelfand’s formula and the associated upper bound in (5) are valid in any
norm, it is natural to consider them with respect to the $\ell^{2}$ norm in
this setting. However, this raises the same issue as encountered in Luz et al.
(2020), that it generally requires an eigendecomposition to compute the norm.
To avoid this, we use a stochastic sampling of
$\left\|\left(T^{(\theta)}\right)^{K}\right\|$, generated by the sample set
$X\in\mathbb{R}^{n\times m}$ for some $m\in\mathbb{N}$, given as

$X=[x_{1},x_{2},\ldots,x_{m}],\forall_{j}\;x_{j}\sim\mathbb{R}^{n}\;\text{uniformly},\|x_{j}\|=1.$
(8)

Here, we randomly select $m$ points uniformly on a unit sphere in
$\mathbb{R}^{n}$, which can be done using the method in Box (1958). We then
define

$Y^{(\theta)}=\left\\{\left\|\left(T^{(\theta)}\right)^{K}x_{1}\right\|,\left\|\left(T^{(\theta)}\right)^{K}x_{2}\right\|,\ldots,\left\|\left(T^{(\theta)}\right)^{K}x_{m}\right\|\right\\},$
(9)

taking each column of $X$ as the initial guess to the solution of the
homogeneous problem $Ax=0$ and taking $K$ steps of the stationary algorithm to
generate $\left(T^{(\theta)}\right)^{K}x_{j}$. Since we normalize each column
of $X$ to have $\|x_{j}\|=1$, the value of
$\left\|\left(T^{(\theta)}\right)^{K}x_{j}\right\|$ serves as a lower bound
for $\left\|\left(T^{(\theta)}\right)^{K}\right\|$. Thus, taking the maximum
of the values in $Y$ provides a practical loss function that we use below,
defining

$\mathcal{L}^{(\theta)}=\max(Y^{(\theta)}).$ (10)

We note similarities between this loss function and that used in Katrutsa et
al. (2020), but that we are able to use the maximum of $Y^{(\theta)}$ (giving
a better approximation to the norm) in our context instead of averaging.

Ultimately, the cost of our algorithm depends strongly on the chosen values of
$m$ and $K$. For sufficiently large values of $m$, we now show that the
maximum value in $Y^{(\theta)}$ is an arbitrarily good approximation to
$\left\|\left(T^{(\theta)}\right)^{K}\right\|^{\frac{1}{K}}$ (in the
statistical sense).

######  Theorem 1.

For any nonzero matrix $T$, $\epsilon>0$, and $\delta<1$, there exist
$M,K\in\mathbb{N}$ such that for any $m>M$, if one chooses $m$ points,
$x_{j}$, uniformly at random from $\\{x\in\mathbb{R}^{n},\|x\|=1\\}$, then
$Y=\left\\{\left\|T^{K}x_{1}\right\|,\left\|T^{K}x_{2}\right\|,\ldots,\left\|T^{K}x_{m}\right\|\right\\}$
satisfies

$P\left(\left|\rho(T)-\max(Y)^{\frac{1}{K}}\right|<\epsilon\right)>1-\delta.$
(11)

###### Proof.

According to Gelfand’s theorem, there exists $L\in\mathbb{N}$ such that
$\forall{\ell>L}$,
$\left|\rho(T)-\displaystyle\sup_{x:\|x\|=1}\|T^{\ell}x\|^{\frac{1}{\ell}}\right|<\frac{\epsilon}{2}$.
Take any $K\geq L$ and let
$\tilde{\epsilon}=\frac{\epsilon}{2\|T^{K}\|^{\frac{1}{K}}}$. Since
$\mathbb{R}^{n}$ is finite-dimensional, there exists an
$x^{*}\in\mathbb{R}^{n},\;\|x^{*}\|=1$ such that
$\displaystyle\sup_{x:\|x\|=1}\|T^{K}x\|=\|T^{K}x^{*}\|$. Denote the volume of
the surface of the $n$-dimensional sphere of unit radius around the origin in
$\mathbb{R}^{n}$ by $C_{\text{tot}}$, and the volume of the region on this
sphere within radius $\tilde{\epsilon}^{K}$ of $x^{*}$ by
$C_{\tilde{\epsilon},K}$.

Given $\delta<1$, let $M\in\mathbb{N}$ satisfy
$M\geq\frac{\log(\delta)}{\log\left(1-\frac{C_{\tilde{\epsilon},K}}{C_{\text{tot}}}\right)}$.
Then,

$P\left(\|x^{*}-x_{i}\|>\tilde{\epsilon}^{K}\text{ for all
}i\right)<\left(1-\frac{C_{\tilde{\epsilon},K}}{C_{\text{tot}}}\right)^{M}\leq\delta.$
(12)

Thus, with probability of at least $1-\delta$, we expect at least one point
from a selection of $m>M$ points uniformly distributed on the sphere of radius
unit radius to be in $C_{\tilde{\epsilon},K}$. Let that point be $x_{r}$,
giving

$\displaystyle\left\|T^{K}x^{*}\right\|^{\frac{1}{K}}-\left\|T^{K}x_{r}\right\|^{\frac{1}{K}}$
$\displaystyle\leq\left(\left\|T^{K}x^{*}\right\|-\left\|T^{K}x_{r}\right\|\right)^{\frac{1}{K}}$
(13)
$\displaystyle\leq\left(\left\|T^{K}(x^{*}-x_{r})\right\|\right)^{\frac{1}{K}}$
(14)
$\displaystyle\leq\left\|T^{K}\right\|^{\frac{1}{K}}\|x^{*}-x_{r}\|^{\frac{1}{K}}\leq\|T^{K}\|^{\frac{1}{K}}\tilde{\epsilon}=\frac{\epsilon}{2}$
(15)

using Lemma 3 (see Supplementary Materials) and the reverse triangle
inequality. Finally, by the triangle inequality, with probability of at least
$1-\delta$, we have:

$\left|\rho(T)-\text{max}(Y)^{\frac{1}{K}}\right|\leq\left|\rho(T)-\sup_{x:\|x\|=1}\|T^{K}x\|^{\frac{1}{K}}\right|+\left|\sup_{x:\|x\|=1}\|T^{K}x\|^{\frac{1}{K}}-\text{max}(Y)^{\frac{1}{K}}\right|\leq\epsilon.$
(16)

∎

######  Remark.

According to Loukas (2019), since the optimal interface values are Turing-
computable, if the depth of the GNN is at least the diameter of the graph, and
a TAGConv layer followed by a feature encoder is Turing-complete, the optimal
interface values can be learned. In our setting, for a problem on a structured
grid of size $N\times N$ with two identical rectangular subdomains, this
implies that the GNN will be able to learn the optimal interface values given,
if and only the GNN has depth at least $2N$, has deep enough feature encoders,
and the width of the layers is unbounded.

######  Remark.

Theorem 1 guarantees convergence of the loss function to the spectral radius,
in the limits of many samples and many stationary iterations. To the best of
our knowledge, such a guarantee is not known for the previous loss functions
used in the area Luz et al. (2020); Greenfeld et al. (2019). Moreover, there
are substantial improvements in the numerical results using the new loss
function in comparison to that of Luz et al. (2020), as shown in Figure 6.

######  Theorem 2.

Assuming bounded subdomain size, the time complexity to evaluate the optimal
Schwarz parameters using our method is $O(n)$, where $n$ is the number of
nodes in the grid.

###### Proof.

Given bounded Lloyd subdomain size and fixed number of Lloyd aggregation
cycles, subdomain generation has $O(n)$ time complexity Bell (2008) (see the
Supplementary Material). To evaluate each TAGConv layer, one computes
$y=\sum_{\ell=1}^{L}G_{\ell}x_{\ell}+b\mathbf{1}_{n}$, where $L$ is the number
of node features, $G_{\ell}\in\mathbb{R}^{n\times n}$ is the graph filter, $b$
is a learnable bias, and $x_{\ell}\in\mathbb{R}^{n}$ are the node features.
Moreover, the graph filter is a polynomial in the adjacency matrix $M$ of the
graph: $G_{\ell}=\sum_{j=0}^{J}g_{\ell,j}M^{j}$ where $J$ is a constant and
$g_{\ell,j}$ are the coefficients of filter polynomial. Since the graph has
bounded node degrees, it implies that $M$ is sparse and the action of $M^{j}$
has $O(n)$ cost, and therefore, the full TAGConv evaluation also has $O(n)$
cost. Moreover, the cost of edge feature and the feature networks are $O(n)$,
resulting in overall $O(n)$ cost. ∎

###  3.3 Training Details

We use a graph neural network based on four TAGConv layers and ResNet node
feature encoders consisting of eight blocks after each TAGConv layer; see
Supplementary Materials for more details on the structure of the GNN, and
Section 4.3 for an ablation study. The training set in our study consists of
1000 unstructured grids with piecewise linear finite elements and with grids
ranging from 90–850 nodes (and an average of 310 nodes). The grids are
generated by choosing either a regular grid (randomly selected 60% of the
time) or a randomly generated convex polygon; pygmsh Schlömer (2021) is used
to generate the mesh on the polygon interior. A sample of the training grids
is shown in Figure 2. We train the GNN for four epochs with a mini batch size
of 25 using the ADAM optimizer Kingma and Ba (2014) with a fixed learning rate
of $10^{-4}$. For the numerical evaluation of the loss function (10) we use
$K=4$ iterations and $m=500$ samples. The code111All code and data for this
paper is at https://github.com/compdyn/learning-oras (MIT licensed). was
implemented using PyTorch Geometric Fey and Lenssen (2019), PyAMG Bell et al.
(2022), and NetworkX Hagberg et al. (2008). All training was performed on an
8-core i9 Macbook Pro using CPU only.

Figure 2: Example grids from the training set.

Moreover, we train two special networks for use in Section 4.1. The first is a
“Brute Force” network, which is a single-layer neural network without an
activation function, trained only on a structured $10\times 10$ grid with two
identical subdomains using the ADAM optimizer. The purpose of training this is
to obtain the optimal interface values as a benchmark for comparison,
including against our method. The second is the same network used for our
method, MLORAS, but overtrained on only the problem in Section 4.1, to
understand the learning capabilities of the method and choice of GNN
architecture.

##  4 Results

###  4.1 Two-subdomain structured grids

We first consider rectangular structured grids with two subdomains. Although
restrictive, these problems allow us to directly compare to existing OSM
parameters, which are only available for structured grids and exactly two
subdomains. We follow St-Cyr et al. (2007) and consider the Helmholtz problem
(1) on the unit square with $\eta=1$ and homogeneous Dirichlet boundary
conditions. We discretize the problem on an $N\times N$ rectangular grid with
$h=1/(N+1)$ grid spacing, using the standard five-point finite difference
stencil.

|
---|---
Figure 3: Results for the Helmholtz problem on a $10\times 10$ structured grid
(left) with two identical subdomains, with convergence plots for the methods
used as preconditioners for FGMRES (center) and as stationary algorithms
(right).

Figure 3 shows the results of different methods on a $10\times 10$ structured
grid with two identical subdomains. We see that the overtrained MLORAS network
learns interface parameters that outperform the Restrictive Additive Schwarz
(RAS) Cai and Sarkis (1999) method; more significantly, MLORAS also
outperforms the existing Optimized-RAS (ORAS) methods that were analytically
derived for this specific problem (zeroth-order OO0 and second-order OO2 from
St-Cyr et al. (2007)). Moreover, in these results, the MLORAS network almost
reaches the performance of the “Brute Force” network (obtained by directly
optimizing all interface values for this single structured grid), indicating
that using a GNN to encode the optimal interface values does not significantly
restrict the search space.

Figure 4: Interface values for the $10\times 10$ structured grid with two
identical subdomains. From left to right: brute force optimization of the
interface values, MLORAS, and second-order ORAS (OO2) St-Cyr et al. (2007).

To understand the performance of the methods in more detail, Figure 4 plots
the interface values output by the brute force optimization, the MLORAS
network, and the OO2 ORAS algorithm. This shows that the MLORAS network is
choosing interface values very close to those directly optimized by the brute
force method, unlike those selected by the OO2 method.

###  4.2 Unstructured grids

To evaluate the performance of the MLORAS network on unstructured grids, we
consider 16 unstructured triangular grids in 2D with sizes ranging from about
90 to 40 000 nodes. These grids are defined on convex subsets of
$(0,1)\times(0,1)$; we solve the Helmholtz problem (1) on these domains with
$\eta=1$ and homogeneous Dirichlet boundary conditions, discretized with
piecewise-linear finite elements.

Example convergence plots are shown in Figure 5, where our method (MLORAS) is
compared to RAS (Restricted Additive Schwarz Cai and Sarkis (1999)), AS
(Additive Schwarz), and Jacobi methods as a preconditioner for FGMRES. Here we
see that the MLORAS network is able to learn optimized interface parameters
for unstructured grids that outperform RAS, and that MLORAS can scale to
problems that are much larger than those in the training set, which are all
below $n=1000$ nodes. Importantly, MLORAS retains an advantage over RAS even
as the grid size increases.

Figure 5: Example convergence on smaller (left) and larger (right)
unstructured grids. Figure 6: Convergence on all unstructured grids in the
testing set. Left: the number of preconditioned FGMRES steps required to solve
the problem to within a relative error of $10^{-12}$. Center: error reduction
by the stationary iteration after 10 iterations. Right: Frobenius norm for
each method on each test problem.

Figure 6 shows the performance of three different methods across all
unstructured test grids. The three methods are: (1) “F-norm optimized”: the
same network as MLORAS, but trained to directly optimize the Frobenius norm of
the error propagation matrix, $\|A\|_{F}$, (2) the RAS Cai and Sarkis (1999)
method, and (3) the MLORAS method. We do not show ORAS results here, as it
cannot be applied to unstructured grids. Figure 6 reveals two important facts.
First, MLORAS consistently outperforms RAS over the entire testing set, both
as an FGMRES preconditioner and as a stationary algorithm, and this advantage
is maintained even for large grids. Second, we see that the Frobenius norm is
indeed a worse choice for the loss function than our new loss in (6). We see
this from the fact that MLORAS has a worse Frobenius norm than either of the
other two methods, but it has the best convergence rate. In addition, when we
explicitly optimize the Frobenius norm (the “F-norm optimized” method), we see
that we do obtain the lowest Frobenius norm, but this translates to the worst
convergence rate.

###  4.3 Ablation study

To understand the impact of the network structure on performance, we conduct
an ablation study by varying the ResNet block length and the number of TAGConv
layers. In each case, we train a network on five different training sets, each
consisting of 1000 grids generated as described in Section 3.3. The trained
networks are tested on 50 unstructured grids, each with 2400 to 2600 nodes.
For each architecture, the mean performance is computed, together with error
bars from the fivefold repetition, with results shown in Figure 7. We see that
a higher residual block length is always better, but that the optimal number
of TAGConv layers is 4, and using more layers actually decreases performance.
This phenomenon has been observed in other contexts (e.g., Zhao and Akoglu
(2019) and Chen et al. (2020)) and can be attributed to over-smoothing in
GNNs. The use of residual blocks in our architecture is, thus, important to
allow us to increase network depth without needing to increase the number of
GNN layers.

Figure 7: Ablation study results. The left panel varies the residual block
length while keeping the number of TAGConv layers fixed at 4. The right panel
varies the number of TAGConv layers while keeping the residual block length
fixed at 8. The solid blue lines (left axis) show the average number of FGMRES
steps needed to reduce the relative error below $10^{-12}$, while the dashed
green lines (right axis) show the average stationary algorithm error reduction
after 10 iterations. The red circles mark the results for the network with the
best performance (residual block length of 8 and 4 TAGConv layers), which was
used for all other studies in this paper. Error bars show one standard error
of the mean.

##  5 Conclusion

We propose an unsupervised learning method based on Graph Convolutional Neural
Networks (GCNNs) for extending optimized restricted additive Schwarz (ORAS)
methods to multiple subdomains and unstructured grid cases. Our method is
trained with a novel loss function, stochastically minimizing the spectral
radius of the error propagation operator obtained using the learned interface
values. The time complexity of evaluating the loss function, as well as
obtaining the interface values using our neural network are all linear in
problem size. Moreover, the proposed method is able to outperform ORAS, both
as a stationary algorithm and preconditioner for FGMRES on structured grids
with two subdomains, as considered in the conventional ORAS literature. On
more general cases, such as unstructured grids with arbitrary subdomains of
bounded size, our method outperforms RAS consistently, both as a stationary
algorithm and preconditioner for FGMRES. The main limitations of the current
work are that the method was studied for two PDE cases, namely the Helmholtz
problem in the main paper and the non-uniform Poisson problem in Appendix D.
We also defer the study of nonlinear or time-dependent PDEs to future work.

## Acknowledgement

The work of S.P.M. was partially supported by an NSERC Discovery Grant. The
authors thank the referees for their insightful comments. The authors have no
competing interests to declare.

## References

  * Toselli and Widlund [2005] Andrea Toselli and Olof Widlund.  _Domain decomposition methods—algorithms and theory_ , volume 34 of _Springer Series in Computational Mathematics_.  Springer-Verlag, Berlin, 2005.  ISBN 3-540-20696-5.  doi: 10.1007/b137868.
  * Quarteroni and Valli [1999] Alfio Quarteroni and Alberto Valli.  _Domain decomposition methods for partial differential equations_.  Numerical Mathematics and Scientific Computation. The Clarendon Press, Oxford University Press, New York, 1999.  ISBN 0-19-850178-1.  Oxford Science Publications.
  * Dolean et al. [2015] Victorita Dolean, Pierre Jolivet, and Frédéric Nataf.  _An introduction to domain decomposition methods_.  Society for Industrial and Applied Mathematics (SIAM), Philadelphia, PA, 2015.  ISBN 978-1-611974-05-8.  doi: 10.1137/1.9781611974065.ch1.
  * Gander et al. [2000] M.J. Gander, L. Halpern, and F. Nataf.  Optimized Schwarz methods.  In _Proceedings of the 12th International Conference on Domain Decomposition_ , pages 15–27. ddm.org, 2000.
  * Gander and Kwok [2011] Martin J. Gander and Felix Kwok.  Optimal interface conditions for an arbitrary decomposition into subdomains.  In _Domain decomposition methods in science and engineering XIX_ , volume 78 of _Lect. Notes Comput. Sci. Eng._ , pages 101–108. Springer, Heidelberg, 2011.  doi: 10.1007/978-3-642-11304-8\\_9.
  * Cai and Sarkis [1999] Xiao-Chuan Cai and Marcus Sarkis.  A restricted additive Schwarz preconditioner for general sparse linear systems.  _Siam Journal on Scientific Computing_ , 21(2):792–797, 1999.
  * St-Cyr et al. [2007] Amik St-Cyr, Martin J Gander, and Stephen J Thomas.  Optimized multiplicative, additive, and restricted additive Schwarz preconditioning.  _SIAM Journal on Scientific Computing_ , 29(6):2402–2425, 2007.
  * Magoulès et al. [2017] Frédéric Magoulès, Daniel B. Szyld, and Cédric Venet.  Asynchronous optimized Schwarz methods with and without overlap.  _Numer. Math._ , 137(1):199–227, 2017.  ISSN 0029-599X.  doi: 10.1007/s00211-017-0872-z.
  * Raissi et al. [2019] Maziar Raissi, Paris Perdikaris, and George E Karniadakis.  Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations.  _Journal of Computational Physics_ , 378:686–707, 2019\.
  * Li et al. [2020a] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Fourier neural operator for parametric partial differential equations.  _arXiv preprint arXiv:2010.08895_ , 2020a.
  * Greenfeld et al. [2019] Daniel Greenfeld, Meirav Galun, Ronen Basri, Irad Yavneh, and Ron Kimmel.  Learning to optimize multigrid PDE solvers.  In _International Conference on Machine Learning_ , pages 2415–2423. PMLR, 2019.
  * Luz et al. [2020] Ilay Luz, Meirav Galun, Haggai Maron, Ronen Basri, and Irad Yavneh.  Learning algebraic multigrid using graph neural networks.  In _International Conference on Machine Learning_ , pages 6489–6499. PMLR, 2020.
  * Taghibakhshi et al. [2021] Ali Taghibakhshi, Scott MacLachlan, Luke Olson, and Matthew West.  Optimization-based algebraic multigrid coarsening using reinforcement learning.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Heinlein et al. [2021] Alexander Heinlein, Axel Klawonn, Martin Lanser, and Janine Weber.  Combining machine learning and domain decomposition methods for the solution of partial differential equations—a review.  _GAMM-Mitteilungen_ , 44(1):e202100001, 2021.
  * Heinlein et al. [2019] Alexander Heinlein, Axel Klawonn, Martin Lanser, and Janine Weber.  Machine learning in adaptive domain decomposition methods—predicting the geometric location of constraints.  _SIAM Journal on Scientific Computing_ , 41(6):A3887–A3912, 2019.
  * Li et al. [2019] Ke Li, Kejun Tang, Tianfan Wu, and Qifeng Liao.  D3M: A deep domain decomposition method for partial differential equations.  _IEEE Access_ , 8:5283–5294, 2019.
  * Li et al. [2020b] Wuyang Li, Xueshuang Xiang, and Yingxiang Xu.  Deep domain decomposition method: Elliptic problems.  In _Mathematical and Scientific Machine Learning_ , pages 269–286. PMLR, 2020b.
  * Mercier et al. [2021] Valentin Mercier, Serge Gratton, and Pierre Boudier.  A coarse space acceleration of deep-DDM.  _arXiv preprint arXiv:2112.03732_ , 2021.
  * Wu et al. [2020] Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Philip.  A comprehensive survey on graph neural networks.  _IEEE Transactions on Neural Networks and Learning Systems_ , 2020\.
  * Bruna et al. [2013] Joan Bruna, Wojciech Zaremba, Arthur Szlam, and Yann LeCun.  Spectral networks and locally connected networks on graphs.  _arXiv preprint arXiv:1312.6203_ , 2013.
  * Defferrard et al. [2016] Michaël Defferrard, Xavier Bresson, and Pierre Vandergheynst.  Convolutional neural networks on graphs with fast localized spectral filtering.  _arXiv preprint arXiv:1606.09375_ , 2016.
  * Kipf and Welling [2016] Thomas N Kipf and Max Welling.  Semi-supervised classification with graph convolutional networks.  _arXiv preprint arXiv:1609.02907_ , 2016.
  * Gilmer et al. [2017] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl.  Neural message passing for quantum chemistry.  In _International Conference on Machine Learning_ , pages 1263–1272. PMLR, 2017.
  * Battaglia et al. [2018] Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al.  Relational inductive biases, deep learning, and graph networks.  _arXiv preprint arXiv:1806.01261_ , 2018.
  * Du et al. [2017] Jian Du, Shanghang Zhang, Guanhang Wu, José MF Moura, and Soummya Kar.  Topology adaptive graph convolutional networks.  _arXiv preprint arXiv:1710.10370_ , 2017.
  * Wang et al. [2019] Wei Wang, Zheng Dang, Yinlin Hu, Pascal Fua, and Mathieu Salzmann.  Backpropagation-friendly eigendecomposition.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Karypis and Kumar [1998] George Karypis and Vipin Kumar.  A fast and high quality multilevel scheme for partitioning irregular graphs.  _SIAM Journal on scientific Computing_ , 20(1):359–392, 1998.
  * Bell [2008] William N Bell.  _Algebraic multigrid for discrete differential forms_.  PhD thesis, University of Illinois at Urbana-Champaign, 2008.
  * Lloyd [1982] Stuart Lloyd.  Least squares quantization in PCM.  _IEEE Transactions on Information Theory_ , 28(2):129–137, 1982.
  * Box [1958] George EP Box.  A note on the generation of random normal deviates.  _Ann. Math. Statist._ , 29:610–611, 1958.
  * Katrutsa et al. [2020] Alexandr Katrutsa, Talgat Daulbaev, and Ivan Oseledets.  Black-box learning of multigrid parameters.  _Journal of Computational and Applied Mathematics_ , 368:112524, 2020.  ISSN 0377-0427.  doi: https://doi.org/10.1016/j.cam.2019.112524.
  * Loukas [2019] Andreas Loukas.  What graph neural networks cannot learn: depth vs width.  _arXiv preprint arXiv:1907.03199_ , 2019.
  * Schlömer [2021] Nico Schlömer.  pygmsh: A Python frontend for Gmsh, 2021.  URL https://github.com/nschloe/pygmsh.  (GPL-3.0 License).
  * Kingma and Ba [2014] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Fey and Lenssen [2019] Matthias Fey and Jan E. Lenssen.  Fast graph representation learning with PyTorch Geometric.  In _ICLR Workshop on Representation Learning on Graphs and Manifolds_ , 2019.  (code is MIT licensed).
  * Bell et al. [2022] Nathan Bell, Luke N. Olson, and Jacob Schroder.  PyAMG: Algebraic multigrid solvers in python.  _Journal of Open Source Software_ , 7(72):4142, 2022.  doi: 10.21105/joss.04142.  URL https://doi.org/10.21105/joss.04142.  (code is MIT licensed).
  * Hagberg et al. [2008] Aric Hagberg, Pieter Swart, and Daniel S Chult.  Exploring network structure, dynamics, and function using networkx.  Technical report, Los Alamos National Lab.(LANL), Los Alamos, NM (United States), 2008.  (code is BSD licensed).
  * Zhao and Akoglu [2019] Lingxiao Zhao and Leman Akoglu.  Pairnorm: Tackling oversmoothing in GNNs.  _arXiv preprint arXiv:1909.12223_ , 2019.
  * Chen et al. [2020] Deli Chen, Yankai Lin, Wei Li, Peng Li, Jie Zhou, and Xu Sun.  Measuring and relieving the over-smoothing problem for graph neural networks from the topological view.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, pages 3438–3445, 2020.
  * Cormen et al. [2001] Thomas H Cormen, Charles E Leiserson, Ronald L Rivest, and Clifford Stein.  _Introduction to algorithms (Second edition)_.  MIT Press, 2001.

##  Appendix A Useful lemma

This lemma is useful in the proof of Theorem 1. While its proof is quite
simple, we are not aware of the result in the literature, and include it here
for completeness.

######  Lemma 3.

For $x,y\in\mathbb{R}$, with $0\leq y\leq x$ and any $K\in\mathbb{N}$,
$x^{\frac{1}{K}}-y^{\frac{1}{K}}\leq(x-y)^{\frac{1}{K}}$

###### Proof.

Since $y\leq x$, the binomial theorem gives that

$x\leq(x-y)+y+\sum\limits_{i=1}^{K-1}\binom{K}{i}y^{\frac{i}{K}}(x-y)^{\frac{K-i}{K}}=(y^{\frac{1}{K}}+(x-y)^{\frac{1}{K}})^{K}.$
(17)

Taking the $K^{\text{th}}$ root of both sides and rearranging gives the stated
result. ∎

##  Appendix B Solution plots

We have compared the solution plot of our method with RAS for the Holmholtz
problem. We consider a $100\times 100$ structured grid on the
$(0,1)\times(0,1)$ domain, and consider the following true solution for the
problem: $u^{*}=\text{sin}(8\pi x)+\text{sin}(8\pi y)$. We then start with an
initial random guess with $L_{2}$ norm of 1. We apply 10 iterations of RAS and
MLORAS (ours) as stationary algorithm and obtain solution for each method.
Moreover, we run 10 iterations of FGMRES with MLORAS and RAS preconditioners
on the initial guess and obtain predictions for both methods. The results are
shown in Figure 8 for MLORAS and Figure 9 for RAS.

Figure 8: MLORAS (ours) solution plots. Top left: ground truth, top middle:
MLORAS stationary solution after 10 iterations, top right: FGMRES solution
with MLORAS preconditioner after 10 steps, bottom left: initial guess, bottom
middle: error of MLORAS stationary solution ($L_{2}$ norm of error = 0.231),
bottom right: error of FGMRES with MLORAS preconditioner solution ($L_{2}$
norm of error = 0.084). Figure 9: RAS solution plots. Top left: ground truth,
top middle: RAS stationary solution after 10 iterations, top right: FGMRES
solution with RAS preconditioner after 10 steps, bottom left: initial guess,
bottom middle: error of RAS stationary solution ($L_{2}$ norm of error =
6.526), bottom right: error of FGMRES with RAS preconditioner solution
($L_{2}$ norm of error = 0.146).

##  Appendix C Neural network

###  C.1 Inputs and Output

#### Inputs:

The GNN takes the grid $G$ as its input which has three main components,
namely $D_{\text{node}}$ (node feature matrix), $D_{\text{edge}}$ (edge
feature matrix), and $A$ (the graph adjacency matrix). Every node has a binary
feature, and its value is one if it is on a boundary of a subdomain and zero
otherwise. Therefore, for every node, the corresponding element in
$D_{\text{node}}$ indicates whether that point is on a boundary or not. In
other words, the binary node feature determines the grid decomposition.
$D_{\text{edge}}$ consists of the edge values obtained from discretization of
the underlying PDE.

#### Outputs:

After passing the input to the GNN architecture which consists of node and
edge convolution blocks and is fully described in the following subsection,
the learned edge weights are obtained for every edge in the grid. However,
only the edges between the nodes on the subdomain boundaries (considering
self-loops) are of our interest so we mask the rest of the edges. Therefore,
the output of the GNN is the learned values for edges connecting nodes on the
boundary of a subdomain which are referred to as interface values in the
paper. For the $i$-th subdomain, the interface value matrix is referred to as
$L_{i}^{\theta}$, where $\theta$ represents the GNN learnable parameters (see
Equation 7). Figure 3 shows an example of the sparsity pattern of an $L$
matrix for each of the identical subdomains of the $10\times 10$ structured
grid in Figure 4. Also note that the interface values are the nonzero elements
of the corresponding $L$ matrix.

###  C.2 Architecture

The overall architecture of the GNN is shown in Figure 10. The GNN takes a
graph as its input and sends node and edge features to the node convolution
and edge feature preprocessing blocks, respectively, both of which are shown
in Figure 11. Each node convolution block consists of a TAGConv layer with
2-size filters and 128 hidden units, followed by a ReLU activation, an
instance norm layer, and a feature network block. The feature network block,
as shown in Figure 12, consists of 8 blocks with residual connections between
each; furthermore, each of the blocks consists of a layer norm followed by a
fully connected layer of size 128, followed by a ReLU activation, and finally
another fully connected layer of size 128. The edge feature preprocessing
block takes the edge features and applies fully connected layers, ReLU
nonlinearities, and instance normalization, before the graph convolution pass.

Figure 10: Overall GNN architecture.

Figure 11: Left: Node convolution block. Right: Edge feature preprocessing
block.  Figure 12: Feature network blocks.

After passing through all of the node convolution blocks, the edge and node
features are sent to a stack block. This block simply stacks node features
onto the edges adjacent to each node. For every edge, $(u,v)$, where $u$ and
$v$ are the nodes on that edge, denote the node and edge features by
$E_{u},E_{v},E_{(u,v)}$, which are the inputs to the stack block. The block
then stacks these features and outputs them as the new edge features for the
edge $(u,v)$. Following the stack layer is the edge convolution block,
depicted in Figure 13, which takes the stacked edge and node features and
passes them through a series of fully connected layers, ReLU activation
functions, and layer norms. It is noteworthy that the size of the input to the
edge convolution block is 257, since following the description of the stack
block, the two node features, each of size 128, and the edge feature of size 1
are stacked together. The output from the edge convolution block is, finally,
passed through a masking block that outputs the interface values. This masks
the interior edge values, i.e., takes the output of the GNN (one value per
each edge in the graph) and multiplies those values that are not on the
boundary of a subdomain by zero, to restrict the output from the GNN to the
desired edges in the graph.

Figure 13: Edge convolution block.

##  Appendix D Poisson problem with discontinuous diffusion coefficient

We also consider the 2D Poisson problem with discontinuous diffusion
coefficient which is formulated as follows:

$-\divergence\kappa(x,y)\nabla
u=f\quad\text{in}\;\Omega,\;\;\;\kappa(x,y)=\begin{cases}1000&0<x<0.5\\\
1&0.5\leq x<1.\end{cases}$ (18)

where $\Omega$ is, as before, defined as a convex subset of $(0,1)\times(0,1)$
and $\kappa(x,y)$ is the discontinuous diffusion coefficient. For this
problem, we consider nine domains with unstructured triangular grids, with
sizes ranging from about 100 to over 30k nodes. The subdomains are generated
using Lloyd aggregation (fully explained in Appendix E), with a fixed ratio of
0.015. We note that the subdomains are not constrained in any way, and a
single subdomain may contain parts of the domain with different diffusion
coefficients. The results of our method compared with RAS baseline for both
the stationary algorithm and the FGMRES preconditioner are shown in Figure 14,
and show little qualitative difference with the earlier results for Helmholtz.

Figure 14: Discontinuous diffusion coefficient for Poisson problem on various
size grids. Left: the number of preconditioned FGMRES steps required to solve
the problem to within a relative error of $10^{-12}$. Right: error reduction
by the stationary iteration after 10 iteration.

##  Appendix E Lloyd aggregation

Lloyd’s algorithm Lloyd [1982] is a standard approach for partitioning data,
closely related to $k$-means clustering, that can be used (for example) to
find close approximations to centroidal Voroni tesselations. Here, we use a
modified form of Lloyd’s algorithm, known as Lloyd aggregation Bell [2008], to
partition a given set of degrees of freedom, $V$, into the non-overlapping
subdomains, $V_{i}^{0}$, for $i\in\\{1,2,\ldots,S\\}$, needed as input to our
algorithm. (Note that, here, we change notation from the paper and use $V$ to
denote the vertices in the graph associated with the matrix, rather than $D$
to denote the index set of degrees of freedom.) Consider a 2D planar graph,
$G$, the set of all edges $E$, the set of all of its nodes $V$, and
$V_{c}\subseteq V$. The nodes in $V_{c}$ serve as the centers of “clusters” in
the graph that will define the subdomains, $V_{i}^{0}$. These regions are
obtained based on the closest center to each graph node, where the distance is
measured by the number of edges covered in the shortest path between two nodes
(denote distance in the graph between node $i$ and $j$ by $g_{ij}$). Define
the centroid of a region as the farthest node from the boundary, breaking
possible ties by random choice. We use a modified version of the Bellman-Ford
algorithm, commonly used for obtaining the nearest center to every node in $V$
and its associated distance Cormen et al. [2001]. Let $\vec{n}$ be a list of
graph nodes whose $i$-th element is the nearest center to the $i$-th node of
the graph, and let $g_{j}$ be the graph distance from node $j$ to $n_{j}$;
then, the modified Bellman-Ford algorithm is shown in Algorithm 1.

Algorithm 1 Modified Bellman-Ford

1:Input $E$: The set all edges, $V$: The set of all nodes, $V_{c}$: The set of
initial center nodes.

2:$g_{i}=\infty\;\;\forall_{i\in\\{1,2,...,|V|\\}}$

3:$n_{i}=-1\;\forall_{i\in\\{1,2,...,|V|\\}}$

4:for $c\in V_{c}$ do

5: $g_{c}\leftarrow 0$

6: $n_{c}\leftarrow c$

7:end for

8:while True do

9: Finished $\leftarrow$ True

10: for $(i,j)\in E$ do

11: if $g_{i}+g_{ij}<g_{j}$ then

12: $g_{j}\leftarrow g_{i}+g_{ij}$

13: $n_{j}\leftarrow n_{i}$

14: Finished $\leftarrow$ False

15: end if

16: end for

17: if Finished then

18: return $\vec{g},\vec{n}$

19: end if

20:end while

While this is an iterative computation, it has finite termination when the
values in $\vec{g}$ and $\vec{n}$ stop changing. After running this modified
Bellman-Ford algorithm, Lloyd’s algorithm modifies the clusters by selecting
the centroid of every subdomain as its new center, then iterates, using the
modified Bellman-Ford algorithm to calculate new distances and nearest
centers. Given updated center positions, it forms the new subdomains. The full
Lloyd algorithm is shown in Algorithm 2, where we define the set of border
nodes, $B$, as the set of all nodes that are connected by an edge to a node
that has a different nearest center node. The key point here is that we use
Modified Bellman-Ford to assign closest centers, then compute the set of
border nodes, then find the new centers as those that are further from the
border set within each of the original subdomains (using graph distances from
$B$, but original center assignment in $\vec{n}$).

Algorithm 2 Lloyd Aggregation

1:Input $K$: Number of iterations, $E$: The set of all edges, $V$: The set of
all nodes, $V_{c}$: The set of initial center nodes.

2:for $i=1,2,3,...,K$ do

3: $\vec{g},\vec{n}\leftarrow$ Modified Bellman-Ford$(E,V,V_{c})$

4: $B\leftarrow\emptyset$

5: for $(i,j)\in E$ do

6: if $n_{i}\neq n_{j}$ then

7: $B\leftarrow B\cup\\{i,j\\}$

8: end if

9: end for

10: $\vec{g},\vec{x}\leftarrow$ Modified Bellman-Ford$(E,V,B)$

11: $V_{c}\leftarrow\\{i\in V:g_{i}>g_{j}\;\;\forall{n_{i}=n_{j}}\\}$

12:end for

13:return $\vec{n}$

#### Time Complexity:

Assuming each node’s initial distance to a center node is bounded
independently of $|V|$, and also assuming that each node’s degree is bounded
independently of $|V|$, Algorithm 1 runs a $|V|$-independent number of
iterations to determine one nearest center node for every point. Thus,
Algorithm 1 is $O(|V|)$ in our case. This is run a $|V|$-independent number of
times in Algorithm 2, to generates the subdomains, resulting in an overall
algorithmic cost of $O(|V|)$.