Issues with using self created PETSc mats instead of dolfinx PETSc mats in SNES

Linear Algebra
1

Hi, as I understand the “Real” element hasn’t yet been implemented in Dolfinx. So to solve Navier-Stokes equations coupled with rigid body motion equations in a Newton solver I have gone to assemble the NS equations using dolfinx and adding the four extra degrees of freedom myself onto the PETSc matrix using CSR matrix to feed into PETSc4Py’s SNES via a class interfacing with SNES. I can’t seem to get the problem to converge, so I played around with solving it with just the NS equations and it converges with a dolfinx create_matrix_block but not a PETSc.Mat().createAIJ() that is based on the size of the create_matrix_block in snes.setJacobian(). In fact, there seems to be different behaviour when I mix and match empty PETSc created matrices/vectors with dolfinx created matrices/vectors. As such I don’t think it’s a problem with my equations/mesh but rather with how I implement dolfinx/petsc. Below I have an MWE for a non-linear Poisson equation (from the tutorials) that Seg Faults for some unknown reason (see first example), and another MWE where I actually get and reassemble a CSR that just doesn’t converge.

Here is the MWE that Seg Faults by using a PETSc.Mat().createAIJ()

import ufl
import numpy

from mpi4py import MPI
from petsc4py import PETSc

from dolfinx import mesh, fem, io, nls, log

def q(u):
    return 1 + u**2

domain = mesh.create_unit_square(MPI.COMM_WORLD, 10, 10)
x = ufl.SpatialCoordinate(domain)
u_ufl = 1 + x[0] + 2*x[1]
f = - ufl.div(q(u_ufl)*ufl.grad(u_ufl))
V = fem.FunctionSpace(domain, ("CG", 1))
u_exact = lambda x: eval(str(u_ufl))
u_D = fem.Function(V)
u_D.interpolate(u_exact)
fdim = domain.topology.dim - 1
boundary_facets = mesh.locate_entities_boundary(domain, fdim, lambda x: numpy.full(x.shape[1], True, dtype=bool))
bc = fem.dirichletbc(u_D, fem.locate_dofs_topological(V, fdim, boundary_facets))

uh = fem.Function(V)
v = ufl.TestFunction(V)
F = q(uh)*ufl.dot(ufl.grad(uh), ufl.grad(v))*ufl.dx - f*v*ufl.dx

class NonlinearPDE_SNESProblem:
    def __init__(self, F, u, bc):
        V = u.function_space
        du = ufl.TrialFunction(V)
        self.L = fem.form(F)
        self.a = fem.form(ufl.derivative(F, u, du))
        self.bc = bc
        self._F, self._J = None, None
        self.u = u

    def F(self, snes, x, F):
        """Assemble residual vector."""
        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        x.copy(self.u.vector)
        self.u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)

        with F.localForm() as f_local:
            f_local.set(0.0)
        fem.petsc.assemble_vector(F, self.L)
        fem.petsc.apply_lifting(F, [self.a], bcs=[[self.bc]], x0=[x], scale=-1.0)
        F.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
        fem.petsc.set_bc(F, [self.bc], x, -1.0)

    def J(self, snes, x, J, P):
        """Assemble Jacobian matrix."""
        J.zeroEntries()
        fem.petsc.assemble_matrix(J, self.a, bcs=[self.bc])
        J.assemble()

problem = NonlinearPDE_SNESProblem(F, uh, bc)
J_mat_dolfinx = fem.petsc.create_matrix(problem.a)
F_vec_dolfinx = fem.petsc.create_vector(problem.L)
J_mat = PETSc.Mat().createAIJ(J_mat_dolfinx.getSizes())
J_mat.setPreallocationNNZ(30)
J_mat.setUp()
F_vec = J_mat.createVecRight()

snes = PETSc.SNES().create(MPI.COMM_WORLD)
snes.getKSP().setType("preonly")
snes.getKSP().getPC().setType("lu")
snes.getKSP().getPC().setFactorSolverType("mumps")
snes.setFunction(problem.F, F_vec)

# This works and converges in 5 iterations
#snes.setJacobian(problem.J, J=J_mat_dolfinx, P=None)
# This doesn't
snes.setJacobian(problem.J, J=J_mat, P=None)

snes.setMonitor(lambda _, it, residual: print(it, residual))
solution = fem.petsc.create_vector(problem.L)
snes.solve(None, solution)
uh.x.array[:] = solution.getArray()[:]

Here is the MWE that doesn’t converge:

import sys
import ufl
import numpy

from mpi4py import MPI
from petsc4py import PETSc

from dolfinx import mesh, fem, io, nls, log

def q(u):
    return 1 + u**2

