Using petsc4py.PETSc.SNES directly

Hi,
I would like to interface with the petsc4py.PETSc.SNES object directly to access some specific PETSc functions, but I cannot find any documentation on how to use to the SNES object in FEniCS e.g. how to convert the linear form F, Jacobian J, and the boundary conditions to PETSc objects etc. Take for example the fenics hyperelasticty demo as a starting point:

from dolfin import *
# Optimization options for the form compiler
parameters["form_compiler"]["cpp_optimize"] = True
ffc_options = {"optimize": True, \
               "eliminate_zeros": True, \
               "precompute_basis_const": True, \
               "precompute_ip_const": True}

# Create mesh and define function space
mesh = UnitCubeMesh(24, 16, 16)
V = VectorFunctionSpace(mesh, "Lagrange", 1)
# Mark boundary subdomians
left =  CompiledSubDomain("near(x[0], side) && on_boundary", side = 0.0)
right = CompiledSubDomain("near(x[0], side) && on_boundary", side = 1.0)
# Define Dirichlet boundary (x = 0 or x = 1)
c = Expression(("0.0", "0.0", "0.0"), degree=2)
r = Expression(("scale*0.0",
                "scale*(y0 + (x[1] - y0)*cos(theta) - (x[2] - z0)*sin(theta) - x[1])",
                "scale*(z0 + (x[1] - y0)*sin(theta) + (x[2] - z0)*cos(theta) - x[2])"),
                scale = 0.5, y0 = 0.5, z0 = 0.5, theta = pi/3, degree=2)
bcl = DirichletBC(V, c, left)
bcr = DirichletBC(V, r, right)
bcs = [bcl, bcr]
# Define functions
du = TrialFunction(V)            # Incremental displacement
v  = TestFunction(V)             # Test function
u  = Function(V)                 # Displacement from previous iteration
B  = Constant((0.0, -0.5, 0.0))  # Body force per unit volume
T  = Constant((0.1,  0.0, 0.0))  # Traction force on the boundary
# Kinematics
d = u.geometric_dimension()
I = Identity(d)             # Identity tensor
F = I + grad(u)             # Deformation gradient
C = F.T*F                   # Right Cauchy-Green tensor
# Invariants of deformation tensors
Ic = tr(C)
J  = det(F)
# Elasticity parameters
E, nu = 10.0, 0.3
mu, lmbda = Constant(E/(2*(1 + nu))), Constant(E*nu/((1 + nu)*(1 - 2*nu)))
# Stored strain energy density (compressible neo-Hookean model)
psi = (mu/2)*(Ic - 3) - mu*ln(J) + (lmbda/2)*(ln(J))**2
# Total potential energy
Pi = psi*dx - dot(B, u)*dx - dot(T, u)*ds

# Compute first variation of Pi (directional derivative about u in the direction of v)
F = derivative(Pi, u, v)

# Compute Jacobian of F
J = derivative(F, u, du)
solve(F == 0, u, bcs, J=J, form_compiler_parameters=ffc_options)

Now, to replace the magic solve call with a custom SNES solver… I know how to use NonlinearProblem together with PETScSNESSolver to replace the magic solve as below, but this is NOT what I want:

class Problem(NonlinearProblem):
    def __init__(self, J, F, bcs):
        self.bilinear_form = J
        self.linear_form = F
        self.bcs = bcs
        NonlinearProblem.__init__(self)

    def F(self, b, x):
        assemble(self.linear_form, tensor=b)
        for bc in self.bcs:
            bc.apply(b, x)

    def J(self, A, x):
        assemble(self.bilinear_form, tensor=A)
        for bc in self.bcs:
            bc.apply(A)
            
problem = Problem(J, F, bcs)
solver  = PETScSNESSolver()
solver.solve(problem, u.vector())

Instead I would like to use petsc4py directly as suggested by @nate here:
Set krylov linear solver paramters in newton solver starting with:

import petsc4py, sys
petsc4py.init(sys.argv)
from petsc4py import PETSc
snes = PETSc.SNES().create(MPI.comm_world)
....

Somehow somewhere there have to be calls to snes.setFunction() and snes.setJacobian() and handling bcs etc:

snes.setFunction(???)
snes.setJacobian(???)
....
snes.solve(??)

Does anyone have any code that does something like this that they could share? It would be much appreciated! Thanks!

2 Likes

This isn’t exactly what you need, but you can use it to figure out how to do it in old dolfin.

Using a PETSc SNES directly with dolfinx here. This shows how to set up the SNES.

The class to interface with the SNES solver here. Should be very similar to the old dolfin custom nonlinear problem that you linked to.

Nate,
What is the old dolfin equivalent of:

b = dolfinx.cpp.la.create_vector(V.dofmap.index_map)
J = dolfinx.cpp.fem.create_matrix(problem.a_comp._cpp_object)

