# Augmented Lagrangian preconditioning for P2-P0 discretisation of the Stokes problem

In this tutorial, we explore constructing preconditioners for discretisations of the Stokes problem with Boffi-Lovadina augmentation. In particular, we will consider a P2-P0 inf-sup stable discretisation of the Stokes problem, i.e.

\[ \begin{align}\begin{aligned}\text{find } (\vec{u},p) \in [H^1_{0}(\Omega)]^d\times L^2(\Omega) \text{ s.t. }\\\begin{split}\begin{cases} (\nabla \vec{u},\nabla \vec{v})_{L^2(\Omega)} + (\nabla\cdot \vec{v}, p)_{L^2(\Omega)} = (\vec{f},\vec{v})_{L^2(\Omega)} \qquad \forall v\in H^1_{0}(\Omega)\\ (\nabla\cdot \vec{u},q)_{L^2(\Omega)} = 0 \qquad \forall q\in L^2(\Omega) \end{cases}\end{split}\end{aligned}\end{align} \]

Such a discretisation can easily be constructed using NGSolve as follows:

    from ngsolve import *
    from netgen.occ import *
    import netgen.gui
    import netgen.meshing as ngm
    from mpi4py.MPI import COMM_WORLD

    if COMM_WORLD.rank == 0:
       shape = Rectangle(2,0.41).Circle(0.2,0.2,0.05).Reverse().Face()
       shape.edges.name="wall"
       shape.edges.Min(X).name="inlet"
       shape.edges.Max(X).name="outlet"
       geo = OCCGeometry(shape, dim=2)
       ngmesh = geo.GenerateMesh(maxh=0.1)
       mesh = Mesh(ngmesh.Distribute(COMM_WORLD))
    else:
       mesh = Mesh(ngm.Mesh.Receive(COMM_WORLD))
    nu = Parameter(1.0)
    V = VectorH1(mesh, order=2, dirichlet="wall|inlet|cyl", autoupdate=True)
    Q = L2(mesh, order=0, autoupdate=True)
    u,v = V.TnT(); p,q = Q.TnT()
    a = BilinearForm(nu*InnerProduct(Grad(u),Grad(v))*dx)
    a.Assemble()
    b = BilinearForm(div(u)*q*dx)
    b.Assemble()
    gfu = GridFunction(V, name="u")
    gfp = GridFunction(Q, name="p")
    uin = CoefficientFunction( (1.5*4*y*(0.41-y)/(0.41*0.41), 0) )
    gfu.Set(uin, definedon=mesh.Boundaries("inlet"))
    f = LinearForm(V).Assemble()
    g = LinearForm(Q).Assemble();

We can now explore what happens if we use a Schur complement preconditioner for the saddle point problem. We can construct the (dense!) Schur complement preconditioner using the following code:

    K = BlockMatrix( [ [a.mat, b.mat.T], [b.mat, None] ] )
    from ngsPETSc.pc import *
    apre = Preconditioner(a, "PETScPC", pc_type="lu")
    S = (b.mat @ apre.mat @ b.mat.T).ToDense().NumPy()
    from numpy.linalg import inv
    from scipy.sparse import coo_matrix
    from ngsolve.la import SparseMatrixd
    Sinv = coo_matrix(inv(S))
    Sinv = la.SparseMatrixd.CreateFromCOO(indi=Sinv.row,
                                          indj=Sinv.col,
                                          values=Sinv.data,
                                          h=S.shape[0],
                                          w=S.shape[1])
    C = BlockMatrix( [ [a.mat.Inverse(V.FreeDofs()), None],
                       [None, Sinv] ] )

    rhs = BlockVector (  [f.vec, g.vec] )
    sol = BlockVector( [gfu.vec, gfp.vec] )
    print("-----------|Schur|-----------")
    solvers.MinRes (mat=K, pre=C, rhs=rhs, sol=sol, tol=1e-8,
                    printrates=True, initialize=False)
    Draw(gfu)

The Schur complement preconditioner converges in a few iterations, but as written it is not very efficient since we need to assemble and invert the dense Schur complement.

