Custom Newton solvers#

Author: Jørgen S. Dokken

Newtons method, as used in the non-linear Poisson problem, is a way of solving a non-linear equation as a sequence of linear equations.

Given a function \(F:\mathbb{R}^M\mapsto \mathbb{R}^M\), we have that \(u_k, u_{k+1}\in \mathbb{R}^M\) is related as:

\[u_{k+1} = u_{k} - J_F(u_k)^{-1} F(u_k)\]

where \(J_F\) is the Jacobian matrix of \(F\).

We can rewrite this equation as \(\delta u_k = u_{k+1} - u_{k}\),

\[ J_F(u_k)\delta u_k = - F(u_k) \]

and

\[ u_{k+1} = u_k + \delta u_k. \]

Problem specification#

We start by importing all packages needed to solve the problem.

import dolfinx
import dolfinx.fem.petsc
import matplotlib.pyplot as plt
import numpy as np
import pyvista
import ufl
from mpi4py import MPI
from petsc4py import PETSc

We will consider the following non-linear problem:

\[ u^2 - 2 u = x^2 + 4x + 3 \text{ in } [0,1] \]

For this problem, we have two solutions, \(u=-x-1\), \(u=x+3\). We define these roots as python functions, and create an appropriate spacing for plotting these soultions.

def root_0(x):
    return 3 + x[0]


def root_1(x):
    return -1 - x[0]


N = 10
roots = [root_0, root_1]
x_spacing = np.linspace(0, 1, N)

We will start with an initial guess for this problem, \(u_0 = 0\). Next, we define the mesh, and the appropriate function space and function uh to hold the approximate solution.

mesh = dolfinx.mesh.create_unit_interval(MPI.COMM_WORLD, N)
V = dolfinx.fem.FunctionSpace(mesh, ("Lagrange", 1))
uh = dolfinx.fem.Function(V)

Definition of residual and Jacobian#

Next, we define the variational form, by multiplying by a test function and integrating over the domain \([0,1]\)

v = ufl.TestFunction(V)
x = ufl.SpatialCoordinate(mesh)
F = uh**2 * v * ufl.dx - 2 * uh * v * ufl.dx - (x[0]**2 + 4 * x[0] + 3) * v * ufl.dx
residual = dolfinx.fem.form(F)

Next, we can define the jacobian \(J_F\), by using ufl.derivative.

J = ufl.derivative(F, uh)
jacobian = dolfinx.fem.form(J)

As we will solve this problem in an iterative fashion, we would like to create the sparse matrix and vector containing the residual only once.

Setup of iteration-independent structures#

A = dolfinx.fem.petsc.create_matrix(jacobian)
L = dolfinx.fem.petsc.create_vector(residual)

Next, we create the linear solver and the vector to hold du.

solver = PETSc.KSP().create(mesh.comm)
solver.setOperators(A)
du = dolfinx.fem.Function(V)

We would like to monitor the evolution of uh for each iteration. Therefore, we get the dof coordinates, and sort them in increasing order.

i = 0
coords = V.tabulate_dof_coordinates()[:, 0]
sort_order = np.argsort(coords)
max_iterations = 25
solutions = np.zeros((max_iterations + 1, len(coords)))
solutions[0] = uh.x.array[sort_order]

We are now ready to solve the linear problem. At each iteration, we reassemble the Jacobian and residual, and use the norm of the magnitude of the update (dx) as a termination criteria.

The Newton iterations#

i = 0
while i < max_iterations:
    # Assemble Jacobian and residual
    with L.localForm() as loc_L:
        loc_L.set(0)
    A.zeroEntries()
    dolfinx.fem.petsc.assemble_matrix(A, jacobian)
    A.assemble()
    dolfinx.fem.petsc.assemble_vector(L, residual)
    L.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)

    # Scale residual by -1
    L.scale(-1)
    L.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)

    # Solve linear problem
    solver.solve(L, du.vector)
    du.x.scatter_forward()
    # Update u_{i+1} = u_i + delta u_i
    uh.x.array[:] += du.x.array
    i += 1

    # Compute norm of update
    correction_norm = du.vector.norm(0)
    print(f"Iteration {i}: Correction norm {correction_norm}")
    if correction_norm < 1e-10:
        break
    solutions[i, :] = uh.x.array[sort_order]