also what are the calls to vector.ghostUpdate ? are they necessary in the old dolfin?
Great to see the increased transparency wrt PETSc in dolfinx (although I’m sure there will be a learning curve to go along with it…)
Thanks!

b = PETScVector()
J = PETScMatrix()

old dolfin does ghost updates behind the scenes (although many are very unnecessary). You don’t need to worry about it if you’re not doing anything too crazy.

Thanks, now I’m getting a TypeError TypeError: F() takes 3 positional arguments but 4 were given when I call solve. What did I miss?
(plugging into end of demo code:)

class Problem(NonlinearProblem):
    def __init__(self, J, F, bcs):
        self.bilinear_form = J
        self.linear_form = F
        self.bcs = bcs
        NonlinearProblem.__init__(self)

    def F(self, b, x):
        assemble(self.linear_form, tensor=b)
        for bc in self.bcs:
            bc.apply(b, x)

    def J(self, A, x):
        assemble(self.bilinear_form, tensor=A)
        for bc in self.bcs:
            bc.apply(A)
    
import petsc4py, sys
petsc4py.init(sys.argv)
from petsc4py import PETSc
    
problem = Problem(J, F, bcs)
b = PETScVector().vec()
J_mat = PETScMatrix().mat()
snes = PETSc.SNES().create(MPI.comm_world)    
snes.setFunction(problem.F, b)
snes.setJacobian(problem.J, J_mat)
snes.solve(None, u.vector().vec())

The snes context variable. See this line.

Also be careful about the order of your arguments and what PETSc SNES expects.

Do I have to create my own snes context? It seem that NonlinearProblem doesn’t let me change the signature of F

Good point. Your Problem doesn’t need to inherit NonlinearProblem. See this line. SNES just needs to know how to formulate the residual and Jacobian, which you’ve defined through the F and J methods.

Ok, I followed the dolfinx example and made a custom SNESProblem but now my python kernel crashes when I try to run the following code, so I’m still missing something…

class SNESProblem():
    def __init__(self, F, u, bcs):
        V = u.function_space()
        du = TrialFunction(V)
        self.L = F
        self.a = derivative(F, u, du)
        self.bcs = bcs
        self.u = u

    def F(self, snes, x, F):
        F  = PETScVector()  # had to create a new vector here b/c assemble was receiving PETSc vec instead
        assemble(self.L, tensor=F)
        for bc in self.bcs:
            bc.apply(F)                

    def J(self, snes, x, J):
        J = PETScMatrix()
        assemble(J, self.a)
        for bc in self.bcs:
            bc.apply(A)
            
problem = SNESProblem(F, u, bcs)
import petsc4py, sys
petsc4py.init(sys.argv)
from petsc4py import PETSc
    
b = PETScVector()  # same as b = PETSc.Vec()
J_mat = PETScMatrix()

snes = PETSc.SNES().create(MPI.comm_world)    
snes.setFunction(problem.F, b.vec())
snes.setJacobian(problem.J, J_mat.mat())
snes.solve(None, problem.u.vector().vec())

Can you turn your code into a MWE and I’ll take a look.

Hi Nate, thanks for offering to help, my MWE is below (top is directly from hyperelasticity demo):

from dolfin import *

parameters["form_compiler"]["cpp_optimize"] = True
ffc_options = {"optimize": True, \
               "eliminate_zeros": True, \
               "precompute_basis_const": True, \
               "precompute_ip_const": True}

# Create mesh and define function space
mesh = UnitCubeMesh(24, 16, 16)
V = VectorFunctionSpace(mesh, "Lagrange", 1)

# Mark boundary subdomians
left =  CompiledSubDomain("near(x[0], side) && on_boundary", side = 0.0)
right = CompiledSubDomain("near(x[0], side) && on_boundary", side = 1.0)

# Define Dirichlet boundary (x = 0 or x = 1)
c = Expression(("0.0", "0.0", "0.0"), degree=2)
r = Expression(("scale*0.0",
                "scale*(y0 + (x[1] - y0)*cos(theta) - (x[2] - z0)*sin(theta) - x[1])",
                "scale*(z0 + (x[1] - y0)*sin(theta) + (x[2] - z0)*cos(theta) - x[2])"),
                scale = 0.5, y0 = 0.5, z0 = 0.5, theta = pi/3, degree=2)

bcl = DirichletBC(V, c, left)
bcr = DirichletBC(V, r, right)
bcs = [bcl, bcr]

# Define functions
du = TrialFunction(V)            # Incremental displacement
v  = TestFunction(V)             # Test function
u  = Function(V)                 # Displacement from previous iteration
B  = Constant((0.0, -0.5, 0.0))  # Body force per unit volume
T  = Constant((0.1,  0.0, 0.0))  # Traction force on the boundary

