TRON solver alternatives

Hi all,

I’m currently working with a TAO tron solver to solve some viscosity problems. It’s working fairly well but I’m encountering some hiccups here and there so I’d like to switch to using a TAO BNK solver, like BNTR. However, I’m finding it difficult to switch between the two. I’ve provided a MWE for 2D linear elasticity uniaxial tension below. The TAO problem is defined very similar to the one here: [Problem with parallel computation and dolfinx (TAO solver)].

# Import required modules
import gmsh
import numpy as np
import ufl
import dolfinx
from dolfinx import fem, io, mesh, plot, nls, log, la
from dolfinx.fem.petsc import (apply_lifting, assemble_matrix, assemble_vector,
                               create_matrix, create_vector, set_bc)
from mpi4py import MPI
from petsc4py.PETSc import ScalarType
import petsc4py
from petsc4py import PETSc
import meshio
import as gmshio

class TAO_Problem:

    """Nonlinear problem class compatible with PETSC.TAO solver"""

    def __init__(self, f, F, J, u, bcs):

        self.bcs = bcs                 # bcs = boundary conditions
        self.u   = u                   # u = solution 
        self.L   = dolfinx.fem.form(F) # F = residual
        self.a   = dolfinx.fem.form(J) # J = jacobian
        self.obj = dolfinx.fem.form(f) # f = objective function

        # .. Generation of matrices in the no-linear problem
        self.A = dolfinx.fem.petsc.create_matrix(self.a)
        self.b = dolfinx.fem.petsc.create_vector(self.L)

    def f(self, tao, x):

        # .... Connection between ranks in the vector x: FORWARD
        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        # ..... We copy the vector x in the vector u
        # .... Connection between ranks in the vector u in the class
        self.u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)

        return dolfinx.fem.assemble_scalar(dolfinx.fem.form(self.obj))

    def F(self, tao, x, F):

        # .... Connection between ranks in the vector x: FORWARD
        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        # ..... We copy the vector x in the vector u
        # .... Connection between ranks in the vector u in the class
        self.u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)

        with F.localForm() as f_local:

        dolfinx.fem.petsc.assemble_vector(F, self.L)
        # .... Include the boundary conditions using the lifting technique
        dolfinx.fem.petsc.apply_lifting(F, [self.a], bcs=[self.bcs], x0=[x], scale=-1.0)
        # .... Connection between ranks in the vector x: REVERSE
        F.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
        # .... Redefine the function F with the boundary condition.
        dolfinx.fem.petsc.set_bc(F, self.bcs, x, -1.0)

    def J(self, tao, x, A, P):

        dolfinx.fem.petsc.assemble_matrix(A, self.a, self.bcs)

gmsh_model_rank = 0
mesh_comm = MPI.COMM_WORLD
gdim = 2

name = "BNTRtest"
msh = mesh.create_unit_square(MPI.COMM_WORLD, 1, 1)

# Create ufl cells
cell = ufl.Cell("triangle", gdim)

# Create elements for displacement and pressure calculation
V_el = ufl.VectorElement("CG", cell, 1)
u_space = fem.FunctionSpace(msh,V_el)
u = fem.Function(u_space)

ux_space, ux_submap = u_space.sub(0).collapse()
uy_space, uy_submap = u_space.sub(1).collapse()

# Identify the boundaries
def boundary_left(x):
    return np.isclose(x[0], 0)
def boundary_right(x):
    return np.isclose(x[0], 1)
def boundary_bottom(x):
    return np.isclose(x[1], 0)

# Identify the dofs
dofs_left_x = fem.locate_dofs_geometrical((u_space.sub(0), ux_space), boundary_left)
dofs_right_x = fem.locate_dofs_geometrical((u_space.sub(0), ux_space), boundary_right)
dofs_bottom_y = fem.locate_dofs_geometrical((u_space.sub(1), uy_space), boundary_bottom)

#Define the displacement boundary interpolation expression
class Expression_x:
    def __init__(self):
        self.m = 0.0
    def eval(self, x):
        return np.full(x.shape[1], self.m*x[0])
class Expression_y:
    def __init__(self):
        self.m = 0.0
    def eval(self, x):
        return np.full(x.shape[1], self.m*x[1])

no_disp = 0.0
u_right_x = fem.Function(ux_space)
u_right_x_expr = Expression_x()
u_right_x_expr.m = no_disp

