Replicating Firedrake SNES solver in dolfin 2019.1 with petsc4py

Hi all,
I have a slightly unusual goal of exactly replicating a firedrake simulation using dolfin 2019.1 with custom SNES problem and solver classes inspired primarily by this post. I will demonstrate my question with a minimal working example of a lid driven cavity Navier Stokes simulation. The following code excerpt is written with dolfin, and writing the analogous firedrake code is simply a matter of replacing `dl.` with `fe.`, except in the lines with comments that specify otherwise. I am running firedrake version 0.13.0+5443.g76237bef.dirty.

``````import dolfin as dl
import ufl
from petsc4py import PETSc
from petsc_utils import *

inner, dot, grad, div, sym, split = \
ufl.inner, ufl.dot, ufl.grad, ufl.div, ufl.sym, ufl.split

# initialize the simulation
reynolds_number=dl.Constant(10)
solver_parameters = {
"snes_monitor": None,
"ksp_monitor": None,
"ksp_norm_type":-1,
"snes_converged_reason": None,
"snes_type": "newtonls",
"snes_linesearch_type": "basic",
"ksp_type": "preonly",
"pc_type": "lu",
"mat_type": "aij",
"pc_factor_mat_solver_type": "mumps",
}

# initialize the mesh, element, solution space
mesh_dimensions=(50,50)
mesh = dl.UnitSquareMesh(*mesh_dimensions, diagonal="left")
taylor_hood_pressure_degree = 1
element = dl.MixedElement(
dl.FiniteElement("P", mesh.ufl_cell(), taylor_hood_pressure_degree),
dl.VectorElement("P", mesh.ufl_cell(), taylor_hood_pressure_degree + 1))

# solution is initialized trivally by default - u and p are both zero vectors
solution = dl.Function(dl.FunctionSpace(mesh, element))
solution_space = solution.function_space()

# pressure nullspace: this applies the constant 1 to pressure dofs, 0 to velocity dofs
# pressure = sub(0) = (1d) scalar field
# velocity = sub(1) = (2d) vector field
w=dl.interpolate(dl.Constant((1.0,0.0,0.0)),solution_space)
nullspace=dl.VectorSpaceBasis([w.vector().copy()])
nullspace.orthonormalize()
# in firedrake, we can simply use:
# nullspace=fe.MixedVectorSpaceBasis(solution_space,(fe.VectorSpaceBasis(constant=True),solution_space.sub(1)))

# dirichlet bc
left = 'near(x[0], 0)'
right = 'near(x[0], 1)'
bottom = 'near(x[1], 0)'
top = 'near(x[1], 1)'
# in firedrake, we have
# left = 1
# right = 2
# bottom = 3
# top = 4
noslip = dl.Constant((0.,)*mesh.geometric_dimension())
veltop=dl.Constant((1.0,0.0))
dirichlet_boundary_conditions=[dl.DirichletBC(solution_space.sub(1), noslip, left),
dl.DirichletBC(solution_space.sub(1), noslip, right),
dl.DirichletBC(solution_space.sub(1), noslip, bottom),
dl.DirichletBC(solution_space.sub(1), veltop, top),
]

# weak form residual
dx = ufl.dx(degree = quadrature_degree)
p,u=dl.split(solution)
psi_p,psi_u=dl.TestFunctions(solution_space)

# set up nonlinear problem
problem = SNESProblem(weak_form_residual,
solution,
dirichlet_boundary_conditions,
nullspace=nullspace)
# set up SNES solver and solve
snes=CustomSNESSolver(problem, solver_parameters, nullspace, dl.MPI.comm_world)
snes.solve()
# In firedrake, we'd instead use
# # set up nonlinear problem
# problem = fe.NonlinearVariationalProblem(
#     F = weak_form_residual,
#     u = solution,
#     bcs = dirichlet_boundary_conditions,
#     J = fe.derivative(weak_form_residual, solution))
# # set up SNES solver and solve
# solver = fe.NonlinearVariationalSolver(
#     problem = problem,
#     nullspace = nullspace,
#     solver_parameters = solver_parameters)
# solver.solve()
``````

the imported file `petsc_utils.py` contains the custom SNES Problem and Solver classes:

``````from dolfin import PETScVector, PETScMatrix, TestFunction, TestFunctions, TrialFunctions, TrialFunction, derivative, assemble_system, assemble, PETScOptions, inner
from dolfin import as_backend_type, MPI
from petsc4py import PETSc
import ufl

class SNESProblem():
def __init__(self, F, u, bcs, nullspace=None):
V = u.function_space()
self.V=V
du = TrialFunction(V)
self.u_test = TestFunction(V)
self.L = F
self.a = derivative(F, u, du)
self.p = derivative(F, u, du)
if len(V.ufl_element().sub_elements())==2:
ptest,_=TestFunctions(V)
ptrial,_=TrialFunctions(V)
elif len(V.ufl_element().sub_elements())==3:
ptest,_,_=TestFunctions(V)
ptrial,_,_=TrialFunctions(V)
self.bcs = bcs
self.u = u
self.nullspace=nullspace