Iteration 1: Correction norm 29.415833333333346
Iteration 2: Correction norm 10.776793575710077
Iteration 3: Correction norm 2.048892531820762
Iteration 4: Correction norm 0.08991946662079764
Iteration 5: Correction norm 0.0002277569702082286
Iteration 6: Correction norm 2.211496656626386e-09
Iteration 7: Correction norm 1.7457052351197382e-15

We now compute the magnitude of the residual.

dolfinx.fem.petsc.assemble_vector(L, residual)
print(f"Final residual {L.norm(0)}")
Final residual 6.379796870252556e-16

Visualization of Newton iterations#

We next look at the evolution of the solutions and the error of the solution when compared to the two exact roots of the problem.

# Plot solution for each of the iterations
fig = plt.figure(figsize=(15, 8))
for j, solution in enumerate(solutions[:i]):
    plt.plot(coords[sort_order], solution, label=f"Iteration {j}")

# Plot each of the roots of the problem, and compare the approximate solution with each of them
args = ("--go",)
for j, root in enumerate(roots):
    u_ex = root(x)
    L2_error = dolfinx.fem.form(ufl.inner(uh - u_ex, uh - u_ex) * ufl.dx)
    global_L2 = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(L2_error), op=MPI.SUM)
    print(f"L2-error (root {j}) {np.sqrt(global_L2)}")

    kwargs = {} if j == 0 else {"label": "u_exact"}
    plt.plot(x_spacing, root(x_spacing.reshape(1, -1)), *args, **kwargs)
plt.grid()
plt.legend()
L2-error (root 0) 5.033222956847167
L2-error (root 1) 1.1102230246251564e-16
<matplotlib.legend.Legend at 0x7fcca8723580>
../_images/800abcde5a4b7e32340963504258e8443da45202503e05d9960d5153428db336.png

Newton’s method with DirichletBC#

In the previous example, we did not consider handling of Dirichlet boundary conditions. For this example, we will consider the non-linear Poisson-problem. We start by defining the mesh, the analytical solution and the forcing term \(f\).

def q(u):
    return 1 + u**2


domain = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10)
x = ufl.SpatialCoordinate(domain)
u_ufl = 1 + x[0] + 2 * x[1]
f = - ufl.div(q(u_ufl) * ufl.grad(u_ufl))


def u_exact(x):
    return eval(str(u_ufl))

Next, we define the boundary condition bc, the residual F and the Jacobian J.

V = dolfinx.fem.FunctionSpace(domain, ("Lagrange", 1))
u_D = dolfinx.fem.Function(V)
u_D.interpolate(u_exact)
fdim = domain.topology.dim - 1
domain.topology.create_connectivity(fdim, fdim + 1)
boundary_facets = dolfinx.mesh.exterior_facet_indices(domain.topology)
bc = dolfinx.fem.dirichletbc(u_D, dolfinx.fem.locate_dofs_topological(V, fdim, boundary_facets))

uh = dolfinx.fem.Function(V)
v = ufl.TestFunction(V)
F = q(uh) * ufl.dot(ufl.grad(uh), ufl.grad(v)) * ufl.dx - f * v * ufl.dx
J = ufl.derivative(F, uh)
residual = dolfinx.fem.form(F)
jacobian = dolfinx.fem.form(J)

Next, we define the matrix A, right hand side vector L and the correction function du

du = dolfinx.fem.Function(V)
A = dolfinx.fem.petsc.create_matrix(jacobian)
L = dolfinx.fem.petsc.create_vector(residual)
solver = PETSc.KSP().create(mesh.comm)
solver.setOperators(A)

Since this problem has strong Dirichlet conditions, we need to apply lifting to the right hand side of our Newton problem. We previously had that we wanted to solve the system:

\[\begin{split} \begin{align} J_F(u_k)\delta u_k &= - F(u_k)\\ u_{k+1} &= u_k + \delta u_k \end{align} \end{split}\]

we want \(u_{k+1}\vert_{bc}= u_D\). However, we do not know if \(u_k\vert_{bc}=u_D\). Therefore, we want to apply the following boundary condition for our correction \(\delta u_k\)

\[ \delta u_k\vert_{bc} = u_D-u_k\vert_{bc} \]

We therefore arrive at the following Newton scheme

i = 0
error = dolfinx.fem.form(ufl.inner(uh - u_ufl, uh - u_ufl) * ufl.dx<