domain = mesh.create_unit_square(MPI.COMM_WORLD, 10, 10)
x = ufl.SpatialCoordinate(domain)
u_ufl = 1 + x[0] + 2*x[1]
V = fem.FunctionSpace(domain, ("CG", 1))
u_exact = lambda x: eval(str(u_ufl))
u_D = fem.Function(V)
u_D.interpolate(u_exact)
fdim = domain.topology.dim - 1
boundary_facets = mesh.locate_entities_boundary(domain, fdim, lambda x: numpy.full(x.shape[1], True, dtype=bool))
bc = fem.dirichletbc(u_D, fem.locate_dofs_topological(V, fdim, boundary_facets))

uh = fem.Function(V)
v = ufl.TestFunction(V)
F = q(uh)*ufl.dot(ufl.grad(uh), ufl.grad(v))*ufl.dx

class NonlinearPDE_SNESProblem:
    def __init__(self, F, u, bc):
        V = u.function_space
        du = ufl.TrialFunction(V)
        self.L = fem.form(F)
        self.a = fem.form(ufl.derivative(F, u, du))
        self.bc = bc
        self._F, self._J = None, None
        self.u = u
        self.J_mat_dolfinx = fem.petsc.create_matrix(self.a)

    def F(self, snes, x, F):
        """Assemble residual vector."""
        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        x.copy(self.u.vector)
        self.u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)

        with F.localForm() as f_local:
            f_local.set(0.0)
        fem.petsc.assemble_vector(F, self.L)
        fem.petsc.apply_lifting(F, [self.a], bcs=[[self.bc]], x0=[x], scale=-1.0)
        F.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
        fem.petsc.set_bc(F, [self.bc], x, -1.0)

    def J(self, snes, x, J, P):
        """Assemble Jacobian matrix."""
        J.zeroEntries()
        fem.petsc.assemble_matrix(self.J_mat_dolfinx, self.a, bcs=[self.bc])
        self.J_mat_dolfinx.assemble()
        ai, aj, av = self.J_mat_dolfinx.getValuesCSR()
        J.setPreallocationNNZ(numpy.diff(ai))
        J.setValuesCSR(ai,aj,av)
        J.assemble()

problem = NonlinearPDE_SNESProblem(F, uh, bc)
J_mat_dolfinx = fem.petsc.create_matrix(problem.a)
F_vec_dolfinx = fem.petsc.create_vector(problem.L)
J_mat = PETSc.Mat().createAIJ(J_mat_dolfinx.getSizes())
J_mat.setPreallocationNNZ(30)
J_mat.setUp()
F_vec = J_mat.createVecRight()

snes = PETSc.SNES().create(MPI.COMM_WORLD)
snes.setTolerances(max_it=1000, atol=1.e-14, rtol=1.e-14)
snes.getKSP().setType("preonly")
snes.getKSP().getPC().setType("lu")
snes.getKSP().getPC().setFactorSolverType("mumps")
snes.setFunction(problem.F, F_vec)

# Now this converges realatively but not absolutely and in 670 iterations)
#snes.setJacobian(problem.J,  J=J_mat_dolfinx, P=None)
# And this has the exact same behaviour as using uncommenting the above line (i.e. no seg fault)
snes.setJacobian(problem.J, J=J_mat, P=None)

snes.setMonitor(lambda _, it, residual: print(it, residual))
solution = fem.petsc.create_vector(problem.L)
snes.solve(None, solution)
uh.x.array[:] = solution.getArray()[:]

Also saw this thread which is similar to my first problem, but the solution is untenable as I need to change the size of the PETSc vector/matrix to add degrees of freedom.

2

The first problem seems to be fixable by adding

J_mat.setLGMap(J_mat_dolfinx.getLGMap()[0],J_mat_dolfinx.getLGMap()[1])

after the creation of J_mat. This method converges in 8 steps instead of the 5 it took to converge with J_mat_dolfinx for this problem. This does not fix the second problem with CSR. Moreover I’m not sure how to edit an LGMap when I need to add degrees of freedom manually.

3

I think I may have a solution for the second problem. It is apparently necessary to set self.J_mat_dolfinx to zero with self.J_mat_dolfinx.zeroEntries() before the assemble_matrix call in the self.J method setJacobian i.e. inserting

def J(self, snes, x, J, P):
        """Assemble Jacobian matrix."""
        J.zeroEntries()
        self.J_mat_dolfinx.zeroEntries()
        fem.petsc.assemble_matrix(self.J_mat_dolfinx, self.a, bcs=[self.bc])

solves the issue.

