Issue in using Petsc Tao solver for a bounded field problem

Hello Everyone,

I am trying to use petsc tao solver to obtain the solution of a bounded field problem. I am following the approach as suggested by

My code is as below

from dolfinx import fem, plot, mesh, cpp , la           # dolfinx library
from mpi4py import MPI
import matplotlib.pyplot as plt    # Plotting Library
import numpy as np
from ufl import replace, TestFunction, TrialFunction, nabla_div, Measure, grad, inner, dot, derivative, lhs, rhs 

import dolfinx

import ufl
import pyvista
from mpi4py import MPI
from petsc4py import PETSc
import sys, os, shutil, math
import numpy as np
from math import pi
from dolfinx.mesh import locate_entities, meshtags
from dolfinx.fem import FunctionSpace, Function, locate_dofs_geometrical, dirichletbc, Constant
from dolfinx.fem.petsc import LinearProblem
from petsc4py.PETSc import ScalarType
import petsc4py
from tao_problem import TAO_Problem

length = 200. 
num_elem = 400
mesh = mesh.create_interval(MPI.COMM_WORLD,num_elem,[0., length])

V_u = FunctionSpace(mesh, ('Lagrange', 1))               # for displacement field)
V_alpha = FunctionSpace(mesh, ('Lagrange', 1) )             # for phase field

# Define the function, test and trial fields
u, du, v = Function(V_u), TrialFunction(V_u), TestFunction(V_u)                             # Functions, Trial Functions and Test functions for displacement field
alpha, dalpha, beta = Function(V_alpha), TrialFunction(V_alpha), TestFunction(V_alpha)      # Functions, Trial Functions and Test functions for phase field


lb = Function(V_alpha)
ub = Function(V_alpha)
alpha_0 = Function(V_alpha)
u_R = Function(V_u)
disp = Function(V_u)
dx = Measure("dx",domain=mesh)
ds = Measure("ds",domain=mesh)

# Identify the boundaries to apply boundary conditions
def on_boundary_L(x):	
    return np.isclose(x[0],0.0)

dof_left_u = locate_dofs_geometrical(V_u, on_boundary_L)
bcl = dirichletbc(ScalarType(0), dof_left_u, V_u)

# Define Material Properties
E = Constant(mesh,30000.0) 
nu = Constant(mesh,0.2)  

# Define the data for phase field
ell = Constant(mesh,5.)
kapa = 1e-7
Gc = Constant(mesh,0.12) 
c_w = 2.0 

# define functions
def w(alpha):
    return alpha**2


def a(alpha):
    k_ell = Constant(mesh,1.e-6)
    return (1-alpha)**2 +k_ell

def eps(u):
    return grad(u)

def sigma_0(u):
    mu    = E / (2.0*(1.0 + nu))
    lmbda = (E * nu) / ((1.0 - 2.0*nu)*(1.0 + nu))
    return E*(eps(u))

def sigma(u,alpha):
    return (a(alpha))*sigma_0(u)

# Apply the limit on the phase field
zero_alpha = Function(V_alpha)
with zero_alpha.vector.localForm() as bc_local:
    bc_local.set(0.0)

one_alpha = Function(V_alpha)
with one_alpha.vector.localForm() as bc_local:
    bc_local.set(1.0)

ub.interpolate(one_alpha)
alpha_0.interpolate(zero_alpha)

# Governing Equations
f_ext = Constant(mesh,0.0)
external_work = dot(f_ext, u) * dx 
elastic_energy = 0.5*inner(sigma(u,alpha), eps(u))*dx 
dissipated_energy = Gc/float(c_w)*(w(alpha)/ell + ell*dot(grad(alpha), grad(alpha)))*dx
total_energy = elastic_energy + dissipated_energy - external_work 

# First directional derivative wrt displacement field
E_u = derivative(total_energy,u,v)

# First and second directional derivative wrt phase field
E_alpha = derivative(total_energy,alpha,beta)
E_alpha_alpha = derivative(E_alpha,alpha,dalpha)


disp_a = Constant(mesh,0.0012)
    
def on_boundary(x):		
    return np.isclose(x[0],length)

dof_right_u = locate_dofs_geometrical(V_u, on_boundary)

bcr = dirichletbc(disp_a, dof_right_u, V_u)
bc_disp = [bcl, bcr]

# Define solver parameters for displacement field
E_du = replace(E_u,{u:du})

