SNES solver fails when using the line search in FEniCSx

Hi all, I found SNES solver is even worse when using the line search in FEniCSx.

I can reproduce the strange behavior by slightly modifying the FEniCSx source code at dolfinx/test_newton.py at main · FEniCS/dolfinx (github.com).

Here is the MWE:

"""Unit tests for Newton solver assembly"""

import numpy as np

import ufl
from dolfinx import cpp as _cpp
from dolfinx import fem, la
from dolfinx.fem import (Function, FunctionSpace, apply_lifting,
                         assemble_matrix, assemble_vector, create_matrix,
                         create_vector, DirichletBC, Form,
                         locate_dofs_geometrical, set_bc)
from dolfinx.generation import UnitSquareMesh
from ufl import TestFunction, TrialFunction, derivative, dx, grad, inner

from mpi4py import MPI
from petsc4py import PETSc

class NonlinearPDE_SNESProblem:
    def __init__(self, F, u, bc):
        V = u.function_space
        du = TrialFunction(V)
        self.L = Form(F)
        self.a = Form(derivative(F, u, du))
        self.bc = bc
        self._F, self._J = None, None
        self.u = u

    def F(self, snes, x, F):
        """Assemble residual vector."""
        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        x.copy(self.u.vector)
        self.u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)

        with F.localForm() as f_local:
            f_local.set(0.0)
        assemble_vector(F, self.L)
        apply_lifting(F, [self.a], bcs=[[self.bc]], x0=[x], scale=-1.0)
        F.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
        set_bc(F, [self.bc], x, -1.0)

    def J(self, snes, x, J, P):
        """Assemble Jacobian matrix."""
        J.zeroEntries()
        assemble_matrix(J, self.a, bcs=[self.bc])
        J.assemble()

def test_nonlinear_pde_snes():
    """Test Newton solver for a simple nonlinear PDE"""
    # Create mesh and function space
    mesh = UnitSquareMesh(MPI.COMM_WORLD, 12, 15)
    V = FunctionSpace(mesh, ("Lagrange", 1))
    u = Function(V)
    v = TestFunction(V)
    F = inner(5.0, v) * dx - ufl.sqrt(u * u) * inner(
        grad(u), grad(v)) * dx - inner(u, v) * dx

    u_bc = Function(V)
    u_bc.x.array[:] = 1e3
    bc = DirichletBC(u_bc, locate_dofs_geometrical(V, lambda x: np.logical_or(np.isclose(x[0], 0.0),
                                                                              np.isclose(x[0], 1.0))))

    # Create nonlinear problem
    problem = NonlinearPDE_SNESProblem(F, u, bc)

    u.x.array[:] = 0.9
    b = la.create_petsc_vector(V.dofmap.index_map, V.dofmap.index_map_bs)
    J = fem.create_matrix(problem.a)

    # Create Newton solver and solve
    snes = PETSc.SNES().create()
    opts = PETSc.Options() # read the prefix of solver
    opts['snes_linesearch_type'] = 'basic'
    # opts['snes_linesearch_type'] = 'bt'
    opts['snes_monitor'] = None
    opts['snes_linesearch_monitor'] = None
    snes.setFromOptions() # make above settings effective

    snes.setFunction(problem.F, b)
    snes.setJacobian(problem.J, J)

    snes.setTolerances(rtol=1.0e-9, max_it=10)
    snes.getKSP().setType("preonly")
    snes.getKSP().setTolerances(rtol=1.0e-9)
    snes.getKSP().getPC().setType("lu")

    snes.solve(None, u.vector)

test_nonlinear_pde_snes()

If we run the code without the line search (opts['snes_linesearch_type'] = 'basic'), everything works normally:

  0 SNES Function norm 6.849186430212e+03 
  1 SNES Function norm 6.976779625428e+04 
  2 SNES Function norm 7.515891730709e+03 
  3 SNES Function norm 4.767650881912e+01 
  4 SNES Function norm 1.358290774702e-03 
  5 SNES Function norm 1.706659139551e-09 

