Conjugate-Gradient Method in 2 Dimensions¶
AMath 585, Winter Quarter 2020 at the University of Washington. Developed by R.J. LeVeque and distributed under the BSD license. You are free to modify and use as you please, with attribution.
These notebooks are all available on Github.
This notebook illustrates an implmentation of the Conjugate-Gradient method on a two-dimensional Poisson problem.
We solve the Poisson problem $u_{xx} + u_{yy} = f(x,y)$ with Dirichlet boundary conditions and choose a problem where the truncation error is zero so that the exact solution of the linear system is the solution of the PDE evaluated at the grid points.
Note that this notebook is set up to solve a particular problem on an mx by my grid with things hardwired. If you want to change f (and the true solution) and/or the grid resolution or domain, make sure to re-execute all cells below where these are defined.
%matplotlib inline
from pylab import *
Animating the iteration¶
To use widgets, set use_widgets = True. For javascript animation, set it to False.
use_widgets = False
if use_widgets:
from ipywidgets import interact
import ipywidgets as widgets
def animate_figs(figs):
show_frame = lambda frameno: display(figs[frameno])
interact(show_frame, frameno=widgets.IntSlider(min=0,max=len(figs)-1, value=0))
else:
from jsanimate_figs import animate_figs
Define the matrix-vector multiply function¶
To implement the C-G algorithm we need to be able to multiply $Av$ for various vectors $v$. The following matvec routine computes and returns this product. This is the only place the matrix $A$ is needed, so not that in general we do not need to form or store the large sparse matrix.
Note that in the 2D case we keep the vector of unknowns as a 2D grid function U rather than reshaping it into a long vector as we did in the notebook LaplacianMatrix.html. Not only is this simpler to manage than reshaping grid functions back and forth into vectors, it is also easier to write the matrix-vector multiply function by keeping U as a grid function and applying the stencil based on neighboring points on the grid rather than having to determine where each neighbor winds up in the reshaped vector. The number of equations in the linear system is the same as the number of unknowns and we can associate each equation with a grid point by applying the finite-difference method at this point.
The matvec function below applies to 5-point Laplacian to any grid function V defined on the mx by my grid of interior points.
def matvec(V):
"""
Given a grid function V on an mx by my grid.
Return b = A*v
"""
# pad V with zeros around border, needed for computing centered differences:
V_full = zeros((mx+2, my+2))
V_full[1:-1, 1:-1] = V
# set array B of shape (mx,my) at interior points (same shape as V)
# Note that slicing with [1:-1] corresponds to interior points,
# [:-2] correspond to points to the left (or below) and
# [2:] are points to the right (or above)
B = (V_full[:-2,1:-1] - 2*V_full[1:-1,1:-1] + V_full[2:,1:-1]) / dx**2 \
+(V_full[1:-1,:-2] - 2*V_full[1:-1,1:-1] + V_full[1:-1,2:]) / dy**2
return B
Inner-product function¶
The C-G algorithm requires taking inner products of two vectors in various places. In the 1D case the inner product of vectors v and w can be computed using dot(v,w), which we used in ConjugateGradient.html.
This won't work in 2D since we want to leave V and W as 2D grid functions rather than reshaping them into 1D vectors, and so dot(V,W) would mean the matrix product $VW$ which would not even be defined if mx is not equal to my and in any case would be the wrong thing.
Instead the inner product can be defined by sum(V*W). Recall that in Python V*W means the element-wise product of the numpy arrays V and W, so it is an array of the same shape as V and W (in Matlab this would be V.*W). Then sum() computes the sum of all elements in an array and returns a scalar. This is the proper inner product in 2D and this would also work in 1D instead of dot (and works in 3D if V and W are 3D numpy arrays). Note also that numpy.inner does not return the right thing the way we are using these arrays.
def inner_product(V,W):
"""
Given two grid functions V and W on an mx by my grid.
Return their inner product when viewed as vectors of unknowns
"""
return sum(V*W)
Define the domain, grid, and true solution¶
ax = 0.
bx = 2.
mx = 19
dx = (bx-ax)/(mx+1.)
x_full = linspace(0,2,mx+2)
ay = 0.
by = 1.
my = 9
dy = (by-ay)/(my+1.)
y_full = linspace(0,1,my+2)
X_full, Y_full = meshgrid(x_full, y_full, indexing='ij')
utrue_fcn = lambda X,Y: X**2 + Y**2 - 3.
f_fcn = lambda X,Y: 4*ones(X.shape)
Utrue_full = utrue_fcn(X_full, Y_full)
Utrue = Utrue_full[1:-1, 1:-1]
A function to plot the approximate and true solution¶
See Grids2D.html for more on plotting grid functions. Note that this function uses some variables that are set above, e.g. Utrue_full used for setting the boundary values when extending U to U_full for plotting.
def make_plot(U,k,rnorm,enorm):
levels = linspace(-4,4,17)
cmap = get_cmap('bwr') # blue-white-red
vmin = -4; vmax = 4 # symmetric so white is in center
fig = figure(figsize=(10,6))
subplot(2,1,1)
U_full = Utrue_full.copy()
U_full[1:-1, 1:-1] = U
contourf(X_full,Y_full,U_full,levels,
cmap=cmap,vmin=vmin,vmax=vmax)
axis('scaled')
colorbar(shrink=0.8)
contour(X_full,Y_full,Utrue_full,levels,colors='k')
title('iteration %3i: 2-norm(r) = %.2e, max-norm(E) = %.2e\n' \
% (k,rnorm,enorm) \
+ 'Approximate solution and contours of true solution')
subplot(2,1,2)
contourf(X_full,Y_full,Utrue_full,levels,
cmap=cmap,vmin=vmin,vmax=vmax)
axis('scaled')
colorbar(shrink=0.8)
contour(X_full,Y_full,Utrue_full,levels,colors='k')
title('True solution')
tight_layout() # make room for titles on subplots
close(fig)
return fig
The Conjugate-Gradient algorithm¶
The next cell implements the C-G algorithm.
maxiter = 20
kplot = 1
verbose = True
tol = 1e-14 # stop if the residual falls below tol
print('Solving on %i by %i grid with %i unknowns' \
% (mx,my,mx*my))
print(' using tol = %.2e with maxiter = %i' % (tol,maxiter))
# right-hand side:
F_full = f_fcn(X_full, Y_full)
F = F_full[1:-1, 1:-1] # at interior points
# adjust for Dirichlet BCs:
F[0,:] = F[0,:] - Utrue_full[0,1:-1] / dx**2
F[-1,:] = F[-1,:] - Utrue_full[-1,1:-1] / dx**2
F[:,0] = F[:,0] - Utrue_full[1:-1,0] / dy**2
F[:,-1] = F[:,-1] - Utrue_full[1:-1,-1] / dy**2
# initial guess:
U0_full = 0*X_full
U0 = U0_full[1:-1, 1:-1] # interior points
U = U0.copy() # current iterate
r = F - matvec(U) # initial residual
p = r.copy() # initial direction
k = 0
enorm = abs(U-Utrue).max()
errors = [enorm]
rTr_k = inner_product(r,r)
rnorm = sqrt(rTr_k)
if verbose:
print('iteration %3i: 2-norm(r) = %.2e, max-norm(E) = %.2e' \
% (k,rnorm,enorm))
figs = [] # for the list of figures we generate
fig = make_plot(U,k,rnorm,enorm)
figs.append(fig)
rTr_km = rTr_k # in general will hold r^T * r at iteration k-1
for k in range(1,maxiter+1):
w = matvec(p) # the only matrix-vector multiply
a = rTr_km / inner_product(p,w) # alpha_{k-1} in CG algorithm
U = U + a*p
r = r - a*w
enorm = abs(U-Utrue).max()
errors.append(enorm)
rTr_k = inner_product(r,r)
rnorm = sqrt(rTr_k)
if mod(k,kplot)==0 or k==maxiter:
# every kplot iterations create a plot:
fig = make_plot(U,k,rnorm,enorm)
figs.append(fig)
if verbose:
print('iteration %3i: 2-norm(r) = %.2e, max-norm(E) = %.2e' \
% (k,rnorm,enorm))
# check for convergence:
if rnorm < tol:
print('Satisfied rnorm < tol after %i iterations' % k)
break
# determine next search direction:
b = rTr_k / rTr_km # beta_{k-1} in CG algorithm
rTr_km = rTr_k # for next iteration
p = r + b*p # next search direction
print('Stopped after %i iterations with 2-norm(r) = %.2e' \
% (k,rnorm))
semilogy(range(1,len(errors)+1), errors, 'b-x')
grid(True)
animate_figs(figs)
Try the following:¶
- Take more iterations: set e.g.
maxiter = 200andkplot = 10. You should find that it converges to machine precision in about 70 iterations. - Refine the grid, e.g. by a factor of 2 in each direction by setting
mx = 39, my = 19. You should find it converges to machine precision in about 149 iterations. Note that in this case there are 741 unknowns, about 4 times as many as whenmx = 19, my = 9. This pattern continues on finer grids.
Recall that the convergence rate can be estimated in terms of the condition number of the matrix, and we expect the number of iterations required to grow like $k = O\left(\sqrt{\kappa(A)}\right)$ from the discussion following (4.60) in the text.
Here $\kappa(A) = O(1/h^2)$ if $h = \Delta x = \Delta y$, so we expect $k$ to grow like $O(1/h)$, as observed. In three dimensions the condition number is roughly the same, so roughly twice as many iterations are required if we cut $h = \Delta x = \Delta y = \Delta z$ in half, even though now this increases the size of the system by a factor of 8.
Note also that requiring such a small tolerance is not reasonable for real problems, and the number of iterations required to reach an error consistent with the global error for most problems would be much less. (Here we chose a problem where the global error is 0.)