Multigrid_3d_notebook-checkpoint


Multigrid 3D

This notebook documents the 3D multigrid solver. There are two versions available in the repo. One is vectorized, resulting in a significant speed up. The vectorized one is documented here. It is fast enough to be usable for small problems on a desktop.

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

Restriction and Prolongation

Straight forward extension from the 2d case for both of these. Trilinear interpolation weights for the prolonation step.


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

V cycle and FMG

This is pretty much same as that in the 2D case.


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)


mgd3d.py FMG solver:
NX: 128 , NY: 128 NZ: 128 , tol: 1e-07 levels:  8
  Elapsed time:  1.0792038440704346  seconds
  L_inf(true error):  0.000731015254836

128^3 = 2 Million unknowns and solved in one second on my desktop (2010 intel i7)!


In [ ]:


In [ ]:


In [ ]: