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