def F(self, snes, x, F):
x = PETScVector(x)
F  = PETScVector(F)
x.vec().copy(self.u.vector().vec())
self.u.vector().apply("")
assemble(self.L, tensor=F)
for bc in self.bcs:
bc.apply(F, x)

def J(self, snes, x, J, P):
J = PETScMatrix(J)
P = PETScMatrix(P)
x.copy(self.u.vector().vec())
self.u.vector().apply("")
dummy = ufl.inner(self.u, self.u_test)*ufl.dx

# index set for pressure components, i.e. V.sub(0)
pc=snes.ksp.pc
fields=[]
subdofs=self.V.sub(0).dofmap().dofs()
IS=PETSc.IS().createGeneral(subdofs)
fields.append((str(0), IS))
pc.setFieldSplitIS(*fields)

# assemble
assemble_system(self.a, dummy, self.bcs, A_tensor = J)
assemble_system(self.p, dummy, self.bcs, A_tensor = P)

# attach nullspace
if self.nullspace is not None:
as_backend_type(J).set_nullspace(self.nullspace)
self.attach_nullspace(J, P)

# seems like I shouldnt have to do this bc loop given assemble_system takes bcs, yet solver gets stuck if you comment it out
for bc in self.bcs:
bc.apply(J)
bc.apply(P)

def attach_nullspace(self, J, P):
# attach nullspace for pressure dofs
self.comm=MPI.comm_world
nsp=PETSc.NullSpace().create(comm=self.comm)
J.mat().setNullSpace(nsp)
P.mat().setNullSpace(nsp)
is_=PETSc.IS().createGeneral(self.u.sub(0).function_space().dofmap().dofs(),comm=self.comm)
is_.compose('nullspace',nsp)

class CustomSNESSolver():
def __init__(self, problem, solver_parameters, nullspace, comm):
self.problem=problem
self.solver_parameters=solver_parameters
self.nullspace=nullspace
self.comm=comm

# create SNES solver
self.snes = PETSc.SNES().create(self.comm)

# set PETSc options from solver_parameters dict
self.set_petsc_options_from_dict(self.solver_parameters)
self.snes.setFromOptions()

# set function, jacobian
b = PETScVector()
J_mat = PETScMatrix()
P_mat = PETScMatrix()
self.snes.setFunction(self.problem.F, b.vec())
self.snes.setJacobian(self.problem.J, J_mat.mat(), P_mat.mat())

# save matrices for debugging when SNES.Problem is called during solve()
self.b=b
self.J_mat=J_mat
self.P_mat=P_mat

def solve(self):
# apply bcs
for bc in self.problem.bcs:
bc.apply(self.problem.u.vector())
self.snes.solve(None, self.problem.u.vector().vec())

def set_petsc_options_from_dict(self,solver_parameters_dict):
for key, value in solver_parameters_dict.items():
if value is None:
PETScOptions.set(key)
else:
PETScOptions.set(key,value)
``````

In dolfin, I get the following output:

``````  0 SNES Function norm 1.347013156745e+00
0 KSP Residual norm 1.347013156745e+00
1 KSP Residual norm 6.598859834202e-14
1 SNES Function norm 5.785669420919e-03
0 KSP Residual norm 5.785669420919e-03
1 KSP Residual norm 1.378854855395e-13
2 SNES Function norm 9.341656752388e-06
0 KSP Residual norm 9.341656752388e-06
1 KSP Residual norm 2.353857907220e-13
3 SNES Function norm 8.756520653743e-12
Nonlinear solve converged due to CONVERGED_FNORM_RELATIVE iterations 3
``````

whereas in firedrake, I get

``````  0 SNES Function norm 1.347013156745e+00
0 KSP Residual norm 1.347013156745e+00
1 KSP Residual norm 1.571543510327e-12
1 SNES Function norm 5.785669420917e-03
0 KSP Residual norm 5.785669420917e-03
1 KSP Residual norm 2.467259171461e-13
2 SNES Function norm 9.341656753825e-06
0 KSP Residual norm 9.341656753825e-06
1 KSP Residual norm 8.956712819439e-13
3 SNES Function norm 8.070849898795e-12
Nonlinear firedrake_0_ solve converged due to CONVERGED_FNORM_RELATIVE iterations 3
``````

In both cases, we have convergence, but I am wondering why the SNES function norms differ slightly by the third iteration between firedrake and dolfin. This has consequences for me later on when running a more complex Navier Stokes based simulation, so I am wondering if someone can offer tips to help get these two to be identical. My guess is I need to change something in my `SNESProblem.J` function. Note that I use `dl.assemble_system` instead of `dl.assemble`, because both firedrake and `dl.assemble_system` perform symmetric elimination of Dirichlet boundary conditions. I have gone in with a debugger and poked around inside `SNESProblem.J` and in the firedrake source code. When I assemble my Jacobians and convert them to dense matrices (numpy arrays), they look to match up to permutation of columns/rows (presumably because dolfin orders dofs differently than firedrake). Also, running `snes.view()` looks to show that the solver parameters are getting set the same in both codes.

Thank you for your time and assistance!