However, it blows up when using the line search (opts['snes_linesearch_type'] = 'bt'):

  0 SNES Function norm 6.849186430212e+03 
      Line search: gnorm after quadratic fit 4.361060415642e+05
      Line search: Cubic step no good, shrinking lambda, current gnorm 4.869482320218e+05 lambda=1.0000000000000002e-02
      Line search: Cubic step no good, shrinking lambda, current gnorm 4.920841347160e+05 lambda=1.0000000000000002e-03
      Line search: Cubic step no good, shrinking lambda, current gnorm 4.925982367782e+05 lambda=1.0000000000000003e-04
      Line search: Cubic step no good, shrinking lambda, current gnorm 4.926496520979e+05 lambda=1.0000000000000004e-05
      Line search: Cubic step no good, shrinking lambda, current gnorm 4.926547936810e+05 lambda=1.0000000000000004e-06
      Line search: Cubic step no good, shrinking lambda, current gnorm 4.926553078398e+05 lambda=1.0000000000000005e-07
      Line search: Cubic step no good, shrinking lambda, current gnorm 4.926553592557e+05 lambda=1.0000000000000005e-08
      Line search: Cubic step no good, shrinking lambda, current gnorm 4.926553643973e+05 lambda=1.0000000000000005e-09
      Line search: Cubic step no good, shrinking lambda, current gnorm 4.926553649114e+05 lambda=1.0000000000000006e-10
      Line search: Cubic step no good, shrinking lambda, current gnorm 4.926553649629e+05 lambda=1.0000000000000006e-11
      Line search: Cubic step no good, shrinking lambda, current gnorm 4.926553649680e+05 lambda=1.0000000000000006e-12
      Line search: Cubic step no good, shrinking lambda, current gnorm 4.926553649685e+05 lambda=1.0000000000000007e-13
      Line search: unable to find good step length! After 12 tries 
      Line search: fnorm=6.8491864302115573e+03, gnorm=4.9265536496850761e+05, ynorm=1.3326628127967717e+04, minlambda=9.9999999999999998e-13, lambda=1.0000000000000007e-13, initial slope=-4.6911354755794138e+07

A similar problem also occurs in FEniCS as reported in the post Using petsc4py.PETSc.SNES directly - #17 by nabarnaf.

Could someone please show me how to debug it? Thanks in advance!

By inspecting the result, it seems the solution variable x or u is not updated correctly when using the line search. I would like to know if there is a way to fix it?

If you consider the petsc documentation: https://petsc.org/release/docs/manualpages/SNES/SNESLINESEARCHBT.html#SNESLINESEARCHBT
It states that

Backtracking line search. This line search finds the minimum of a polynomial fitting of the L2 norm of the function or the objective function if it is provided with SNESSetObjective(). If this fit does not satisfy the conditions for progress, the interval shrinks and the fit is reattempted at most max_it times or until lambda is below minlambda.

This means that you need to use snes.setObjective()

Thanks for your reply, @dokken. I try to set up a function of the norm of residual as follows:

    def func(snes, x):
        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        x.copy(u.vector)
        u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        
        R = inner(5.0, v) * dx - ufl.sqrt(u * u) * inner(
            grad(u), grad(v)) * dx - inner(u, v) * dx
        J = ufl.derivative(R, u)
        Vec = create_vector(R)
        with Vec.localForm() as f_local:
            f_local.set(0.0)
        assemble_vector(Vec, R)
        apply_lifting(Vec, [J], bcs=[[bc]], x0=[x], scale=-1.0)
        Vec.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
        set_bc(Vec, [bc], x, -1.0)
        return Vec.norm()

However, even after I pass it into the solver with snes.setObjective(func), it still does not work. Could you please advise me on how to improve it?

Hi @dokken, I noticed it is also mentioned in the petsc documentation that

If not provided then this defaults to the two norm of the function evaluation (set with SNESSetFunction()).

So even without providing the user-defined objective function, it seems the snes solver should also be able to minimize the function provided by FEniCSx (the residual exactly in this case). But I am still confused why this line search does not work.

You’ve set a boundary condition of u|_{\partial\Omega} = 10^3 but an initial guess of u_0 = 0.9. Try something more reasonable like u_0 = 10^3.

Hi @nate , thanks for your reply! I set a relatively unreasonable initial guess on purpose to check whether the line search is effective. It turns out that, if the initial guess u_0=10^3 is close to the boundary condition u|_{\partial \Omega}=10^3, the step length is just one. That means the problem can converge without the line search.

When I gradually decrease the initial guess as u_0=10, the line search just breaks down, and it is quite difficult to see a search step between (0,1).

Could you please take a quick look which part of my code is incorrect? Here is the MWE:

"""Unit tests for Newton solver assembly"""

import numpy as np

import ufl
from dolfinx import cpp as _cpp
from dolfinx import fem, la
from dolfinx.fem import (Function, FunctionSpace, apply_lifting,
                         assemble_matrix, assemble_vector, create_matrix,
                         create_vector, DirichletBC, Form,
                         locate_dofs_geometrical, set_bc)
from dolfinx.generation import UnitSquareMesh
from ufl import TestFunction, TrialFunction, derivative, dx, grad, inner

from mpi4py import MPI
from petsc4py import PETSc