It is unclear to me why this is the case because the same non-convergence occurs when I assemble into self.J_mat_dolfinx and then copy into J with self.J_mat_dolfinx.copy(J) and checking that they are equal with self.J_mat_dolfinx.equal(J). But for some reason setting self.J_mat_dolfinx.zeroEntries() beforehand will make it converge so I will stick with it. It is also curiously unnecessary to use setLGMap() in this case.

Moreover, the same thing happens with the residual if instead of assembling into F as it is in my MWE, we create a self.F_vec_dolfinx = fem.petsc.create_vector(self.L) and assemble into that in the method self.F then it is also necessary to set self.F_vec_dolfinx.set(0.) (or with localForm()) before assemble_vector(self.F_vec_dolfinx, self.L). In all, here is a final form of my MWE that works

import sys
import ufl
import numpy

from mpi4py import MPI
from petsc4py import PETSc

from dolfinx import mesh, fem, io, nls, log
PETSc.Options().setValue('-snes_linesearch_type', 'basic')

def q(u):
    return 1 + u**2

domain = mesh.create_unit_square(MPI.COMM_WORLD, 10, 10)
x = ufl.SpatialCoordinate(domain)
u_ufl = 1 + x[0] + 2*x[1]
f = - ufl.div(q(u_ufl)*ufl.grad(u_ufl))
V = fem.FunctionSpace(domain, ("CG", 1))
u_exact = lambda x: eval(str(u_ufl))
u_D = fem.Function(V)
u_D.interpolate(u_exact)
fdim = domain.topology.dim - 1
boundary_facets = mesh.locate_entities_boundary(domain, fdim, lambda x: numpy.full(x.shape[1], True, dtype=bool))
bc = fem.dirichletbc(u_D, fem.locate_dofs_topological(V, fdim, boundary_facets))

uh = fem.Function(V)
v = ufl.TestFunction(V)
F = q(uh)*ufl.dot(ufl.grad(uh), ufl.grad(v))*ufl.dx - f*v*ufl.dx

class NonlinearPDE_SNESProblem:
    def __init__(self, F, u, bc):
        V = u.function_space
        du = ufl.TrialFunction(V)
        self.L = fem.form(F)
        self.a = fem.form(ufl.derivative(F, u, du))
        self.bc = bc
        self._F, self._J = None, None
        self.u = u
        self.J_mat_dolfinx = fem.petsc.create_matrix(self.a)
        self.F_vec_dolfinx = fem.petsc.create_vector(self.L)

    def F(self, snes, x, F):
        """Assemble residual vector."""
        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        x.copy(self.u.vector)
        self.u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)

        self.F_vec_dolfinx.set(0.)
        with F.localForm() as f_local:
            f_local.set(0.0)
        fem.petsc.assemble_vector(self.F_vec_dolfinx, self.L)
        fem.petsc.apply_lifting(self.F_vec_dolfinx, [self.a], bcs=[[self.bc]], x0=[x], scale=-1.0)
        self.F_vec_dolfinx.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
        fem.petsc.set_bc(self.F_vec_dolfinx, [self.bc], x, -1.0)
        F.getArray()[:] = self.F_vec_dolfinx.getArray()[:]

    def J(self, snes, x, J, P):
        """Assemble Jacobian matrix."""
        J.zeroEntries()
        self.J_mat_dolfinx.zeroEntries()

        fem.petsc.assemble_matrix(self.J_mat_dolfinx, self.a, bcs=[self.bc])
        self.J_mat_dolfinx.assemble()
        ai, aj, av = self.J_mat_dolfinx.getValuesCSR()
        J.setPreallocationNNZ(numpy.diff(ai))
        J.setValuesCSR(ai,aj,av)
        J.assemble()

problem = NonlinearPDE_SNESProblem(F, uh, bc)
J_mat_dolfinx = fem.petsc.create_matrix(problem.a)
F_vec_dolfinx = fem.petsc.create_vector(problem.L)
J_mat = PETSc.Mat().createAIJ(J_mat_dolfinx.getSizes())
#J_mat.setLGMap(J_mat_dolfinx.getLGMap()[0],J_mat_dolfinx.getLGMap()[1]) # Only needed when assembling straight into J in problem.J
J_mat.setPreallocationNNZ(30)
J_mat.setUp()
F_vec = J_mat.createVecRight()

snes = PETSc.SNES().create(MPI.COMM_WORLD)
snes.setTolerances(max_it=4000)
snes.getKSP().setType("preonly")
snes.getKSP().getPC().setType("lu")
snes.getKSP().getPC().setFactorSolverType("mumps")
snes.setFunction(problem.F, F_vec)

snes.setJacobian(problem.J, J=J_mat, P=None)

snes.setMonitor(lambda _, it, residual: print(it, residual))
solution = fem.petsc.create_vector(problem.L)
snes.solve(None, solution)
uh.x.array[:] = solution.getArray()[:]