Preconditioner performance
Preconditioner | Iterations
--- | ---
Schur complement | 4 (9.15e-14)

Since our discretisation is inf-sup stable it is possible to prove that the mass matrix of the pressure space is spectrally equivalent to the Schur complement. This means that we can use the mass matrix of the pressure space as a preconditioner for the Schur complement. Notice that we still need to invert the mass matrix and we will do so using a PETSc PC of type Jacobi, which is the exact inverse since we are using P0 pressure elements. We will also invert the Laplacian block using a PETSc PC of type LU.

    m = BilinearForm((1/nu)*p*q*dx).Assemble()
    mpre = Preconditioner(m, "PETScPC", pc_type="jacobi")
    apre = a.mat.Inverse(V.FreeDofs())
    C = BlockMatrix( [ [apre, None], [None, mpre] ] )

    gfu.vec.data[:] = 0; gfp.vec.data[:] = 0;
    gfu.Set(uin, definedon=mesh.Boundaries("inlet"))
    sol = BlockVector( [gfu.vec, gfp.vec] )

    print("-----------|Mass & LU|-----------")
    solvers.MinRes (mat=K, pre=C, rhs=rhs, sol=sol, tol=1e-8,
                    maxsteps=100, printrates=True, initialize=False)
    Draw(gfu)

Preconditioner performance
Preconditioner | Iterations
--- | ---
Schur complement | 4 (9.15e-14)
Mass & LU | 66 (2.45e-08)

We can also construct a multi-grid preconditioner for the top left block of the saddle point problem, as we have seen in Vertex patch smoothing for p-multigrid preconditioners for the Poisson problem.

    def DoFInfo(mesh, fes):
       blocks = []
       freedofs = fes.FreeDofs()
       vertexdofs = BitArray(fes.ndof)
       vertexdofs[:] = False
       for v in mesh.vertices:
          vdofs = set()
          vdofs |= set(d for d in fes.GetDofNrs(v) if freedofs[d])
          for ed in mesh[v].edges:
             vdofs |= set(d for d in fes.GetDofNrs(ed) if freedofs[d])
          for fc in mesh[v].faces:
             vdofs |= set(d for d in fes.GetDofNrs(fc) if freedofs[d])
          blocks.append(vdofs)
          for d in fes.GetDofNrs(v):
             vertexdofs[d] = True
       vertexdofs &= fes.FreeDofs()
       return vertexdofs, blocks

    vertexdofs, blocks = DoFInfo(mesh, V)
    blockjac = a.mat.CreateBlockSmoother(blocks)
    preH = PETScPreconditioner(a.mat, vertexdofs, solverParameters={"pc_type":"hypre"})
    twolvpre = preH + blockjac
    C = BlockMatrix( [ [twolvpre, None], [None, mpre] ] )
    gfu.vec.data[:] = 0; gfp.vec.data[:] = 0;
    gfu.Set(uin, definedon=mesh.Boundaries("inlet"))
    print("-----------|Mass & Two Level Additive Schwarz|-----------")
    solvers.MinRes (mat=K, pre=C, rhs=rhs, sol=sol, tol=1e-8,
                    maxsteps=100, printrates=True, initialize=False)

Preconditioner performance
Preconditioner | Iterations
--- | ---
Schur complement | 4 (9.15e-14)
Mass & LU | 66 (2.45e-08)
Mass & Two Level Additive Schwarz | 100 (4.68e-06)

The preconditioner constructed so far doesn’t seem to be ideal, in fact, our Krylov solver took many iterations to converge with a direct LU factorisation of the velocity block and did not converge at all with HYPRE. To resolve this issue we resort to an augmented Lagrangian formulation, i.e.

\[\begin{split}\begin{cases} (\nabla \vec{u},\nabla \vec{v})_{L^2(\Omega)} + (\nabla\cdot \vec{v}, p)_{L^2(\Omega)} + \gamma (\nabla\cdot \vec{u},\nabla\cdot\vec{v})_{L^2(\Omega)} = (\vec{f},\vec{v})_{L^2(\Omega)} \qquad \forall v\in H^1_{0}(\Omega)\\ (\nabla\cdot \vec{u},q)_{L^2(\Omega)} = 0 \qquad \forall q\in L^2(\Omega) \end{cases}\end{split}\]

