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:
where \(J_F\) is the Jacobian matrix of \(F\).
We can rewrite this equation as \(\delta u_k = u_{k+1} - u_{k}\),
and
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:
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