# Kinematics
d = u.geometric_dimension()
I = Identity(d)             # Identity tensor
F = I + grad(u)             # Deformation gradient
C = F.T*F                   # Right Cauchy-Green tensor

# Invariants of deformation tensors
Ic = tr(C)
J  = det(F)

# Elasticity parameters
E, nu = 10.0, 0.3
mu, lmbda = Constant(E/(2*(1 + nu))), Constant(E*nu/((1 + nu)*(1 - 2*nu)))

# Stored strain energy density (compressible neo-Hookean model)
psi = (mu/2)*(Ic - 3) - mu*ln(J) + (lmbda/2)*(ln(J))**2

# Total potential energy
Pi = psi*dx - dot(B, u)*dx - dot(T, u)*ds

# Compute first variation of Pi (directional derivative about u in the direction of v)
F = derivative(Pi, u, v)

class SNESProblem():
    def __init__(self, F, u, bcs):
        V = u.function_space()
        du = TrialFunction(V)
        self.L = F
        self.a = derivative(F, u, du)
        self.bcs = bcs
        self.u = u

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

    def J(self, snes, x, J):
        J = PETScMatrix()
        assemble(J, self.a)
        for bc in self.bcs:
            bc.apply(A)
            
problem = SNESProblem(F, u, bcs)
import petsc4py, sys
petsc4py.init(sys.argv)
from petsc4py import PETSc
    
b = PETScVector()  # same as b = PETSc.Vec()
J_mat = PETScMatrix()

snes = PETSc.SNES().create(MPI.comm_world)    
snes.setFunction(problem.F, b.vec())
snes.setJacobian(problem.J, J_mat.mat())
snes.solve(None, problem.u.vector().vec())

Ahhh, you were so close. Just had to use the dolfin wrappers for the PETSc objects, apply the BCs to the residual in a homogeneous sense, and allow SNES to pass the preconditioner matrix P also.

Corrected code follows (where I’ve used the excellent vtkplotter for visualisation which you can get with pip3 install vtkplotter --user).

from dolfin import *

parameters["form_compiler"]["cpp_optimize"] = True
ffc_options = {"optimize": True, \
               "eliminate_zeros": True, \
               "precompute_basis_const": True, \
               "precompute_ip_const": True}

# Create mesh and define function space
mesh = UnitCubeMesh(24, 16, 16)
V = VectorFunctionSpace(mesh, "Lagrange", 1)

# Mark boundary subdomians
left =  CompiledSubDomain("near(x[0], side) && on_boundary", side = 0.0)
right = CompiledSubDomain("near(x[0], side) && on_boundary", side = 1.0)

# Define Dirichlet boundary (x = 0 or x = 1)
c = Expression(("0.0", "0.0", "0.0"), degree=2)
r = Expression(("scale*0.0",
                "scale*(y0 + (x[1] - y0)*cos(theta) - (x[2] - z0)*sin(theta) - x[1])",
                "scale*(z0 + (x[1] - y0)*sin(theta) + (x[2] - z0)*cos(theta) - x[2])"),
                scale = 0.5, y0 = 0.5, z0 = 0.5, theta = pi/3, degree=2)

bcl = DirichletBC(V, c, left)
bcr = DirichletBC(V, r, right)
bcs = [bcl, bcr]

# Define functions
du = TrialFunction(V)            # Incremental displacement
v  = TestFunction(V)             # Test function
u  = Function(V)                 # Displacement from previous iteration
B  = Constant((0.0, -0.5, 0.0))  # Body force per unit volume
T  = Constant((0.1,  0.0, 0.0))  # Traction force on the boundary

# Kinematics
d = u.geometric_dimension()
I = Identity(d)             # Identity tensor
F = I + grad(u)             # Deformation gradient
C = F.T*F                   # Right Cauchy-Green tensor

# Invariants of deformation tensors
Ic = tr(C)
J  = det(F)

# Elasticity parameters
E, nu = 10.0, 0.3
mu, lmbda = Constant(E/(2*(1 + nu))), Constant(E*nu/((1 + nu)*(1 - 2*nu)))

# Stored strain energy density (compressible neo-Hookean model)
psi = (mu/2)*(Ic - 3) - mu*ln(J) + (lmbda/2)*(ln(J))**2

# Total potential energy
Pi = psi*dx - dot(B, u)*dx - dot(T, u)*ds

# Compute first variation of Pi (directional derivative about u in the direction of v)
F = derivative(Pi, u, v)