class NonlinearPDE_SNESProblem:
    def __init__(self, F, u, bc):
        V = u.function_space
        du = TrialFunction(V)
        self.L = Form(F)
        self.a = Form(derivative(F, u, du))
        self.bc = bc
        self._F, self._J = None, None
        self.u = u

    def F(self, snes, x, F):
        """Assemble residual vector."""
        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        x.copy(self.u.vector)
        self.u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)

        with F.localForm() as f_local:
            f_local.set(0.0)
        assemble_vector(F, self.L)
        apply_lifting(F, [self.a], bcs=[[self.bc]], x0=[x], scale=-1.0)
        F.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
        set_bc(F, [self.bc], x, -1.0)

    def J(self, snes, x, J, P):
        """Assemble Jacobian matrix."""
        J.zeroEntries()
        assemble_matrix(J, self.a, bcs=[self.bc])
        J.assemble()

def test_nonlinear_pde_snes():
    """Test Newton solver for a simple nonlinear PDE"""
    # Create mesh and function space
    mesh = UnitSquareMesh(MPI.COMM_WORLD, 12, 15)
    V = FunctionSpace(mesh, ("Lagrange", 1))
    u = Function(V)
    v = TestFunction(V)
    F = inner(5.0, v) * dx - ufl.sqrt(u * u) * inner(
        grad(u), grad(v)) * dx - inner(u, v) * dx

    u_bc = Function(V)
    u_bc.x.array[:] = 1e3
    bc = DirichletBC(u_bc, locate_dofs_geometrical(V, lambda x: np.logical_or(np.isclose(x[0], 0.0),
                                                                              np.isclose(x[0], 1.0))))

    # Create nonlinear problem
    problem = NonlinearPDE_SNESProblem(F, u, bc)

    u.x.array[:] = 1e1
    b = la.create_petsc_vector(V.dofmap.index_map, V.dofmap.index_map_bs)
    J = fem.create_matrix(problem.a)

    # Create Newton solver and solve
    snes = PETSc.SNES().create()
    opts = PETSc.Options() # read the prefix of solver
    # opts['snes_linesearch_type'] = 'basic'
    opts['snes_linesearch_type'] = 'bt'
    opts['snes_monitor'] = None
    opts['snes_linesearch_monitor'] = None
    snes.setFromOptions() # make above settings effective

    snes.setFunction(problem.F, b)
    snes.setJacobian(problem.J, J)

    def func(snes, x):
        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        x.copy(u.vector)
        u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        
        R = inner(5.0, v) * dx - ufl.sqrt(u * u) * inner(
            grad(u), grad(v)) * dx - inner(u, v) * dx
        J = ufl.derivative(R, u)
        Vec = create_vector(R)
        with Vec.localForm() as f_local:
            f_local.set(0.0)
        assemble_vector(Vec, R)
        apply_lifting(Vec, [J], bcs=[[bc]], x0=[x], scale=-1.0)
        Vec.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
        set_bc(Vec, [bc], x, -1.0)
        return Vec.norm()

    snes.setObjective(func)
    snes.setTolerances(rtol=1.0e-9, max_it=10)
    snes.getKSP().setType("preonly")
    snes.getKSP().setTolerances(rtol=1.0e-9)
    snes.getKSP().getPC().setType("lu")

    snes.solve(None, u.vector)

test_nonlinear_pde_snes()

Hi @dokken and @nate, could you please take a quick look about this issue? I investigated this problem for many days but I still cannot find the reason why it cannot converge when using the line search of SNES with the parameters passed from FEniCSx.

It is quite strange that, if we do not use the line search, it can converge normally even with the “unreasonable” initial guess u_0=0.9. But the residual can never go down when using the line search.

The theory underlying methods encouraging convergence of Newton-type iterative methods is complex and interesting. I suggest you speak to your advisor or review the literature.

Thanks for the reply very much @nate. Previously I have already successfully implemented the line search (minimizing the 2-norm of the residual function) by myself by adapting the NewtonSolver in FEniCSx. In that case, the residual can go down to zero and the line search parameter should be equal to one for a well-posed problem.

Later I found we can use the line search by calling the SNES solver directly, and that is the reason why I would like to post this question. And as you mentioned, the line search should ENCOURAGE the convergence, and so it confused me why the code gave me a worse result when using the line search. I am not familiar with the structure of FEniCSx and PETSc, and just want to make sure the coding part is not the main issue.

In case somebody stumbles across this post, I just wanted to share that this was indeed an issue in the code, or specifically the way the PETSc problem class was set up with Dolfinx.

In short, the issue is that the solution vector on the Dolfinx side and the solution vector on the PETSc side share the same memory if the NonlinearPDE_SNESProblem class is set up as above.
The PETSc SNES line search routines however need to have an independent solution vector to test different step sizes while having a fixed reference point for these tests, see detailed description in the comments of this PR.

It can be fixed by passing in another independent vector to snes.solve (most straightforwardly a copy of u), e.g.

problem = NonlinearPDE_SNESProblem(F, u, bc)
...
snes = PETSc.SNES().create()
x = u.x.petsc_vec.copy()
x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
snes.solve(None, x)

while instead

snes.solve(None, u.x.petsc_vec)

causes the line search routines of PETSc to fail.

See details here.

4 Likes