# 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 dolfinx.io.gmshio 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
# ..... We copy the vector x in the vector u
x.copy(self.u.vector)
# .... Connection between ranks in the vector u in the class

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

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

# .... Connection between ranks in the vector x: FORWARD
# ..... We copy the vector x in the vector u
x.copy(self.u.vector)
# .... Connection between ranks in the vector u in the class

with F.localForm() as f_local:
f_local.set(0.0)

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
# .... 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):

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

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_right_x.interpolate(u_right_x_expr.eval)

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

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

# 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)

# Kinematics
def eps(u):

def dotC(e):
return la_E*ufl.tr(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)

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 = dolfinx.la.create_petsc_vector( 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.setType("tron")
solver_tao.setObjective(problem.f)
solver_tao.setHessian(problem.J, J)
solver_tao.setMaximumIterations(20)
solver_tao.getKSP().setType("nash")
solver_tao.getKSP().getPC().setType("lu")

# 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': []}

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

u_right_x_expr.m = factor
u_right_x.interpolate(u_right_x_expr.eval)

# Solve the problem
solver_tao.solve(u.vector)
print(f'elastic potential = {fem.assemble_scalar(fem.form(strain_energy(E)*dx))}')

results['E'].append(fem.assemble_scalar(fem.form(ufl.dot(E*nx,nx)*dx))/V)
results['S'].append(fem.assemble_scalar(fem.form(ufl.dot(S*nx,nx)*dx))/V)

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.grid(linewidth=0.25)
fig.tight_layout()

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,

Magnus