problem_u = LinearProblem(a = lhs(E_du), L = rhs(E_du), bcs = bc_disp,u = u, petsc_options={"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "umfpack"})
solver_u = problem_u.solve()

bc_alpha_L = dirichletbc(Constant(mesh,0.0), dof_left_u, V_alpha)
bc_alpha_R = dirichletbc(Constant(mesh,0.0), dof_right_u, V_alpha)
bcs_alpha = [bc_alpha_L,bc_alpha_R]
dolfinx.fem.set_bc(ub.vector, bcs_alpha)

damage_problem = TAO_Problem(E_alpha, E_alpha_alpha, total_energy, alpha, bcs_alpha)

b = la.create_petsc_vector(V_alpha.dofmap.index_map, V_alpha.dofmap.index_map_bs)
J = dolfinx.fem.petsc.create_matrix(damage_problem.a)
T_E = la.create_petsc_vector(V_alpha.dofmap.index_map, V_alpha.dofmap.index_map_bs)

solver_alpha_tao = PETSc.TAO().create()
solver_alpha_tao.setType("tron")
solver_alpha_tao.setFromOptions()
solver_alpha_tao.setObjectiveGradient(damage_problem.F,b)
solver_alpha_tao.setObjective(damage_problem.f,T_E)
solver_alpha_tao.setGradient(damage_problem.J,J)
solver_alpha_tao.setTolerances(gatol=1.0e-5,grtol=1.0e-5)
solver_alpha_tao.getKSP().setType("preonly")
solver_alpha_tao.getKSP().setTolerances(rtol=1.0e-5)
solver_alpha_tao.getKSP().getPC().setType("lu")

solver_alpha_tao.setVariableBounds(lb.vector,ub.vector)

solver_alpha_tao.solve(alpha.vector)

Module tao_problem

from typing import List

from petsc4py import PETSc

import dolfinx
import ufl
        
class TAO_Problem:
    """Nonlinear problem class compatible with PETSC.SNES solver.
    """

    def __init__(self, F: ufl.form.Form, J: ufl.form.Form, f: ufl.form.Form, alpha_u: dolfinx.fem.Function, bcs: List[dolfinx.fem.dirichletbc]):
        """This class set up structures for solving a non-linear problem using Newton's method.

        Parameters
        ==========
        F: Residual.
        J: Jacobian.
        alpha_u: Solution.
        bcs: Dirichlet boundary conditions.
        """
        F, J = dolfinx.fem.form(F), dolfinx.fem.form(J)
        f = dolfinx.fem.form(f)
        #
        self.L = F
        self.a = J
        self.K = f
        self.bcs = bcs
        self.alpha_u = alpha_u
        
        # Create matrix and vector to be used for assembly
        # of the non-linear problem
        self.A = dolfinx.fem.petsc.create_matrix(self.a)
        self.b = dolfinx.fem.petsc.create_vector(self.L)
        self.ob = self.K


    def F(self, x: PETSc.Vec, b: PETSc.Vec):
        """Assemble the gradiant vector.

        Parameters
        ==========
        x: Vector containing the latest solution.
        A: vector to assemble the gradiant into.
        """
        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        x.copy(self.alpha_u.vector)
        self.alpha_u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        
        # Zero the residual vector
        with b.localForm() as b_local:
            b_local.set(0.0)
        dolfinx.fem.petsc.assemble_vector(b, self.L)
        
        # Apply boundary conditions
        dolfinx.fem.petsc.apply_lifting(b, [self.a], [self.bcs], [x])
        b.ghostUpdate(addv=PETSc.InsertMode.ADD,mode=PETSc.ScatterMode.REVERSE)
        dolfinx.fem.petsc.set_bc(b, self.bcs, x)

                
    def J(self, x: PETSc.Vec, A: PETSc.Mat):
        """Assemble the Jacobian matrix.

        Parameters
        ==========
        x: Vector containing the latest solution.
        A: Matrix to assemble the Jacobian into.
        """
        A.zeroEntries()
        dolfinx.fem.petsc.assemble_matrix(A, self.a, self.bcs)
        A.assemble()


    def f(self, x: PETSc.Vec, ob: PETSc.Vec):
        """Assemble the objective vector.

        Parameters
        ==========
        x: Vector containing the latest solution.
        ob: vector to assemble the objective.
        """
        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        x.copy(self.alpha_u.vector)
        self.alpha_u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        
        # Zero the residual vector
        with b.localForm() as b_local:
            b_local.set(0.0)
        dolfinx.fem.petsc.assemble_vector(ob, self.K)
        
        # Apply boundary conditions
        dolfinx.fem.petsc.apply_lifting(ob, [self.a], [self.bcs], [x])
        b.ghostUpdate(addv=PETSc.InsertMode.ADD,mode=PETSc.ScatterMode.REVERSE)
        dolfinx.fem.petsc.set_bc(ob, self.bcs, x)

When I run this code, it keep on running and at last get killed without showing any error. I am not able understand why this is happening.

Can someone please help me to find the reason of this issue and how I can correct it.

Thank you
Manish