This is a demonstration of a multigrid preconditioned krylov solver in python3. The code and more examples are present on github here. The problem solved is a Poisson equation on a rectangular domain with homogenous dirichlet boundary conditions. Finite difference with cell-centered discretization is used to get a second order accurate solution, that is further improved to 4th order using deferred correction.
The first step is a multigrid algorithm. This is the simplest 2D geometric multigrid solver.
We need some terminology before going further.
This is a geometric multigrid algorithm, where a series of nested grids are used. There are four parts to a multigrid algorithm
We will define each of these in sequence. These operators act of different quantities that are stored at the cell center. We will get to exactly what later on. To begin import numpy.
In [119]:
import numpy as np
In [120]:
def GSrelax(nx,ny,u,f,iters=1):
'''
Gauss Seidel smoothing
'''
dx=1.0/nx
dy=1.0/ny
Ax=1.0/dx**2
Ay=1.0/dy**2
Ap=1.0/(2.0*(Ax+Ay))
#BCs. Homogeneous Dirichlet BCs
u[ 0,:] = -u[ 1,:]
u[-1,:] = -u[-2,:]
u[:, 0] = -u[:, 1]
u[:,-1] = -u[:,-2]
for it in range(iters):
for c in [0,1]:#Red Black ordering
for i in range(1,nx+1):
start = 1 + (i%2) if c == 0 else 2 - (i%2)
for j in range(start,ny+1,2):
u[i,j]= Ap*( Ax*(u[i+1,j]+u[i-1,j])
+Ay*(u[i,j+1]+u[i,j-1]) - f[i,j])
#BCs. Homogeneous Dirichlet BCs
u[ 0,:] = -u[ 1,:]
u[-1,:] = -u[-2,:]
u[:, 0] = -u[:, 1]
u[:,-1] = -u[:,-2]
#calculate the residual
res=np.zeros([nx+2,ny+2])
for i in range(1,nx+1):
for j in range(1,ny+1):
res[i,j]=f[i,j] - ((Ax*(u[i+1,j]+u[i-1,j])+Ay*(u[i,j+1]+u[i,j-1]) - 2.0*(Ax+Ay)*u[i,j]))
return u,res
In [121]:
def prolong(nx,ny,v):
'''
interpolate 'v' to the fine grid
'''
v_f=np.zeros([2*nx+2,2*ny+2])
for i in range(1,nx+1):
for j in range(1,ny+1):
v_f[2*i-1,2*j-1] = 0.5625*v[i,j]+0.1875*(v[i-1,j]+v[i,j-1])+0.0625*v[i-1,j-1]
v_f[2*i ,2*j-1] = 0.5625*v[i,j]+0.1875*(v[i+1,j]+v[i,j-1])+0.0625*v[i+1,j-1]
v_f[2*i-1,2*j ] = 0.5625*v[i,j]+0.1875*(v[i-1,j]+v[i,j+1])+0.0625*v[i-1,j+1]
v_f[2*i ,2*j ] = 0.5625*v[i,j]+0.1875*(v[i+1,j]+v[i,j+1])+0.0625*v[i+1,j+1]
return v_f
This is exactly the opposite of the interpolation. It takes values from the find grid and transfers them onto the coarse grid. It is kind of an averaging process. This is fundamentally different from interpolation. Each coarse grid point is surrounded by four fine grid points. So quite simply we take the value of the coarse point to be the average of 4 fine points. Here 'v' is the fine grid quantity and 'v_c' is the coarse grid quantity
In [122]:
def restrict(nx,ny,v):
'''
restrict 'v' to the coarser grid
'''
v_c=np.zeros([nx+2,ny+2])
for i in range(1,nx+1):
for j in range(1,ny+1):
v_c[i,j]=0.25*(v[2*i-1,2*j-1]+v[2*i,2*j-1]+v[2*i-1,2*j]+v[2*i,2*j])
return v_c
Note that we have looped over the coarse grid in both the cases above. It is easier to access the variables this way. The last part is the Bottom Solver. This must be something that gives us the exact/converged solution to what ever we feed it. What we feed to the bottom solver is the problem at the coarsest level. This has generally has very few points (e.g 2x2=4 in our case) and can be solved exactly by the smoother itself with few iterations. That is what we do here but, any other direct method can also be used. 50 Iterations are used here. If we coarsify to just one point, then just one iteration will solve it exactly.
Now that we have all the parts, we are ready to build our multigrid algorithm. First we will look at a V-cycle. It is self explanatory. It is a recursive function ,i.e., it calls itself. It takes as input an initial guess 'u', the rhs 'f', the number of multigrid levels 'num_levels' among other things. At each level the V cycle calls another V-cycle. At the lowest level the solving is exact.
In [123]:
def V_cycle(nx,ny,num_levels,u,f,level=1):
'''
V cycle
'''
if(level==num_levels):#bottom solve
u,res=GSrelax(nx,ny,u,f,iters=50)
return u,res
#Step 1: Relax Au=f on this grid
u,res=GSrelax(nx,ny,u,f,iters=1)
#Step 2: Restrict residual to coarse grid
res_c=restrict(nx//2,ny//2,res)
#Step 3:Solve A e_c=res_c on the coarse grid. (Recursively)
e_c=np.zeros_like(res_c)
e_c,res_c=V_cycle(nx//2,ny//2,num_levels,e_c,res_c,level+1)
#Step 4: Interpolate(prolong) e_c to fine grid and add to u
u+=prolong(nx//2,ny//2,e_c)
#Step 5: Relax Au=f on this grid
u,res=GSrelax(nx,ny,u,f,iters=1)
return u,res
Thats it! Now we can see it in action. We can use a problem with a known solution to test our code. The following functions set up a rhs for a problem with homogenous dirichlet BC on the unit square.
In [124]:
#analytical solution
def Uann(x,y):
return (x**3-x)*(y**3-y)
#RHS corresponding to above
def source(x,y):
return 6*x*y*(x**2+ y**2 - 2)
Let us set up the problem, discretization and solver details. The number of divisions along each dimension is given as a power of two function of the number of levels. In principle this is not required, but having it makes the inter-grid transfers easy. The coarsest problem is going to have a 2-by-2 grid.
In [125]:
#input
max_cycles = 18
nlevels = 6
NX = 2*2**(nlevels-1)
NY = 2*2**(nlevels-1)
tol = 1e-15
In [126]:
#the grid has one layer of ghost cellss
uann=np.zeros([NX+2,NY+2])#analytical solution
u =np.zeros([NX+2,NY+2])#approximation
f =np.zeros([NX+2,NY+2])#RHS
#calcualte the RHS and exact solution
DX=1.0/NX
DY=1.0/NY
xc=np.linspace(0.5*DX,1-0.5*DX,NX)
yc=np.linspace(0.5*DY,1-0.5*DY,NY)
XX,YY=np.meshgrid(xc,yc,indexing='ij')
uann[1:NX+1,1:NY+1]=Uann(XX,YY)
f[1:NX+1,1:NY+1] =source(XX,YY)
Now we can call the solver
In [127]:
print('mgd2d.py solver:')
print('NX:',NX,', NY:',NY,', tol:',tol,'levels: ',nlevels)
for it in range(1,max_cycles+1):
u,res=V_cycle(NX,NY,nlevels,u,f)
rtol=np.max(np.max(np.abs(res)))
if(rtol<tol):
break
error=uann[1:NX+1,1:NY+1]-u[1:NX+1,1:NY+1]
print(' cycle: ',it,', L_inf(res.)= ',rtol,',L_inf(true error): ',np.max(np.max(np.abs(error))))
error=uann[1:NX+1,1:NY+1]-u[1:NX+1,1:NY+1]
print('L_inf (true error): ',np.max(np.max(np.abs(error))))
True error is the difference of the approximation with the analytical solution. It is largely the discretization error. This what would be present when we solve the discrete equation with a direct/exact method like gaussian elimination. We see that true error stops reducing at the 4th cycle. The approximation is not getting any better after this point. So we can stop after 4 cycles. But, in general we dont know the true error. In practice we use the norm of the (relative) residual as a stopping criterion. As the cycles progress the floating point round-off error limit is reached and the residual also stops decreasing.
This was the multigrid V cycle. We can use this as preconditioner to a Krylov solver. But before we get to that let's complete the multigrid introduction by looking at the Full Multi-Grid algorithm. You can skip this section safely.
We started with a zero initial guess for the V-cycle. Presumably, if we had a better initial guess we would get better results. So we solve a coarse problem exactly and interpolate it onto the fine grid and use that as the initial guess for the V-cycle. The result of doing this recursively is the Full Multi-Grid(FMG) Algorithm. Unlike the V-cycle which was an iterative procedure, FMG is a direct solver. There is no successive improvement of the approximation. It straight away gives us an approximation that is within the discretization error. The FMG algorithm is given below.
In [128]:
def FMG(nx,ny,num_levels,f,nv=1,level=1):
if(level==num_levels):#bottom solve
u=np.zeros([nx+2,ny+2])
u,res=GSrelax(nx,ny,u,f,iters=50)
return u,res
#Step 1: Restrict the rhs to a coarse grid
f_c=restrict(nx//2,ny//2,f)
#Step 2: Solve the coarse grid problem using FMG
u_c,_=FMG(nx//2,ny//2,num_levels,f_c,nv,level+1)
#Step 3: Interpolate u_c to the fine grid
u=prolong(nx//2,ny//2,u_c)
#step 4: Execute 'nv' V-cycles
for _ in range(nv):
u,res=V_cycle(nx,ny,num_levels-level,u,f)
return u,res
Lets call the FMG solver for the same problem
In [129]:
print('mgd2d.py FMG solver:')
print('NX:',NX,', NY:',NY,', levels: ',nlevels)
u,res=FMG(NX,NY,nlevels,f,nv=1)
rtol=np.max(np.max(np.abs(res)))
print(' FMG L_inf(res.)= ',rtol)
error=uann[1:NX+1,1:NY+1]-u[1:NX+1,1:NY+1]
print('L_inf (true error): ',np.max(np.max(np.abs(error))))
Well... It works. The residual is large but the true error is within the discretization level. FMG is said to be scalable because the amount of work needed is linearly proportional to the the size of the problem. In big-O notation, FMG is $\mathcal{O}(N)$. Where N is the number of unknowns. Exact methods (Gaussian Elimination, LU decomposition ) are typically $\mathcal{O}(N^3)$
A preconditioner is matrix is an easily invertible approximation of the coefficient matrix. We dont explicitly need a matrix because we dont access the elements by index, coefficient matrix or preconditioner. What we do need is the action of the matrix on a vector. That is we need the matrix-vector product only. The coefficient matrix can be defined as a function that takes in a vector and returns the matrix vector product.
Now, a stationary iterative method for solving an equation can be written as a Richardson iteration. When the initial guess is set to zero and one iteration is performed, what you get is the action of the preconditioner on the RHS vector. That is, we get a preconditioner-vector product, which is what we want.
This allows us to use any blackbox function stationary iterative method as a preconditioner
We can use the multigrid V-cycle as a preconditioner this way. We cant use FMG because it is not an iterative method.
The matrix as a function can be defined using LinearOperator from scipy.sparse.linalg. It gives us an object which works like a matrix in-so-far as the product with a vector is concerned. It can be used as a regular 2D numpy array in multiplication with a vector. This can be passed to GMRES() or BiCGStab() as a preconditioner.
Having a symmetric preconditioner would be nice because it will retain the symmetry if the original problem is symmetric. The multigrid V-cycle above is not symmetric because the Gauss-Seidel preconditioner is unsymmetric. If we were to use jacobi method, or symmetric Gauss-Seidel (SGS) method, then symmetry would be retained. As such Conjugate Gradient method will not work here becuase our preconditioner is not symmetric.
Below is the code for defining a V-Cycle preconditioner
In [130]:
from scipy.sparse.linalg import LinearOperator,bicgstab
def MGVP(nx,ny,num_levels):
'''
Multigrid Preconditioner. Returns a (scipy.sparse.linalg.) LinearOperator that can
be passed to Krylov solvers as a preconditioner.
'''
def pc_fn(v):
u =np.zeros([nx+2,ny+2])
f =np.zeros([nx+2,ny+2])
f[1:nx+1,1:ny+1] =v.reshape([nx,ny])
#perform one V cycle
u,res=V_cycle(nx,ny,num_levels,u,f)
return u[1:nx+1,1:ny+1].reshape(v.shape)
M=LinearOperator((nx*ny,nx*ny), matvec=pc_fn)
return M
Let us define the Poisson matrix also as a Linear Operator
In [131]:
def Laplace(nx,ny):
'''
Action of the Laplace matrix on a vector v
'''
def mv(v):
u =np.zeros([nx+2,ny+2])
ut=np.zeros([nx,ny])
u[1:nx+1,1:ny+1]=v.reshape([nx,ny])
dx=1.0/nx; dy=1.0/ny
Ax=1.0/dx**2;Ay=1.0/dy**2
#BCs. Homogenous Dirichlet
u[ 0,:] = -u[ 1,:]
u[-1,:] = -u[-2,:]
u[:, 0] = -u[:, 1]
u[:,-1] = -u[:,-2]
for i in range(1,nx+1):
for j in range(1,ny+1):
ut[i-1,j-1]=(Ax*(u[i+1,j]+u[i-1,j])+Ay*(u[i,j+1]+u[i,j-1]) - 2.0*(Ax+Ay)*u[i,j])
return ut.reshape(v.shape)
return mv
The nested function is required because "matvec" in LinearOperator takes only one argument-- the vector. But we require the grid details and boundary condition information to create the Poisson matrix. Now will use these to solve a problem. Unlike earlier where we used an analytical solution and RHS, we will start with a random vector which will be our exact solution, and multiply it with the Poisson matrix to get the Rhs vector for the problem. There is no analytical equation associated with the matrix equation.
The scipy sparse solve routines do not return the number of iterations performed. We can use this wrapper to get the number of iterations
In [132]:
def solve_sparse(solver,A, b,tol=1e-10,maxiter=500,M=None):
num_iters = 0
def callback(xk):
nonlocal num_iters
num_iters+=1
x,status=solver(A, b,tol=tol,maxiter=maxiter,callback=callback,M=M)
return x,status,num_iters
Lets look at what happens with and without the preconditioner.
In [134]:
A = LinearOperator((NX*NY,NX*NY), matvec=Laplace(NX,NY))
#Exact solution and RHS
uex=np.random.rand(NX*NY,1)
b=A*uex
#Multigrid Preconditioner
M=MGVP(NX,NY,nlevels)
u,info,iters=solve_sparse(bicgstab,A,b,tol=1e-10,maxiter=500)
print('without preconditioning. status:',info,', Iters: ',iters)
error=uex.reshape([NX,NY])-u.reshape([NX,NY])
print('error :',np.max(np.max(np.abs(error))))
u,info,iters=solve_sparse(bicgstab,A,b,tol=1e-10,maxiter=500,M=M)
print('With preconditioning. status:',info,', Iters: ',iters)
error=uex.reshape([NX,NY])-u.reshape([NX,NY])
print('error :',np.max(np.max(np.abs(error))))
Without the preconditioner ~150 iterations were needed, where as with the V-cycle preconditioner the solution was obtained in far fewer iterations. There we have it. A Multigrid Preconditioned Krylov Solver. We did all this without even having to deal with an actual matrix. How great is that! I think the next step should be solving a non-linear problem without having to deal with an actual Jacobian (matrix).