u_left_x = fem.Function(ux_space)
u_left_x_expr = Expression_x()
u_left_x_expr.m = no_disp

u_bottom_y = fem.Function(uy_space)
u_bottom_y_expr = Expression_y()
u_bottom_y_expr.m = no_disp

# Set Dirichlet BC
bc_left_x = fem.dirichletbc(u_left_x, dofs_left_x, V=u_space.sub(0))
bc_bottom_y = fem.dirichletbc(u_bottom_y, dofs_bottom_y, V=u_space.sub(1))
bc_right_x = fem.dirichletbc(u_right_x, dofs_right_x, V=u_space.sub(0))
bcs = [bc_left_x, bc_bottom_y,bc_right_x]

nu = fem.Constant(msh, ScalarType(0.3))  # Poisson ratio
E1 = fem.Constant(msh, ScalarType(200e9))  # Young's Modulus [GPa]
la_E = nu*1/(1+nu)/(1-nu)
mu_E = 1/2/(1+nu)    

#Loading Parameters
load = np.linspace(0.0, 0.01, num=5)

# Kinematics
def eps(u):
    return ufl.sym(ufl.grad(u))

def dotC(e):
    return la_E**ufl.Identity(gdim) + 2*mu_E*e

def sigma(E):
    return E1*dotC(E)

def strain_energy(E):
    return 0.5*(E1*ufl.inner(E, dotC(E)))

E = eps(u)
S = sigma(E)
metadata = {"quadrature_degree": 4}
dx = ufl.Measure("dx", domain=msh, metadata=metadata)

J_p = strain_energy(E)*dx # Functional
R = ufl.derivative(J_p, u) # Jacobian
dR = ufl.derivative(R, u) # Hessian

# Define the nonlinear problem
problem = TAO_Problem(J_p, R, dR, u, bcs)
b = u_space.dofmap.index_map, u_space.dofmap.index_map_bs)
J = dolfinx.fem.petsc.create_matrix(problem.a)

# Create PETSc TAO
solver_tao = petsc4py.PETSc.TAO().create(comm=mesh_comm)
solver_tao.setGradient(problem.F, b)
solver_tao.setHessian(problem.J, J)

# Create lower and upper bound functions

def V_upper(x): # vector upper bound
    return np.full((gdim,x.shape[1]),1e80)
def V_lower(x): # vector lower bound
    return np.full((gdim,x.shape[1]),-1e80)

lb = fem.Function(u_space)
ub = fem.Function(u_space)

ub.interpolate(V_upper) # applying displacement upper bound
lb.interpolate(V_lower) # applying displacement lower bound

solver_tao.setVariableBounds(lb.vector, ub.vector)
results = {'E': [], 'S': []}

# Process load steps
nx = ufl.as_vector([1, 0])
ny = ufl.as_vector([0, 1])
V =fem.assemble_scalar(fem.form(ScalarType(1.0)*dx))

for step, factor in enumerate(load):
    u_right_x_expr.m = factor
    print(f'Step {step+1} of {len(load)}')
    # Solve the problem
    print(f'elastic potential = {fem.assemble_scalar(fem.form(strain_energy(E)*dx))}')
import matplotlib.pyplot

E = np.array(results['E'])

S = np.array(results['S'])

fig, ax1 = matplotlib.pyplot.subplots()
ax1.set_title("Elasticity 2D", fontsize=12)
ax1.set_ylabel(r'volume-averaged stress $\sigma_{xx}$ [Pa]', fontsize=12)

ax1.plot(S, color='tab:blue', linestyle='-', linewidth=1.0, markersize=4.0, marker='.', label=r'calculated')

ax1.legend(loc='lower right')

When the solver type is TRON [ solver_tao.setType("tron") ], the solver works fine, the strain energy is correct at each step and the solution matches theory. However, switching to solver_tao.setType("bntr") the solver converges on zero displacement and zero energy, which is clearly incorrect and doesn’t even abide by the boundary conditions. I can’t find any working examples of BNTR (or similar) solvers on here or the PETSc TAO website, but I’d very much like to get one working, for a couple of reasons. TRON solvers “will be deprecated in the next version [of PETSc TAO] in favor of the Bounded Newton Trust Region (BNTR) algorithm.” and libraries like SciPy use BNTR, so it’s easier to test problems there when I can use the same type of solver.

Kind regards,