This formulation can easily be constructed in NGSolve, as follows:

    gamma = Parameter(1e6)
    aG = BilinearForm(nu*InnerProduct(Grad(u),Grad(v))*dx+gamma*div(u)*div(v)*dx)
    aG.Assemble()
    aGpre = Preconditioner(aG, "PETScPC", pc_type="lu")
    mG = BilinearForm((1/nu+gamma)*p*q*dx).Assemble()
    mGpre = Preconditioner(mG, "PETScPC", pc_type="jacobi")

    K = BlockMatrix( [ [aG.mat, b.mat.T], [b.mat, None] ] )
    C = BlockMatrix( [ [aGpre.mat, None], [None, mGpre.mat] ] )

    gfu.vec.data[:] = 0; gfp.vec.data[:] = 0;
    gfu.Set(uin, definedon=mesh.Boundaries("inlet"))
    sol = BlockVector( [gfu.vec, gfp.vec] )

    print("-----------|Boffi--Lovadina Augmentation LU|-----------")
    solvers.MinRes (mat=K, pre=C, rhs=rhs, sol=sol, tol=1e-10,
                    printrates=True, initialize=False)
    Draw(gfu)

Using an augmented Lagrangian formulation and a direct inversion of the velocity block, we were able to converge in only two iterations. This is because the augmented Lagrangian formulation improves the spectral equivalence between the mass matrix of the pressure space and the Schur complement.

Preconditioner performance
Preconditioner | Iterations
--- | ---
Schur complement | 4 (9.15e-14)
Mass & LU | 66 (2.45e-08)
Mass & Two Level Additive Schwarz | 100 (4.68e-06)
Augmented Lagrangian LU | 2 (9.24e-8)

Notice that so far we have been inverting the matrix corresponding to the Laplacian block using a direct LU factorisation. This is not ideal for large problems, and we can try to use a HYPRE preconditioner for the Laplacian block.

    smoother = aG.mat.CreateBlockSmoother(blocks)
    preHG = PETScPreconditioner(aG.mat, vertexdofs, solverParameters={"pc_type":"hypre"})
    twolvpre = preHG + smoother
    C = BlockMatrix( [ [twolvpre, None], [None, mGpre] ] )
    gfu.vec.data[:] = 0; gfp.vec.data[:] = 0;
    gfu.Set(uin, definedon=mesh.Boundaries("inlet"))
    print("-----------|Boffi--Lovadina Augmentation Two Level Additive Schwarz|-----------")
    solvers.MinRes (mat=K, pre=C, rhs=rhs, sol=sol, tol=1e-10,
                    printrates=True, initialize=False)
    Draw(gfu)

Preconditioner performance
Preconditi¬oner | Iterations
--- | ---
Schur complement | 4 (9.15e-14)
Mass & LU | 66 (2.45e-08)
Mass & Two Level Additive Schwarz | 100 (4.68e-06)
Augmented Lagrangian LU | 2 (9.24e-08)
Augmented Two Level Additive Schwarz | 100 (1.06e-03)

Our first attempt at using a HYPRE preconditioner for the Laplacian block did not converge. This is because the top left block of the saddle point problem now contains the augmentation term, which has a very large kernel. It is well known that algebraic multi-grid methods do not work well for such problems, and this is what we are observing here. We begin by constructing the augmented Lagrangian formulation in more numerical linear algebra terms, i.e.

\[\begin{split}\begin{bmatrix} A + B^T (\gamma M^{-1}) B & B^T \\ B & 0 \end{bmatrix} \begin{bmatrix} u \\ p \end{bmatrix} = \begin{bmatrix} f \\ 0 \end{bmatrix}\end{split}\]

