\section{Introduction}
	The solution of large, sparse systems of linear equations is ubiquitous when partial differential equations (PDEs) are discretized to simulate complex natural phenomena such as fluid flow \cite{losasso:2006:spatially}, thermodynamics \cite{chen:2021:momentum}, mechanical fracture \cite{paluszny:2011:numerical}, etc.  Discretization often leads to a linear system
	\begin{align}\label{eq:lin_sys}
		\textbf{A}\textbf{x}=\textbf{b}
	\end{align}
	where the dimension $n$ of the matrix $\textbf{A}\in\mathbb{R}^{n \times n}$ and the vector $\textbf{b}\in\mathbb{R}^{n}$ correlate with spatial fidelity of the computational domain. 
	Overall quality and realism of a simulation is proportional to this spatial fidelity and typical modern applications of numerical PDEs require the solution of linear systems with many millions of unknowns.
	In such applications, numerical approximation to the solution of these linear systems is typically the bottleneck in overall performance; accordingly, practitioners have spent decades devising specialized algorithms for their efficient solution \cite{golub2012matrix,saad:2003:sparse}.
	
	The appropriate numerical linear algebra technique depends on the nature of the problem.
	In the present work, we consider sparse linear systems that arise from discrete Poisson equations in incompressible flow applications \cite{chorin:1967:numerical,fedkiw:2001:visual,bridson:2008:fluid-simulation}.
	Direct solvers that utilize matrix factorizations (QR, Cholesky, etc.\ \cite{trefethen1997numerical}) have optimal approximation error, but their computational cost is $O(n^3)$ and they typically require dense storage, even for sparse $\AA$.
	Although Fast Fourier Transforms \cite{nussbaumer:1981:fast} can be used in limited instances (periodic boundary conditions, etc.), iterative techniques are most commonly adopted for these systems given their sparsity.
	Many applications with strict performance constraints (e.g., real-time fluid simulation) utilize basic iterations (Jacobi, Gauss-Seidel, successive over relaxation (SOR), etc.) given limited computational budget \cite{saad:2003:sparse}.
	However, large approximation errors must be tolerated since iteration counts are limited by the performance constraints.  
	This is particularly problematic since the wide elliptic spectrum of these matrices (a condition that worsens with increased spatial fidelity/matrix dimension) leads to poor conditioning and iteration counts.
	Iterative techniques can achieve sub-quadratic convergence if their iteration count does not grow excessively with problem size $n$ since each iteration generally requires $O(n)$ floating point operations for sparse matrices.
	Discrete elliptic operators are typically symmetric positive (semi) definite and the preconditioned conjugate gradients method (PCG) can be used to minimize iteration counts \cite{saad:2003:sparse,hestenes:1952:methods,stiefel:1952:einige}.
	Preconditioners $\mb{P}$ for PCG must simultaneously:  be symmetric positive definite (and therefore admit factorization $\mb{P}=\mb{F}^2$), improve the condition number of the preconditioned system $\mb{F}\AA\mb{F}\textbf{y}=\mb{F}\textbf{b}$, and be computationally cheap to construct and apply; accordingly, designing specialized preconditioners for particular classes of problems is somewhat of an art. 
	Incomplete Cholesky preconditioners (ICPCG) \cite{kershaw:1978:incomplete} use a sparse approximation to the Cholesky factorization and significantly reduce iteration counts in practice; however, their inherent data dependency prevents efficient parallel implementation.
	Nonetheless, these are very commonly adopted for Poisson equations arising in incompressible flow \cite{fedkiw:2001:visual,bridson:2008:fluid-simulation}.
	Multigrid \cite{brandt:1977:multi} and domain decomposition \cite{saad:2003:sparse} preconditioners greatly reduce iterations counts, but they must be updated (with non-trivial cost) each time the problem changes (e.g., in computational domains with time varying boundaries) and/or for different hardware platforms.
	In general, choice of an optimal preconditioner for discrete elliptic operators is an open area of research.
	
	Recently, data-driven approaches that leverage deep learning techniques have shown promise for solving linear systems.
	Various researchers have investigated machine learning estimation of multigrid parameters \cite{greenfeld:2019:learning,grebhahn:2016:performance,luz:2020:learning}.
	Others have developed machine learning methods to estimate preconditioners \cite{gotz:2018:machine,stanaityte:2020:ilu,ichimura:2020:fast} and initial guesses for iterative methods \cite{luna:2021:accelerating,um:2020:solver,ackmann:2020:machine}.
	Thompson et al. \cite{tompson:2017:accelerating} and Yang et al. \cite{yang:2016:data} develop non-iterative machine learning approximations of the inverse of discrete Poisson equations from incompressible flow.
	We leverage deep learning and develop a novel accelerated conjugate gradients iterative method for approximating the solution of symmetric positive definite linear systems which we call the deep gradient correction method (DGCM).
	CG iteratively adds $\textbf{A}$-conjugate search directions while minimizing the matrix norm of the error.  
	We use a convolutional neural network as an approximation of the inverse of the matrix in order to generate more efficient search directions.
	We only ask that our network approximate the inverse up to an unknown scaling since this decreases the degree of nonlinearity and since it does not affect the quality of the search direction (which is scale independent).
	The network is similar to a preconditioner, but it is not a linear function, and our modified conjugate gradients approach is designed to accommodate this nonlinearity.  
	We use unsupervised learning to train our network with a loss function equal to the $L^2$ difference between an input vector and a scaling of $\AA$ times the output of our network.
	To account for this unknown scaling during training, we choose the scale of the output of the network by minimizing the matrix norm of the error.
	Our approach allows for efficient training and generalization to problems unseen (matrices $\AA$ and right hand sides $\textbf{b}$).
	We benchmark our algorithm using the ubiquitous pressure Poisson equation (discretized on regular voxelized domains) and compare against FluidNet \cite{tompson:2017:accelerating}, which is the state of the art learning-based method for these types of problems.  
	Unlike the non-iterative approaches of Thompson et al. \cite{tompson:2017:accelerating} and Yang et al. \cite{yang:2016:data}, our method can reduce the linear system residuals \textit{arbitrarily}.
	We showcase our approach with examples that have over 16 million degrees of freedom.  




	
	\section{Related Work}
	
	Deep learning has been applied in various ways to physical simulation and the numerical solution of PDEs (see, e.g., the reviews in \cite{beck:2020:overview,gibou:2019:sharp,karniadakis:2021:physics}).  From learning how to discretize PDEs \cite{barsinai:2019:learning} to upsampling the results of low-resolution simulations \cite{jakob:2020:fluid,tumanov:2021:data}, practitioners have sought to incorporate learning into most aspects of numerical PDE pipelines.
	
	A number of recent works have sought to eschew the numerical solution of PDEs by using neural networks and appropriate representations of physical quantities \cite{ummenhofer:2020:lagrangian,sanchez:2020:learning,wiewel:2019:latent}, often taking into account PDEs or physics while training.
	Many other works have hybridized traditional techniques and governing equations with neural networks in order to improve the accuracy and/or efficiency of solving PDEs, without replacing PDEs wholesale (e.g., \cite{liang:2019:differentiable,holl:2020:phiflow,sirignano:2020:dpm}).  We note the popular physics-informed neural network framework of Karniadakis and colleagues \cite{cai:2022:physics,raissi:2019:physics}, which uses automatic differentiation to represent all PDE operators and incorporates physical constraints like conservation laws into the network's loss function. 
	
	Notably, several papers have focused on enhancing the solution of linear systems (arising from discretized PDEs) using learning.  For instance, \cite{gotz:2018:machine} generate sparsity patterns for block-Jacobi preconditioners using convolutional neural networks (CNNs), and \cite{stanaityte:2020:ilu} use a CNN to predict non-zero patterns for ILU-type preconditioners for the Navier-Stokes equations (though neither work designs fundamentally new preconditioners).  \cite{ichimura:2020:fast} develop a neural-network based preconditioner where the network is used to predict approximate Green's functions (which arise in the analytical solution of certain PDEs) that in turn yield an approximate inverse of the linear system.  \cite{hsieh:2019:learning} learn an iterator that solves linear systems, performing competitively with classical solvers like multigrid-preconditioned MINRES \cite{paige:1975:solution}. Luz et al.\ \cite{luz:2020:learning} and Greenfield et al.\ \cite{greenfeld:2019:learning} use machine learning to estimate algebraic multigrid (AMG) parameters. 
	They note that AMG approaches rely most fundamentally on effectively chosen (problem-dependent) prolongation sparse matrices and that numerous methods have attempted to automatically create them from the matrix $\AA$. 
	They train a graph neural network to learn (in an unsupervised fashion) a mapping from matrices $\AA$ to prolongation operators.  
	Gebhahn et al.\ \cite{grebhahn:2016:performance} note that geometric multigrid solver parameters can be difficult to choose to guarantee parallel performance on different hardware platforms.
	They use machine learning to create a code generator to help achieve this.
	
	Several works consider accelerating the solution of linear systems by learning an initial guess that is close to the true solution or otherwise helpful to descent algorithms for finding the true solution.  In order to solve the discretized Poisson equation, \cite{luna:2021:accelerating} accelerate the convergence of GMRES \cite{saad:1986:gmres} with an initial guess that is learned in real-time (i.e., as a simulation code runs) with no prior data.  \cite{um:2020:solver} train a network (incorporating differentiable physics, based on the underlying PDEs) in order to produce high-quality initial guesses for a CG solver.  In a somewhat similar vein, \cite{ackmann:2020:machine} use a simple feedforward neural network to predict pointwise solution components, which accelerates the conjugate residual method used to solve a relatively simple shallow-water model (a more sophisticated network and loss function are needed to handle more general PDEs and larger-scale problems).
	
	At least two papers \cite{ruelmann:2018:prospects,sappl:2019:deep} have sought to learn a mapping between a matrix and an associated sparse approximate inverse.  In their investigation, \cite{ruelmann:2018:prospects} propose training a neural network using matrix-inverse pairs as training data.  Although straightforward to implement, the cost of generating training data, let alone training the network, is prohibitive for large-scale 3D problems.  \cite{sappl:2019:deep} seek to learn a mapping between linear system matrices and sparse (banded) approximate inverses.  Their loss function is the condition number of the product of the system matrix and the approximate inverse; the minimum value of the condition number is one, which is achieved exactly when an exact inverse is obtained.  Although this framework is quite simple, evaluation of the condition number of a matrix is asymptotically costly, and in general, the inverse of a sparse matrix can be quite dense.  Accordingly, the method is not efficient or accurate enough for the large-scale 3D problems that arise in real-world engineering problems.
	
	Most relevant to the present work is FluidNet \cite{tompson:2017:accelerating}.  FluidNet develops a highly-tailored CNN architecture that is used to predict the solution of a linear projection operation (specifically, for the discrete Poisson equation) given a matrix and right hand side.  The authors demonstrate fluid simulations where the linear solve is replaced by evaluating their network.  Because their network is relatively lightweight and is only evaluated once per time step, their simulations run efficiently.  However, their design allows the network only one opportunity to reduce the residual for the linear solve; in practice, we observe that FluidNet is able to reduce the residual by no more than about one order of magnitude.  However, in computer graphics applications, at least four orders of magnitude in residual reduction are usually required for visual fidelity, while in scientific and engineering applications, practitioners prefer solutions that reduce the residual by eight or more orders of magnitude (i.e., to within machine precision).  Accordingly, FluidNet's lack of convergence stands in stark contrast to classical, convergent methods like CG.  Our method resolves this gap.
	

	\section{Motivation: Incompressible Flow}\label{sec:flow}
	
	\begin{figure}
		\centering
	
	
		\includegraphics[width=1\linewidth]{figures/mac_grid_labeled.pdf}
		\caption{\textbf{(a)} We illustrate a sample flow domain $\Omega\subset(0,1)^2$ (in 2D for ease of illustration) with internal boundaries (blue lines). \textbf{(b)} We voxelize the domain with a regular grid: white cells represent interior/fluid, and blue cells represent boundary conditions. \textbf{(c)} We train using matrix $\AA^{(0,1)^d}$ from a discretized domain with no interior boundary conditions, where $d$ is the dimension. \textbf{(d)} We illustrate the non-zero entries in an example matrix $\AA^\Omega$ from the voxelized and labeled (white vs. blue) grid for three example interior cells (green, purple, and brown). Each case illustrates the non-zero entries in the row associated with the example cell. All entries in rows corresponding to boundary/blue cells are zero.}
		\label{fig:mac_grid}
	\end{figure}
	
	We demonstrate the efficacy of our approach with the linear systems that arise in incompressible flow applications.
	Specifically, we use our algorithm to solve the discrete Poisson equations in regular-grid-based discretization of the pressure projection equations that arise in Chorin's splitting technique \cite{chorin:1967:numerical} for the inviscid, incompressible Euler equations.
	These equations are
	\begin{align}
		\rho \left(\dfrac{\partial \mb{u}}{\partial t} + \frac{\partial \mb{u}}{\partial \boldsymbol{x}}\mb{u}\right) + \nabla p = \mb{f}^{ext}, \qquad	\nabla \cdot \mathbf{u}  = 0 \label{eq:eulerequation}
	\end{align}
	where $\mb{u}$ is fluid velocity, $p$ is pressure, $\rho$ is density, and $\mb{f}^{ext}$ accounts for external forces like gravity.
	The equations are assumed at all positions $\boldsymbol{x}$ in the spatial fluid flow domain $\Omega$ and for time $t>0$.
	The left term in Equation~\ref{eq:eulerequation} enforces conservation of momentum in the absence of viscosity, and the second part enforces incompressibility and conservation of mass.
	These equations are subject to initial conditions $\rho(\boldsymbol{x},0)=\rho^0$ and $\mb{u}(\boldsymbol{x},0)=\mb{u}^0(\boldsymbol{x})$ as well as boundary conditions $\mb{u}(\boldsymbol{x},t)\cdot\mb{n}(\boldsymbol{x})=u^{\partial \Omega}(\boldsymbol{x},t)$ on the boundary of the domain $\boldsymbol{x}\in\partial \Omega$ (where $\mb{n}$ is the unit outward pointing normal at position $\boldsymbol{x}$ on the boundary).
	Equation~\ref{eq:eulerequation} is discretized in both time and space.
	Temporally, we split the advection $\dfrac{\partial \mb{u}}{\partial t} + \dfrac{\partial \mb{u}}{\partial \boldsymbol{x}}\mb{u} = 0$ and body forces terms $\rho\dfrac{\partial \mb{u}}{\partial t} = \mb{f}^{ext}$, and finally enforce incompressibility via the pressure projection $\dfrac{\partial \mb{u}}{\partial t} + \dfrac{1}{\rho} \nabla p = \mathbf{0} \text{ such that } \nabla \cdot \mb{u} = 0$; this is the standard advection-projection scheme proposed by Chorin \cite{chorin:1967:numerical}.
	Using finite differences in time, we can summarize this as
	\begin{align}
		\rho^0 \left(\frac{\mb{u}^*-\mb{u}^n}{\Delta t}+\frac{\partial \mb{u}^n}{\partial \boldsymbol{x}}\mb{u}^n\right)&=\mb{f}^{ext}\label{eq:advect}\\
		-\nabla\cdot\frac{1}{\rho^0}\nabla p^{n+1}&=-\nabla\cdot\mb{u}^* \label{eq:proj_1}\\
		-\frac{1}{\rho^0}\nabla p^{n+1} \cdot\mb{n}&=\frac{1}{\Delta t}\left(u^{\partial \Omega}-\mb{u}^*\cdot\mb{n}\right). \label{eq:proj_2}
	\end{align} 
	For the spatial discretization, we use a regular marker-and-cell (MAC) grid \cite{harlow:1965:viscous-flow} with cubic voxels whereby velocity components are stored on the face of voxel cells, and scalar quantities (e.g., pressure $p$ or density $\rho$) are stored at voxel centers.
	We use backward semi-Lagrangian advection \cite{fedkiw:2001:visual,gagniere:2020:hybrid} for Equation~\ref{eq:advect}.
	All spatial partial derivatives are approximated using finite differences. 
	Equations~\ref{eq:proj_1} and \ref{eq:proj_2} describe the pressure Poisson equation with Neumann conditions on the boundary of the flow domain.
	We discretize the left hand side of Equation~\ref{eq:proj_1} using a standard 7-point finite difference stencil.
	The right hand side is discretized using the MAC grid discrete divergence finite difference stencils as well as contributions from the boundary condition terms in Equation~\ref{eq:proj_2}.
	We refer the reader to \cite{bridson:2008:fluid-simulation} for more in-depth implementation details.
	Equation~\ref{eq:proj_2} is discretized by modifying the Poisson stencil to enforce Neumann boundary conditions.
	We do this using a simple labeling of the voxels in the domain.
	For simplicity we assume $\Omega\subset(0,1)^3$ is a subset of the unit cube, potentially with internal boundaries.
	We label cells in the domain as either liquid or boundary.
	This simple classification is enough to define the Poisson discretizations (with appropriate Neumann boundary conditions at domain boundaries) that we focus on in the present work; we illustrate the details in Figure~\ref{fig:mac_grid}.
	We use the following notation to denote the discrete Poisson equations associated with Equations~\ref{eq:proj_1}--\ref{eq:proj_2}:
	\begin{align}\label{eq:disc_proj}
		\AA^\Omega \textbf{x}=\textbf{b}^{\nabla \cdot \mb{u}^*} + \textbf{b}^{u^{\partial \Omega}},
	\end{align}
	where $\AA^\Omega$ is the discrete Poisson matrix associated with the voxelized domain, $\textbf{x}$ is the vector of unknown pressure, and $\textbf{b}^{\nabla \cdot \mb{u}^*}$ and  $\textbf{b}^{u^{\partial \Omega}}$ are the right hand side terms from Equations~\ref{eq:proj_1} and  \ref{eq:proj_2} respectively.
	We define a special case of the matrix involved in this discretization to be the Poisson matrix $\AA^{(0,1)^3}$ associated with $\Omega=(0,1)^3$, i.e., a full fluid domain with no internal boundaries.
	We use this matrix for training, yet demonstrate that our network generalizes to all other matrices arising from more complicated flow domains.
	
	
	\section{Deep Gradient Correction Method}
	We present our method for the deep learning acceleration of iterative approximations to the solution of linear systems of the form seen in Equation~\ref{eq:disc_proj}.
	We first discuss relevant details of the method of conjugate gradients, particularly line search and $\AA$-orthogonal search directions.
	We then present a deep learning technique for improving the quality of these search directions that ultimately reduces iteration counts required to achieve satisfactory residual reduction.
	Lastly, we outline the training procedures for our deep convolutional neural network.
	\subsection{Line search and $\AA$-orthogonality}
	Our approach iteratively improves approximations to the solution $\textbf{x}$ of Equation~\ref{eq:disc_proj}.
	We build on the method of conjugate gradients (CG), which requires the matrix $\AA^\Omega$ in Equation~\ref{eq:disc_proj} to be symmetric positive definite (SPD).
	SPD matrices $\AA^\Omega$ give rise to the matrix norm $\left\|\textbf{y}\right\|_{\AA^\Omega}=\sqrt{\textbf{y}^T\AA^\Omega\textbf{y}}$.
	The conjugate gradients method can be derived in terms of iterative line search improvement based on optimality in this norm.
	That is, an iterate $\textbf{x}_{k-1}\approx\textbf{x}$ is updated in search direction $\textbf{d}_k$ by a step size $\alpha_k$ that is chosen to minimize the matrix norm of the error:
	\begin{align}
		\alpha_k = \argmin_\alpha \frac{1}{2}\left\|\textbf{x}-\left(\textbf{x}_{k-1}+\alpha \textbf{d}_k\right)\right\|^2_{\AA^\Omega}= \frac{\textbf{r}_{k-1}^T\textbf{d}_k}{\textbf{d}_k^T\AA^\Omega\textbf{d}_k}, \label{eq:step_size}
	\end{align}
	where $\textbf{r}_{k-1}=\textbf{b} - \AA^\Omega\textbf{x}_{k-1}$ is the $(k-1)^{\text{th}}$ residual.
	Different search directions result in different algorithms. 
	A natural choice is the negative gradient of the matrix norm of the error (evaluated at the current iterate), since this will point in the direction of steepest decrease $\textbf{d}_k=-\frac{1}{2}\nabla \left\|\textbf{x}_{k-1}\right\|^2_{\AA^\Omega}=\textbf{r}_{k-1}$.
	This is the gradient descent method (GD).
	Unfortunately, this approach requires many iterations in practice. 
	A more effective strategy is to choose directions that are $\AA$-orthogonal, i.e.,
	\begin{align*}
		\textbf{d}_i^T\AA^\Omega\textbf{d}_j=0, \ i\neq j.
	\end{align*}
	With this choice, the search directions form a basis for $\mathbb{R}^n$ with which we can write the initial error as 
	\begin{align*}
		\textbf{e}_0=\textbf{x}-\textbf{x}_0=\sum_{i=1}^n e_i \textbf{d}_i,
	\end{align*}
	where $e_i$ are the components of the initial error written in the basis.
	Furthermore, when the search directions are $\AA$-orthogonal, the optimal step sizes $\alpha_k$ at each iteration satisfy
	\begin{align*}
		\alpha_k=\frac{\textbf{r}_{k-1}^T\textbf{d}_k}{\textbf{d}_k^T\AA^\Omega\textbf{d}_k}=\frac{\textbf{d}_k^T\AA^\Omega\textbf{e}_{k-1}}{\textbf{d}_k^T\AA^\Omega\textbf{d}_k}=\frac{\textbf{d}_k^T\AA^\Omega\left(\sum_{i=1}^n e_i \textbf{d}_i-\sum_{j=1}^{k-1}\alpha_j \textbf{d}_j\right)}{\textbf{d}_k^T\AA^\Omega\textbf{d}_k}=e_k.
	\end{align*}
	That is, the optimal step sizes are chosen to exactly eliminate the components of the error in the basis defined by the search directions, and therefore, convergence is determined by the (at most $n$) non-zero components $e_i$ in the initial error.
	Although rounding errors prevent this from happening exactly in practice, this property does greatly reduce the number of required iterations \cite{golub2012matrix}.
	
	CG can be viewed as a modification of GD where the search direction is chosen as the component of the residual (equivalently, the negative gradient of the matrix norm of the error) that is $\AA$-orthogonal to all previous search directions:
	\begin{align*}
		\textbf{d}_k=\textbf{r}_{k-1}-\sum_{i=1}^{k-1} h_{ik}\textbf{d}_i, \qquad h_{ik} =\frac{\textbf{d}_i^T\AA^\Omega\textbf{r}_{k-1}}{\textbf{d}_i^T\AA^\Omega\textbf{d}_i}.
	\end{align*}
	In practice, $h_{ik}=0$ for $i<k-1$, and this iteration can therefore be performed without the need to store all previous search directions $\textbf{d}_i$ and without the need for computing all previous $h_{ik}$.
	

	\subsection{Learning Optimal Search Directions}
	While the residual is a natural choice for generating $\AA$-orthogonal search directions (since it points in the direction of the steepest local decrease), it is not the optimal search direction.
	If $\textbf{d}_k$ is parallel to $(\AA^\Omega)^{-1}\textbf{r}_{k-1}$, then $\textbf{x}_k$ will be equal to $\textbf{x}$ since $\alpha_k$ (computed from Equation \ref{eq:step_size}) will step directly to the solution.
	We can see this by considering the residual and its relation to the search direction:
	\begin{center}
		$\textbf{r}_k = \textbf{b}-\AA^\Omega\textbf{x}_k=\textbf{b}-\AA^\Omega\textbf{x}_{k-1}-\alpha_k\AA^\Omega\textbf{d}_k=\textbf{r}_{k-1}-\alpha_k\AA^\Omega\textbf{d}_k$.
	\end{center}
	In light of this, we use deep learning to create an approximation $\mb{f}(\textbf{c},\textbf{r})$ to $(\AA^\Omega)^{-1}\textbf{r}$, where $\textbf{c}$ denotes the network weights and biases.
	This is analogous to using a preconditioner in PCG; however, our network is not symmetric positive definite (nor even a linear function).
	We simply use this data-driven approach as our means of generating better search directions $\textbf{d}_k$.
	Furthermore, we only need to approximate a vector parallel to $(\AA^\Omega)^{-1}\textbf{r}$ since the step size $\alpha_k$ will account for any scaling in practice.
	In other words, $\mb{f}(\textbf{c},\textbf{r})\approx s_\textbf{r} (\AA^\Omega)^{-1}\textbf{r}$, where the scalar $s_\textbf{r}$ is not defined globally; it only depends on $\textbf{r}$, and the model does not learn it.
	Lastly, as with CG, we enforce $\AA$-orthogonality, yielding search directions
	\begin{align*}
		\textbf{d}_k=\mb{f}(\textbf{c},\textbf{r}_{k-1})-\sum_{i=1}^{k-1}h_{ik}\textbf{d}_i,\qquad	h_{ik}=\frac{\mb{f}(\textbf{c},\textbf{r}_{k-1})^T\AA^\Omega\textbf{d}_i}{\textbf{d}_i^T\AA^\Omega\textbf{d}_i}.
	\end{align*}
	
	
	\begin{wrapfigure}{R}{0.3\textwidth}
		\vspace{-25pt}
		\begin{minipage}{0.3\textwidth}
		
			\begin{algorithm}[H]
				\begin{algorithmic}[1]
					\STATE $\textbf{r}_0 = \textbf{b}-\AA^\Omega\textbf{x}_0$
					\STATE $\textbf{d}_1=\textbf{r}_0$
					\STATE $k=1$
				
					\WHILE{$\|\textbf{r}_{k-1}\|\ge\epsilon$}
					\STATE $\textbf{d}_k =\mb{f}(\textbf{c},\frac{\textbf{r}_{k-1}}{\left\|\textbf{r}_{k-1}\right\|})$
					\FOR{$i_\textrm{start}\leq i <k$} 
					\STATE $h_{ik}=\frac{\textbf{d}_k^T\AA^\Omega\textbf{d}_i}{\textbf{d}_i^T\AA^\Omega\textbf{d}_i}$
					\STATE $\textbf{d}_k$-=$h_{ik}\textbf{d}_i$
					\ENDFOR
					\STATE $\alpha_{k}=\frac{\textbf{r}_{k-1}^T\textbf{d}_k}{\textbf{d}_{k}^T\AA^\Omega\textbf{d}_{k}} $
				
					\STATE $\textbf{x}_{k}=\textbf{x}_{k-1}+\alpha_k \textbf{d}_{k}$
					\STATE $\textbf{r}_k = \textbf{b} - \AA^\Omega\textbf{x}_k$
					\STATE $k=k+1$
					\ENDWHILE
				\end{algorithmic}
				\caption{\textbf{DGCM}}
				\label{alg:mlpcg}
			\end{algorithm}
		
			\vspace{-40pt}
		\end{minipage}
		\vspace{-10pt}
	\end{wrapfigure}
	
	We summarize our approach in Algorithm~\ref{alg:mlpcg}. Note that we introduce the variable $i_\textrm{start}$.
	To guarantee $\AA$-orthogonality between all search directions, we must have $i_\textrm{start}=1$.
	However, this requires storing all prior search directions, which can be costly.
	We found that using $i_\textrm{start}=k-2$ worked nearly as well as $i_\textrm{start}=1$ in practice (in terms of our ability to iteratively reduce the residual of the system).
	We demonstrate this in Figure~\ref{fig:four_plots}c.
	
	\newpage
	\section{Model Architecture, Datasets, and Training}
	Efficient performance of our method requires effective training of our deep convolutional network for weights and biases $\textbf{c}$ such that $\mb{f}(\textbf{c},\textbf{r})\approx s_\textbf{r} (\AA^\Omega)^{-1}\textbf{r}$ (for arbitrary scalar $s_\textbf{r}$).  
	We design a model architecture, loss function, and unsupervised training approach to achieve this.
	Our approach has modest training requirements and allows for effective residual reduction while generalizing well to problems not seen in the training data. 
	
	\subsection{Loss Function and Unsupervised Learning}\label{sec:dataset_creation}
	Although we generalize to arbitrary matrices $\AA^{\Omega}$ from Equation~\ref{eq:disc_proj} that correspond to domains $\Omega\subset(0,1)^3$ that have internal boundaries (see Figure~\ref{fig:mac_grid}), we train using just the matrix $\AA^{(0,1)^3}$ from the full cube domain $(0,1)^3$.
	In contrast, other similar approaches \cite{tompson:2017:accelerating,yang:2016:data} train using matrices $\AA^{\Omega}$ and right hand sides $\textbf{b}^{\nabla \cdot \mb{u}^*} + \textbf{b}^{u^{\partial \Omega}}$ that arise from flow in many domains with internal boundaries. 
	We train our network by minimizing the $L^2$ difference $\|\textbf{r} - \alpha \AA^{(0,1)^3}\mb{f}(\textbf{c},\textbf{r})\|_2$, where $\alpha =\frac{\textbf{r}^T \mb{f}(\textbf{c},\textbf{r})}{\mb{f}(\textbf{c},\textbf{r})^T \AA^{(0,1)^3} \mb{f}(\textbf{c},\textbf{r})}$ from Equation~\ref{eq:step_size}.  
	This choice of $\alpha$ accounts for the unknown scaling in the approximation of $\mb{f}(\textbf{c},\textbf{r})$ to $\left({\AA^{(0,1)^3}}\right)^{-1}\textbf{r}$.
	We use an unsupervised approach and train the model by minimizing
	\begin{center}
		$\text{Loss}(\mb{f},\textbf{c},\mathcal{D}) = \frac{1}{|\mathcal{D}|}\sum_{\textbf{r}\in \mathcal{D}}\|\textbf{r} - \frac{\textbf{r}^T \mb{f}(\textbf{c},\textbf{r})}{\mb{f}(\textbf{c},\textbf{r})^T \AA^{(0,1)^3}\mb{f}(\textbf{c},\textbf{r})}\AA^{(0,1)^3} \mb{f}(\textbf{c},\textbf{r})\|_2$
	\end{center}
	for given dataset $\mathcal{D}$ consisting of training vectors $\textbf{b}^i$.
	In Algorithm~\ref{alg:mlpcg}, the normalized residuals $\frac{\textbf{r}_k}{\|\textbf{r}_k\|}$ are passed as inputs to the model. 
	Unlike in e.g.\ FluidNet \cite{tompson:2017:accelerating}, only the first residual $\frac{\textbf{r}_0}{\|\textbf{r}_0\|}$ is directly related to the problem-dependent original right hand side $\textbf{b}$. 
	Hence we consider a broader range of training vectors than those expected in a given problem of interest, e.g., incompressible flows.
	We observe that generally the residuals $\textbf{r}_k$ in Algorithm~\ref{alg:mlpcg} are skewed to the lower end of the spectrum of the matrix $\AA^{\Omega}$.
	Since $\AA^{\Omega}$ is a discretized elliptic operator, this means that lower end modes are of lower frequency of spatial oscillation.
	We create our training vectors $\textbf{b}^i\in\mathcal{D}$ using $N\ll n$ approximate eigenvectors of the training matrix $\AA^{(0,1)^3}$.
	We use the Rayleigh-Ritz method to create approximate eigenvectors $\textbf{q}_i$, $0\leq i <N$. 
	This approach allows us to effectively approximate the full spectrum of $\AA^{(0,1)^3}$ without computing the full eigendecomposition, which can be expensive ($O(n^3)$) at high resolution.
	We found that using $N=10000$ worked well in practice.
	The Rayleigh-Ritz vectors are orthonormal and satisfy $\textbf{Q}_N^T\AA^{(0,1)^3}\textbf{Q}_N=\boldsymbol{\Lambda}_N$, where $\boldsymbol{\Lambda}_N$ is a diagonal matrix with nondecreasing diagonal entries $\lambda_i$ referred to as Ritz values (approximate eigenvalues) and $\textbf{Q}_N=\left[\textbf{q}_0,\textbf{q}_1,\hdots,\textbf{q}_{N-1}\right]\in\mathbb{R}^{n\times N}$.
	We pick $\textbf{b}^i = \frac{\sum_{j=1}^k c_j^i\textbf{q}_j}{\left\|\sum_{j=1}^k c_j^i\textbf{q}_j\right\|}$, where the coefficients $c_j^i$ are picked from a normal distribution
	\begin{equation*} c_j^i =  
		\begin{cases}
			9\cdot \mathcal{N}(0,1) & \text{  if } \tilde{j}\le j \le \frac{N}{2}+\theta\\
			\mathcal{N}(0,1) & \text{otherwise} 
		\end{cases}
	\end{equation*}
	where $\theta$ is a small number (we used $\theta=500$), and $\tilde{j}$ is the first index that $\lambda_{\tilde{j}}>0$. This choice creates $90\%$ of $\textbf{b}^i$ from the lower end of the spectrum with the remaining $10\%$ from the higher end. 
	Notably, even though this dataset only uses Rayleigh-Ritz vectors from the training matrix $\AA^{(0,1)^3}$, our network can be effectively generalized to flows in irregular domains, e.g., smoke flow past a rotating box and flow past a bunny (see Figure~\ref{fig:examples}).
	

	We generate the Rayleigh-Ritz vectors by first tridiagonalizing the training matrix $\AA^{(0,1)^3}$ with $N$ Lanczos iterations \cite{lanczos:1950:iteration,paige:1971:lanczos} to form $\textbf{T}^N={\textbf{Q}^L_N}^T\AA^{(0,1)^3}\textbf{Q}^L_N\in\mathbb{R}^{N\times N}$.
	We then diagonalize $\textbf{T}^N=\hat{\textbf{Q}}^T\boldsymbol{\Lambda}_N\hat{\textbf{Q}}$ using the QR eigenvalue algorithm with implicit Wilkinson shifts \cite{trefethen1997numerical}.
	While costly, we note that this algorithm is performed on the comparably small $N\times N$ matrix $\textbf{T}^N$ (rather than on the $\AA^{(0,1)^3}\in\mathbb{R}^{n\times n}$).
	This yields the Rayleigh-Ritz vectors as the columns of $\textbf{Q}_N=\textbf{Q}^L_N\hat{\textbf{Q}}$.
	The Lanczos vectors are the columns of the matrix $\textbf{Q}^L_N$ and satisfy a three-term recurrence whereby we can iteratively compute the next Lanczos vector from previous two as
	\begin{align*}
		\beta_j\textbf{q}^L_{j+1} = \AA^{(0,1)^3} \textbf{q}^L_j - \beta_{j-1}\textbf{q}^L_{j-1} - \alpha_j\textbf{q}^L_j ,
	\end{align*}

	where $\alpha_j$ and $\beta_j$ are diagonal and subdiagonal entries of $\textbf{T}^k$. $\beta_j$ is computed so that $\textbf{q}^L_{j+1}$ is a unit vector, and $\alpha_{j+1}=\textbf{q}_{j+1}^T\AA^{(0,1)^3}\textbf{q}_{j+1}$.
	We initialize the iteration with a random $\textbf{q}^L_{0}\in\textrm{span}(\AA^{(0,1)^3})$.
	The Lanczos algorithm can be viewed as a modified Gram-Schmidt technique to create an orthonormal basis for the Krylov space associated with $\textbf{q}^L_0$ and $\AA^{(0,1)^3}$, and it therefore suffers from rounding error sensitivities manifested as loss of orthonormality with vectors that do not appear in the recurrence. 
	We found that the simple strategy described in \cite{paige:1971:lanczos} of orthogonalizing each iterate with all previous Lanczos vectors to be sufficient for our training purposes.

	\subsection{Model Architecture}
	
	\begin{wrapfigure}{R}{0.6\textwidth}
		\vspace{-17pt}
		\begin{framed}
			\begin{center}
				\includegraphics[width=1\linewidth]{figures/ml_diagram_V4.pdf}
			\end{center}
		
			\caption{Model architecture for training with $\AA^{(0,1)^3}$ on a $128^3$ grid.}
			\vspace{-10pt}
			\label{fig:ML_model}
		\end{framed}
		\vspace{-15pt}
	\end{wrapfigure}
	The internal structure of our CNN architecture for a $128^3$ grid can be seen in Figure~\ref{fig:ML_model}. 
	Our model consists of a series of convolutional layers with residual connections. 
	The upper left of Figure~\ref{fig:ML_model} ($K$ Residual Blocks) shows our use of multiple blocks of residually connected layers.
	Importantly, within each block, the first layer directly affects the last layer with an addition operator. 
	All non-input or output convolutions use a $3\times 3 \times 3$ filter, and all layers consist of 16 feature maps. 
	In the middle of the first level, one layer is downsampled (via the average pooling operator with $(2\times2\times2)$ pool size) and another set of convolutional layers is applied with residual connection blocks. 
	The last layer in the second level is upscaled and added to the layer that is downsampled. 
	The last layer in the network is dense with a linear activation function. 
	The activation functions in all convolutional layers are Rectified Linear Units ($\ReLU$), except for the first convolution, which uses a linear activation function.\\
	\\
	Differing resolutions use differing numbers of convolutions, but the fundamental structure remains the same. 
	More precisely, the number of residual connections are changed for different resolutions. 
	For example, a $64^3$ grid uses 1 residual block on the left and 2 residual blocks on the right on the upper level, and 3 residual blocks on the lower level. 
	Furthermore, the weights trained on a lower resolution grid can be used effectively with higher resolutions. 
	Figure~\ref{fig:four_plots}d shows convergence results for a $256^3$ grid, using a model trained for a $64^3$ grid and a $128^3$ grid. 
	The model that we use for $256^3$ grids in our final examples was trained on a $128^3$ grid; however, as the shown in the figure, even training with a $64^3$ grid allows for efficient residual reduction.





	
	\subsection{Training}
         \label{sec:subtraining}
	Using the procedure explained in Section~\ref{sec:dataset_creation}, we create the training dataset $\mathcal{D}\in \textrm{span}(\AA^{(0,1)^3})\cap \mathcal{S}^{n-1}$ of size 20000 generated from 10000 Rayleigh-Ritz vectors. 
	We train our model with TensorFlow \cite{abadi:2015:tensorflow} on a single NVIDIA RTX A6000 GPU with 48GB memory.
	Training is done with standard deep learning techniques---more precisely, back-propagation and the ADAM optimizer \cite{Kingma2015AdamAM} (with starting learning rate 0.0001). 
	Training takes approximately 10 minutes and 1 hour per epoch for grid resolutions $64^3$ and $128^3$, respectively.
	We trained our model for 50 epochs; however, the model from the thirty-first epoch was optimal for $64^3$, and the model from the third epoch was optimal for $128^3$.
	
	\newpage 
	
	\section{Results and Analysis}
	\label{sec:results}
	
	\begin{wrapfigure}{R}{0.6\textwidth}
		\vspace{-17pt}
		\begin{framed}
			\begin{center}
				\centering
			
				\includegraphics[width=0.16\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example40/smoke_plume.mantra1.0200.png}%
				\includegraphics[width=0.16\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example40/smoke_plume.mantra1.0400.png}
			
				\includegraphics[width=0.16\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example42/bunny_smoke.mantra1.0150.png}%
				\includegraphics[width=0.16\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example42/bunny_smoke.mantra1.0300.png}
			
				\includegraphics[width=0.16\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example41/spinning_box.mantra1.0080.png}%
				\includegraphics[width=0.16\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example41/spinning_box.mantra1.0180.png} \\
			
				\includegraphics[width=0.32\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example40/smoke_plume.mantra1.0600.png}
			
				\includegraphics[width=0.32\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example42/bunny_smoke.mantra1.0450.png}	
			
				\includegraphics[width=0.32\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example41/spinning_box.mantra1.0280.png}
			\end{center}
		
			\caption{DGCM for simulating a variety of incompressible flow examples. Left: smoke plume at $t=6.67, 13.33, 20$ seconds. Middle: smoke passing bunny at $t = 5, 10, 15$ seconds. Right: smoke passing a spinning box (time-dependent Neumann boundary conditions) at $t = 2.67, 6, 9.33$ seconds.}
		
			\label{fig:examples}
		\end{framed}
		\vspace{-20pt}
	\end{wrapfigure}
	
	We demonstrate DGCM on three increasingly difficult examples and provide numerical evidence for the efficient convergence of our method.  All examples were run on a workstation with dual AMD EPYC 75F3 processors, 512GB RAM, and an NVIDIA RTX A6000 GPU with 48GB memory.
	
	Figure \ref{fig:examples} showcases DGCM for incompressible smoke simulations.  In each simulation, inlet boundary conditions are set in a circular portion of the bottom of the cubic domain, whereby smoke flows around potential obstacles and fills the domain.  We show a smoke plume (with no obstacles), flow past a complex static geometry (the Stanford bunny), and flow past a dynamic geometry (a rotating cube).  Visually plausible and highly-detailed results are achieved for each simulation (see supplementary material for larger videos).   The plume example uses a computational grid with resolution $128^3$, while the other two examples use grids with resolution $256^3$ (representing over 16 million unknowns).  For each linear solve, DGCM was run until the residual was reduced by four orders of magnitude.
	
	For the bunny example, Figures \ref{fig:four_plots}a--b demonstrate how residuals decrease over the course of a linear solve, comparing DGCM with other methods.  Figure \ref{fig:four_plots}a shows the mean results (with standard deviations) over the course of 400 simulation frames, while in Figure \ref{fig:four_plots}b, we illustrate behavior on a particular frame (frame 150).  For FluidNet, we use the implementation provided by \cite{fluidnetsc22:2022:fluidnet} (available under the CeCILL-B Free Software License).  This implementation includes pre-trained models (usable under the same license) that we use without modification.  In both subfigures, it is evident that the FluidNet residual never changes, since the method is not iterative; FluidNet reduces the initial residual by no more than one order of magnitude.  On the other hand, with DGCM, we can continually reduce the residual (e.g., by four orders of magnitude) as we apply more iterations of our method, just as with classical CG.  In Figure \ref{fig:four_plots}b, we also visualize the convergence of two other classical methods, CG and incomplete Cholesky preconditioned CG (ICPCG); clearly, DGCM reduces the residual in the fewest number of iterations (e.g., approximately one order of magnitude fewer iterations than ICPCG).  Since FluidNet is not an iterative method and lacks a notion of residual reduction, we treat $\textbf{r}_0$ for FluidNet as though an initial guess of zero is used (as is done in our solver).
	
	To clarify these results, Table \ref{tab:results} reports convergence statistics for DGCM compared to standard iterative techniques CG and ICPCG. 
	For both $128^3$ and $256^3$ grids with the bunny example, we measure the time $t_r$ and the number of iterations $n_r$ required to reduce the initial residual on a particular time step of the simulation by four orders of magnitude.  
	DGCM achieves the desired results in by far the fewest number of iterations at both resolutions. At $256^3$, DGCM performs approximately one order of magnitude faster than CG, suggesting a potentially even wider performance advantage at higher resolutions.  ICPCG performs poorly for both resolutions due to its inherent serial nature.
	
	\begin{wraptable}{R}{0.4\textwidth}
		\vspace{0pt}
	
		
			\vspace{-13pt}
			\caption{Timing and iteration comparison for different methods on the bunny example.}
			\label{tab:results}
			\centering
			\begin{tabular}{c c c c c}
				\toprule
				& \multicolumn{2}{c}{$128^3$ Grid} & \multicolumn{2}{c}{$256^3$ Grid}  \\
				\cmidrule(r){2-3} \cmidrule(r){4-5}
			
			
				Method          &$t_r$ & $n_r$      &$t_r$ & $n_r$ \\
				\midrule
				DGCM      & \textbf{$\mathbf{29}$s} &  \textbf{19}       & \textbf{$\mathbf{58}$s} &  \textbf{34}    \\
			
				CG     			     & $34$s   & 387      & $592$s &  656  \\
				ICPCG     	      & $59$s &   216       & $1459$s & 374  \\
				\bottomrule
			\end{tabular}
		
			\vspace{-10pt}
	
	\end{wraptable}



	\iffalse
	\begin{figure}
		\centering
	
		\includegraphics[width=0.16\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example40/smoke_plume.mantra1.0200.png}%
		\includegraphics[width=0.16\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example40/smoke_plume.mantra1.0400.png}
	
		\includegraphics[width=0.16\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example42/bunny_smoke.mantra1.0150.png}%
		\includegraphics[width=0.16\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example42/bunny_smoke.mantra1.0300.png}
	
		\includegraphics[width=0.16\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example41/spinning_box.mantra1.0080.png}%
		\includegraphics[width=0.16\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example41/spinning_box.mantra1.0180.png} \\
	
		\includegraphics[width=0.32\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example40/smoke_plume.mantra1.0600.png}
	
		\includegraphics[width=0.32\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example42/bunny_smoke.mantra1.0450.png}	
	
		\includegraphics[width=0.32\linewidth, trim={300px 40px 300px 20px}, clip]{figures/mlpcg_example41/spinning_box.mantra1.0280.png}
		\caption{DGCM for simulating a variety of incompressible flow examples. Left: smoke plume at $t=6.67, 13.33, 20$ seconds. Middle: smoke passing bunny at $t = 5, 10, 15$ seconds. Right: smoke passing a spinning box (time-dependent Neumann boundary conditions) at $t = 2.67, 6, 9.33$ seconds.}

		\label{fig:examples}

	\end{figure}
	\fi
	\iffalse
	\begin{figure}
		\centering
		\includegraphics[width=0.24\linewidth]{figures/standard_deviation.pdf}		
		\includegraphics[width=0.24\linewidth]{figures/residuals.pdf}		
		\includegraphics[width=0.24\linewidth]{figures/A_orthogonality.pdf}
		\includegraphics[width=0.24\linewidth]{figures/model64_vs_model128.pdf}
		\caption{Residuals during the course of a linear solve for the bunny example.  (Left) Averaged results (mean and std.\ dev.) over all 400 iterations of the simulation.  (Right) Results for frame 150.}
		\label{fig:residuals}
	\end{figure}
	\fi 
	
	\begin{figure}
	\centering


























	\includegraphics[width=\textwidth, trim={0px 15px 0px 20px}, clip]{figures/merged_figure.pdf}
	\caption{Convergence data for the bunny example. \textbf{(a)} Mean and std.\ dev. (over all 400 frames in the simulation) of residual reduction during linear solves (with $128^3$ and $256^3$ grids) using FluidNet (FN) and DGCM. \textbf{(b)} Residual plots with CG, ICPCG, FN, and DGCM at frame 150. Dashed and solid lines represent results for $128^3$ and $256^3$, respectively. \textbf{(c)} Decrease in residuals with varying degrees of $\textbf{A}$-orthogonalization ($i_s = i_\text{start}$). \textbf{(d)} Reduction in residuals when the network is trained with a $64^3$ or $128^3$ grid for the $256^3$ grid simulation shown in Figure \ref{fig:examples} Middle.}
	\label{fig:four_plots}