class SNESProblem():
    def __init__(self, F, u, bcs):
        V = u.function_space()
        du = TrialFunction(V)
        self.L = F
        self.a = derivative(F, u, du)
        self.bcs = bcs
        self.u = u

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

    def J(self, snes, x, J, P):
        J = PETScMatrix(J)
        assemble(self.a, tensor=J)
        for bc in self.bcs:
            bc.apply(J)
            
problem = SNESProblem(F, u, bcs)
import petsc4py, sys
petsc4py.init(sys.argv)
from petsc4py import PETSc
    
b = PETScVector()  # same as b = PETSc.Vec()
J_mat = PETScMatrix()

snes = PETSc.SNES().create(MPI.comm_world)    
snes.setFunction(problem.F, b.vec())
snes.setJacobian(problem.J, J_mat.mat())
snes.solve(None, problem.u.vector().vec())

from vtkplotter.dolfin import plot
plot(u, mode="displace")

2 Likes

Thanks Nate, the code runs but I don’t get the same result as you (using fenics 2019, see attached). Vtkplotter is great (can confirm it’s not a plotting issue, as it looks the same in Paraview). Running with dolfin.solve() instead gives me the same result you posted…

Make sure you set the PETSc options to something sensible. I’m using whatever the default options PETSc gives in version 3.12. Maybe make sure it’s using a direct linear solver, and check the Newton residual at each iteration to ensure it’s converged.

SNES options
KSP options

1 Like

Thanks! Calling snes.setFromOptions() after creating the snes object did the trick for me (using PETSc 3.10.5)

@nate Hi, I’m still having convergence issues. In my last post, I thought the code worked, but because I ran it in Spyder, it actually reused the solution from dolfin.solve() as the initial guess. I am using all default PETSc options as far as I know and running your code verbatim. (I upgraded to PETSc 3.12.) I also tried manually setting all the settings per PETScSNESSolver.cpp. To troubleshoot I tried snes.getKSP().getPC().setType("lu") to make sure the linear solve isn’t the issue, so that only leaves the jacobian calculation or solver set-up as potential issues. With default settings, I get convergence code -6 (line search failed). Can you please confirm you can run the code without issues and share your solver settings snes.view() and snes.getConvergedReason()? It would greatly appreciate it - thanks!

1 Like

@niewiarowski Hi, sorry for the intrusion but were you able to figure this out? I’m getting the exact same problem:

Nonlinear solve did not converge due to DIVERGED_LINE_SEARCH iterations 1

Apparently, SNES is not picking up the residual function correctly for some reason, these are the linesearch iterations:

      Line search: gnorm after quadratic fit 2.031122937974e-01
      Line search: Cubic step no good, shrinking lambda, current gnorm 2.031122937974e-01 lambda=1.0000000000000002e-02
      Line search: Cubic step no good, shrinking lambda, current gnorm 2.031122937974e-01 lambda=1.0000000000000002e-03
      Line search: Cubic step no good, shrinking lambda, current gnorm 2.031122937974e-01 lambda=1.0000000000000003e-04
      Line search: Cubic step no good, shrinking lambda, current gnorm 2.031122937974e-01 lambda=1.0000000000000004e-05
      Line search: Cubic step no good, shrinking lambda, current gnorm 2.031122937974e-01 lambda=1.0000000000000004e-06
      Line search: Cubic step no good, shrinking lambda, current gnorm 2.031122937974e-01 lambda=1.0000000000000005e-07
      Line search: Cubic step no good, shrinking lambda, current gnorm 2.031122937974e-01 lambda=1.0000000000000005e-08
      Line search: Cubic step no good, shrinking lambda, current gnorm 2.031122937974e-01 lambda=1.0000000000000005e-09
      Line search: Cubic step no good, shrinking lambda, current gnorm 2.031122937974e-01 lambda=1.0000000000000006e-10
      Line search: Cubic step no good, shrinking lambda, current gnorm 2.031122937974e-01 lambda=1.0000000000000006e-11
      Line search: Cubic step no good, shrinking lambda, current gnorm 2.031122937974e-01 lambda=1.0000000000000006e-12
      Line search: Cubic step no good, shrinking lambda, current gnorm 2.031122937974e-01 lambda=1.0000000000000007e-13

Best!
Nicolas

Ok, in case someone else arrives: the problem is that SNES moves function evaluation in the ‘x’ vector, whereas the residuals depend on ‘u’, so the solution is to update things correspondingly:

    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)
            bc.apply(F, self.u.vector())



    def J(self, snes, x, J, P):
        J = PETScMatrix(J)
        x.copy(self.u.vector().vec())
        self.u.vector().apply("")
        assemble(self.a, tensor=J)
        for bc in self.bcs:
            bc.apply(J)

The ‘apply’ lines are required to update ghost dofs in parallel… you can actually remove that if your code will only be run in serial.

Best!

4 Likes

Thanks! Your answer solves my problem when using other Quasi-Newton methods.