We can construct this linear algebra problem inside NGSolve as follows:

    d = BilinearForm((1/gamma)*p*q*dx)
    d.Assemble()
    dpre = PETScPreconditioner(d.mat, Q.FreeDofs(), solverParameters={"pc_type":"lu"})
    aG = a.mat + b.mat.T@dpre@b.mat
    aG = coo_matrix(aG.ToDense().NumPy())
    aG = la.SparseMatrixd.CreateFromCOO(indi=aG.row,
                                          indj=aG.col,
                                          values=aG.data,
                                          h=aG.shape[0],
                                          w=aG.shape[1])
    K = BlockMatrix( [ [aG, b.mat.T], [b.mat, None] ] )
    pre = PETScPreconditioner(aG, V.FreeDofs(), solverParameters={"pc_type":"lu"})
    C = BlockMatrix( [ [pre, None], [None, mGpre.mat] ] )

    gfu.vec.data[:] = 0; gfp.vec.data[:] = 0;
    gfu.Set(uin, definedon=mesh.Boundaries("inlet"))
    sol = BlockVector( [gfu.vec, gfp.vec] )

    print("-----------|Boffi--Lovadina Augmentation LU|-----------")
    solvers.MinRes (mat=K, pre=C, rhs=rhs, sol=sol, tol=1e-10,
                    printrates=True, initialize=False)
    Draw(gfu)

Since the augmentation block has a lower rank than the Laplacian block, we can use the Sherman-Morrison-Woodbury formula to invert the augmentation block.

\[(A + B^T(\gamma M^{-1})B)^{-1} = A^{-1} - A^{-1}B^T(\frac{1}{\gamma}M^{-1} + BA^{-1}B^T)^{-1}BA^{-1}\]

We will do this in two different ways first we will invert the \((\frac{1}{\gamma}M^{-1} + BA^{-1}B^T)\) block using a direct LU factorisation. Then we will notice that since the penalisation parameter is large we can ignore the \(\frac{1}{\gamma}M^{-1}\) term and use the mass matrix since it is spectrally equivalent to the Schur complement.

    SM = (d.mat + b.mat@apre@b.mat.T).ToDense().NumPy()
    SM = coo_matrix(SM)
    SM = la.SparseMatrixd.CreateFromCOO(indi=SM.row,
                                          indj=SM.col,
                                          values=SM.data,
                                          h=SM.shape[0],
                                          w=SM.shape[1])

    SMinv = PETScPreconditioner(SM, Q.FreeDofs(), solverParameters={"pc_type":"lu"})

    C = BlockMatrix( [ [apre - apre@b.mat.T@SMinv@b.mat@apre, None], [None, mGpre.mat] ] )

    gfu.vec.data[:] = 0; gfp.vec.data[:] = 0;
    gfu.Set(uin, definedon=mesh.Boundaries("inlet"))
    sol = BlockVector( [gfu.vec, gfp.vec] )

    print("-----------|Boffi--Lovadina Augmentation Sherman-Morrison-Woodbury|-----------")
    solvers.MinRes (mat=K, pre=C, rhs=rhs, sol=sol, tol=1e-10,
                    printrates=True, initialize=False)
    Draw(gfu)

    C = BlockMatrix( [ [apre + apre@(b.mat.T@mpre.mat@b.mat)@apre, None], [None, mGpre.mat] ] )

    gfu.vec.data[:] = 0; gfp.vec.data[:] = 0;
    gfu.Set(uin, definedon=mesh.Boundaries("inlet"))
    sol = BlockVector( [gfu.vec, gfp.vec] )

    print("-----------|Boffi--Lovadina Augmentation Sherman-Morrison-Woodbury|-----------")
    solvers.MinRes (mat=K, pre=C, rhs=rhs, sol=sol, tol=1e-13,
                    printrates=True, initialize=False)
    Draw(gfu)

We see that a purely algebraic approach based on the Sherman-Morrison-Woodbury formula is more efficient for the augmented Lagrangian formulation than a naive two-level additive Schwarz approach.

Preconditioner performance
Preconditioner | Iterations
--- | ---
Schur complement | 4 (9.15e-14)
Mass & LU | 66 (2.45e-08)
Mass & Two Level Additive Schwarz | 100 (4.68e-06)
Augmented Lagrangian LU | 2 (9.24e-08)
Augmented Two Level Additive Schwarz | 100 (1.06e-03)
Augmentation Schermon-Morrison-Woodbury | 84 (1.16e-07)