Smoother being used is Jacobi method. There is an optimization trick here. After entering a coarse grid, we always use a zero initial guess. So, for the first iteration, only the source term is needed in the rhs as everything else is zero. Typically since we might do only one pre relaxation sweep, we can calculate this directly and save some computation. In a V cycle, on the finest grid we have a non zero initial guess or inhomogenous BCs that result in 'u' not being completely zero. So we use an if condition making sure that it is a pre sweep, not in the finest level and do the minimal computation.
Jacobi does not seem to be as good as two color GS for the FMG algorithm. FMG(1,V(1,1)) with Jacobi smoother is not converging sufficiently. FMG(2,V(1,3)) and FMG(3,V(1,2)) are needed. One possible vectorized version of red-black GS can be done using 8 steps. This is turned out to be much slower, presumable due to the memory access pattern. So Jacobi it is.
All the actions are written in vector form.
In [1]:
def Jacrelax(level,nx,ny,nz,u,f,iters=1,pre=False):
'''
Jacobi method smoothing
'''
dx=1.0/nx; dy=1.0/ny; dz=1.0/nz
Ax=1.0/dx**2; Ay=1.0/dy**2; Az=1.0/dz**2
Ap=1.0/(2.0*(1.0/dx**2+1.0/dy**2+1.0/dz**2))
#Dirichlet BC
u[ 0,:,:] = -u[ 1,:,:]
u[-1,:,:] = -u[-2,:,:]
u[: ,0,:] = -u[:, 1,:]
u[:,-1,:] = -u[:,-2,:]
u[:,:, 0] = -u[:,:, 1]
u[:,:,-1] = -u[:,:,-2]
#if it is a pre-sweep, u is fully zero (on the finest grid depends on BC, always true on coarse grids)
# we can save some calculation, if doing only one iteration, which is typically the case.
if(pre and level>1):
u[1:nx+1,1:ny+1,1:nz+1] = -Ap*f[1:nx+1,1:ny+1,1:nz+1]
#Dirichlet BC
u[ 0,:,:] = -u[ 1,:,:]
u[-1,:,:] = -u[-2,:,:]
u[: ,0,:] = -u[:, 1,:]
u[:,-1,:] = -u[:,-2,:]
u[:,:, 0] = -u[:,:, 1]
u[:,:,-1] = -u[:,:,-2]
iters=iters-1
for it in range(iters):
u[1:nx+1,1:ny+1,1:nz+1] = Ap*(Ax*(u[2:nx+2,1:ny+1,1:nz+1] + u[0:nx,1:ny+1,1:nz+1])
+ Ay*(u[1:nx+1,2:ny+2,1:nz+1] + u[1:nx+1,0:ny,1:nz+1])
+ Az*(u[1:nx+1,1:ny+1,2:nz+2] + u[1:nx+1,1:ny+1,0:nz])
- f[1:nx+1,1:ny+1,1:nz+1])
#Dirichlet BC
u[ 0,:,:] = -u[ 1,:,:]
u[-1,:,:] = -u[-2,:,:]
u[: ,0,:] = -u[:, 1,:]
u[:,-1,:] = -u[:,-2,:]
u[:,:, 0] = -u[:,:, 1]
u[:,:,-1] = -u[:,:,-2]
#if residual not needed
if(not pre):
return u,None
res=np.zeros([nx+2,ny+2,nz+2])
res[1:nx+1,1:ny+1,1:nz+1]=f[1:nx+1,1:ny+1,1:nz+1]-(Ax*(u[2:nx+2,1:ny+1,1:nz+1] + u[0:nx,1:ny+1,1:nz+1])
+ Ay*(u[1:nx+1,2:ny+2,1:nz+1] + u[1:nx+1,0:ny,1:nz+1])
+ Az*(u[1:nx+1,1:ny+1,2:nz+2] + u[1:nx+1,1:ny+1,0:nz])
- 2.0*(Ax+Ay+Az)*u[1:nx+1,1:ny+1,1:nx+1])
return u,res
In [2]:
def restrict(nx,ny,nz,v):
v_c=np.zeros([nx+2,ny+2,nz+2])
v_c[1:nx+1,1:ny+1,1:nz+1]=0.125*(v[1:2*nx:2,1:2*ny:2,1:2*nz:2 ]+v[1:2*nx:2,2:2*ny+1:2,1:2*nz:2 ]+v[2:2*nx+1:2,1:2*ny:2,1:2*nz:2 ]+v[2:2*nx+1:2,2:2*ny+1:2,1:2*nz:2 ]
+ v[1:2*nx:2,1:2*ny:2,2:2*nz+1:2]+v[1:2*nx:2,2:2*ny+1:2,2:2*nz+1:2]+v[2:2*nx+1:2,1:2*ny:2,2:2*nz+1:2]+v[2:2*nx+1:2,2:2*ny+1:2,2:2*nz+1:2])
return v_c
def prolong(nx,ny,nz,v):
v_f=np.zeros([2*nx+2,2*ny+2,2*nz+2])
a=27.0/64; b= 9.0/64; c= 3.0/64; d= 1.0/64
v_f[1:2*nx:2 ,1:2*ny:2 ,1:2*nz:2 ] = a*v[1:nx+1,1:ny+1,1:nz+1] + b*(v[0:nx ,1:ny+1,1:nz+1] + v[1:nx+1,0:ny ,1:nz+1] + v[1:nx+1,1:ny+1,0:nz ]) + c*(v[0:nx ,0:ny ,1:nz+1] + v[0:nx ,1:ny+1,0:nz ] + v[1:nx+1,0:ny ,0:nz ]) + d*v[0:nx ,0:ny ,0:nz ]
v_f[2:2*nx+1:2,1:2*ny:2 ,1:2*nz:2 ] = a*v[1:nx+1,1:ny+1,1:nz+1] + b*(v[2:nx+2,1:ny+1,1:nz+1] + v[1:nx+1,0:ny ,1:nz+1] + v[1:nx+1,1:ny+1,0:nz ]) + c*(v[2:nx+2,0:ny ,1:nz+1] + v[2:nx+2,1:ny+1,0:nz ] + v[1:nx+1,0:ny ,0:nz ]) + d*v[2:nx+2,0:ny ,0:nz ]
v_f[1:2*nx:2 ,2:2*ny+1:2,1:2*nz:2 ] = a*v[1:nx+1,1:ny+1,1:nz+1] + b*(v[0:nx ,1:ny+1,1:nz+1] + v[1:nx+1,2:ny+2,1:nz+1] + v[1:nx+1,1:ny+1,0:nz ]) + c*(v[0:nx ,2:ny+2,1:nz+1] + v[0:nx ,1:ny+1,0:nz ] + v[1:nx+1,2:ny+2,0:nz ]) + d*v[0:nx ,2:ny+2,0:nz ]
v_f[2:2*nx+1:2,2:2*ny+1:2,1:2*nz:2 ] = a*v[1:nx+1,1:ny+1,1:nz+1] + b*(v[2:nx+2,1:ny+1,1:nz+1] + v[1:nx+1,2:ny+2,1:nz+1] + v[1:nx+1,1:ny+1,0:nz ]) + c*(v[2:nx+2,2:ny+2,1:nz+1] + v[2:nx+2,1:ny+1,0:nz ] + v[1:nx+1,2:ny+2,0:nz ]) + d*v[2:nx+2,2:ny+2,0:nz ]
v_f[1:2*nx:2 ,1:2*ny:2 ,2:2*nz+1:2] = a*v[1:nx+1,1:ny+1,1:nz+1] + b*(v[0:nx ,1:ny+1,1:nz+1] + v[1:nx+1,0:ny ,1:nz+1] + v[1:nx+1,1:ny+1,2:nz+2]) + c*(v[0:nx ,0:ny ,1:nz+1] + v[0:nx ,1:ny+1,2:nz+2] + v[1:nx+1,0:ny ,2:nz+2]) + d*v[0:nx ,0:ny ,2:nz+2]
v_f[2:2*nx+1:2,1:2*ny:2 ,2:2*nz+1:2] = a*v[1:nx+1,1:ny+1,1:nz+1] + b*(v[2:nx+2,1:ny+1,1:nz+1] + v[1:nx+1,0:ny ,1:nz+1] + v[1:nx+1,1:ny+1,2:nz+2]) + c*(v[2:nx+2,0:ny ,1:nz+1] + v[2:nx+2,1:ny+1,2:nz+2] + v[1:nx+1,0:ny ,2:nz+2]) + d*v[2:nx+2,0:ny ,2:nz+2]
v_f[1:2*nx:2 ,2:2*ny+1:2,2:2*nz+1:2] = a*v[1:nx+1,1:ny+1,1:nz+1] + b*(v[0:nx ,1:ny+1,1:nz+1] + v[1:nx+1,2:ny+2,1:nz+1] + v[1:nx+1,1:ny+1,2:nz+2]) + c*(v[0:nx ,2:ny+2,1:nz+1] + v[0:nx ,1:ny+1,2:nz+2] + v[1:nx+1,2:ny+2,2:nz+2]) + d*v[0:nx ,2:ny+2,2:nz+2]
v_f[2:2*nx+1:2,2:2*ny+1:2,2:2*nz+1:2] = a*v[1:nx+1,1:ny+1,1:nz+1] + b*(v[2:nx+2,1:ny+1,1:nz+1] + v[1:nx+1,2:ny+2,1:nz+1] + v[1:nx+1,1:ny+1,2:nz+2]) + c*(v[2:nx+2,2:ny+2,1:nz+1] + v[2:nx+2,1:ny+1,2:nz+2] + v[1:nx+1,2:ny+2,2:nz+2]) + d*v[2:nx+2,2:ny+2,2:nz+2]
return v_f
In [3]:
def V_cycle(nx,ny,nz,num_levels,u,f,level=1):
if(level==num_levels):#bottom solve
u,res=Jacrelax(level,nx,ny,nz,u,f,iters=1)
return u,res
#Step 1: Relax Au=f on this grid
u,res=Jacrelax(level,nx,ny,nz,u,f,1,pre=True)
#Step 2: Restrict residual to coarse grid
res_c=restrict(nx//2,ny//2,nz//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,nz//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,nz//2,e_c)
#Step 5: Relax Au=f on this grid
u,res=Jacrelax(level,nx,ny,nz,u,f,1)
return u,res
def FMG(nx,ny,nz,num_levels,f,nv=1,level=1):
if(level==num_levels):#bottom solve
u=np.zeros([nx+2,ny+2,nz+2])
u,res=Jacrelax(level,nx,ny,nz,u,f,iters=1)
return u,res
#Step 1: Restrict the rhs to a coarse grid
f_c=restrict(nx//2,ny//2,nz//2,f)
#Step 2: Solve the coarse grid problem using FMG
u_c,_=FMG(nx//2,ny//2,nz//2,num_levels,f_c,nv,level+1)
#Step 3: Interpolate u_c to the fine grid
u=prolong(nx//2,ny//2,nz//2,u_c)
#step 4: Execute 'nv' V-cycles
for _ in range(nv):
u,res=V_cycle(nx,ny,nz,num_levels-level,u,f)
return u,res
In [ ]:
In [4]:
import numpy as np
import time
#analytical solution
def Uann(x,y,z,n):
return np.sin(2*n*np.pi*x)*np.sin(2*n*np.pi*y)*np.sin(2*n*np.pi*z)
#RHS corresponding to above
def source(x,y,z,n):
return -12 * (np.pi)**2 * n**2 * np.sin(2*n*np.pi*x) * np.sin(2*n*np.pi*y) * np.sin(2*n*np.pi*z)
#input
max_cycles = 20 #maximum number of V cycles
nlevels = 8 #total number of grid levels. 1 means no multigrid, 2 means one coarse grid. etc
NX = 1*2**(nlevels-1) #Nx and Ny are given as function of grid levels
NY = 1*2**(nlevels-1)
NZ = 1*2**(nlevels-1)
tol = 1e-7
#the grid has one layer of ghost cells to help apply the boundary conditions
uann=np.zeros([NX+2,NY+2,NZ+2])#analytical solution
u =np.zeros([NX+2,NY+2,NZ+2])#approximation
f =np.zeros([NX+2,NY+2,NZ+2])#RHS
#calcualte the RHS and exact solution
DX=1.0/NX
DY=1.0/NY
DZ=1.0/NZ
xc=np.linspace(0.5*DX,1-0.5*DX,NX)
yc=np.linspace(0.5*DY,1-0.5*DY,NY)
zc=np.linspace(0.5*DZ,1-0.5*DZ,NZ)
XX,YY,ZZ=np.meshgrid(xc,yc,zc)
uann[1:NX+1,1:NY+1,1:NZ+1] = Uann (XX,YY,ZZ,1)
f [1:NX+1,1:NY+1,1:NZ+1] = source(XX,YY,ZZ,1)
print('mgd3d.py FMG solver:')
print('NX:',NX,', NY:',NY,'NZ:',NZ,', tol:',tol,'levels: ',nlevels)
#start solving
tb=time.time()
u,res=FMG(NX,NY,NZ,nlevels,f,2)
print(' Elapsed time: ',time.time()-tb,' seconds')
error=uann[1:NX+1,1:NY+1,1:NZ+1]-u[1:NX+1,1:NY+1,1:NZ+1]
en=np.max(np.max(np.abs(error)))
print(' L_inf(true error): ',en)
128^3 = 2 Million unknowns and solved in one second on my desktop (2010 intel i7)!
In [ ]:
In [ ]:
In [ ]: