Creating a custom PETScSNESSolver

Linear Algebra
5

Hi,

Thanks for your reply.

Currently within the NewtonSolver code the nullspace is attached in the solver block, but I’m not sure how to insert the nullspace within the snes solver.

NewtonSolver syntax:

class CustomSolver(NewtonSolver):
    def __init__(self, mesh, nullspace):
        NewtonSolver.__init__(self, mesh.mpi_comm(),
                              PETScKrylovSolver(), PETScFactory.instance())
        self.nullspace = nullspace
        
    def solver_setup(self, A, P, problem, iteration):
        as_backend_type(A).set_nullspace(self.nullspace) # What is the equivalent for the direct SNES code?
        self.linear_solver().set_operator(A)
        self.linear_solver().set_from_options()

Looking at the PETSc documentation I think snes.setDM sets the nullspace, is this correct? If so can you let me know how to convert from dolfin.cpp.la.VectorSpaceBasis (which is how nullspaces are typically constructed in fences) to petsc4py.PETSc.DM as required by the function. I’ve copied the code from Nate’s example and added a nullspace constructor below.

Many thanks,
K

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Apr 25 07:12:48 2023
"""

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)

def create_nullspace(w, W):
    nullspace_basis = [w.vector().copy() for i in range(3)]
    W.sub(0).dofmap().set(nullspace_basis[0], 1.0)
    W.sub(1).dofmap().set(nullspace_basis[1], 1.0)
    W.sub(2).dofmap().set(nullspace_basis[2], 1.0);
    for x in nullspace_basis:
        x.apply("insert")

    nullspace = VectorSpaceBasis(nullspace_basis)
    nullspace.orthonormalize()
    return nullspace

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)
        
u  = Function(V)  
nullspace = create_nullspace(u, V)
print (u.vector()[:])
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.setDM(nullspace) # TypeError: Argument 'dm' has incorrect type (expected petsc4py.PETSc.DM, got dolfin.cpp.la.VectorSpaceBasis)
snes.setJacobian(problem.J, J_mat.mat())
snes.solve(None, problem.u.vector().vec())

print (u.vector()[:])

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