\end{figure}
	
	\section{Conclusions}
	\label{sec:conclusions}
	
	We presented DGCM, incorporating CNNs into a CG-style algorithm that yields efficient, convergent behavior for solving linear systems.  Our method is evaluated on linear systems with over 16 million degrees of freedom and converges to a desired tolerance in merely tens of iterations.  Furthermore, despite training the underlying network on domains without obstacles, our network is able to successfully predict search directions that enable efficient linear solves on domains with complex and dynamic geometries.  Moreover, the training data for our network does not require running fluid simulations or solving linear systems ahead of time; our Rayleigh-Ritz vector approach enables us to quickly generate very large training datasets, unlike approaches seen in other works.
	
	Our network was designed for and trained exclusively using data related to the discrete Poisson matrix, which likely limits the generalizability of our present model.  However, we believe our \textit{method} is readily applicable to other classes of PDEs (or general problems with graph structure) that give rise to large, sparse, symmetric linear systems.  We note that our method is unlikely to work well for dense $\AA$ since that would likely entail using an impractical number of convolutional layers.  An interesting question to consider is how well our method and current models would apply to discrete Poisson matrices arising from non-uniform grids, e.g., quadtrees or octrees \cite{losasso:2004:simulating}.
	
	
	\begin{ack}
		The authors would like to thank Guido Mont\'ufar for the suggestion to use the weights and biases from networks trained for lower-resolution grids to make predictions on higher-resolution grids.
	\end{ack}
	
	\bibliographystyle{plain}
	