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
x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.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
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
x.copy(self.u.vector)
# .... 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:
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
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):
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)
#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.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)
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 = 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.setGradient(problem.F, b)
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': []}
# 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)}')
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