多重网格法-松弛迭代法-二维泊松方程-python实现

这几天在家躲避疫情，闲来无事，写了这个多重网格法求解泊松方程的算法的代码。
多重网格法可能是目前为止解泊松方程最快的算法，n个格点需要n次计算就可以收敛，而快速傅里叶变换的收敛速度是n*logn, 共轭梯度法是n^2.。多重网格法可以方便的应对各种边界条件，这一点比傅里叶变换之类的谱方法要好得多。

多重网格法可以这么理解。泊松方程化为差分方程后，每个格点都可以写成一个方程，因此得到一个方程组。使用迭代法解这个方程组的时候，粗网格可以快速消除低频误差，细网格消除高频误差，收敛速度远快于直接使用细网格做迭代。

本程序写的时候只考虑了正方形网格的情形，但是可以很方便的改为长方形网格以及不等距长方形网格（目前不支持极坐标）。迭代方法采用了SOR松弛迭代法。例题选取的是我之前用共轭梯度法解的一个泊松方程。
4.

import numpy as np  
import matplotlib.pyplot as plt
import math
from pylab import *  

def fine2(grid0,hx0,hy0):   #化成细网格
    nx_grid2  =  int((hx0.size)*2 + 1) #x 方向粗化后网格数量
    ny_grid2  =  int((hy0.size)*2 + 1)        
    grid2  =  np.zeros((nx_grid2,ny_grid2),  dtype   =   float)
    hx2 = np.zeros((nx_grid2-1),  dtype   =   float)
    hy2 = np.zeros((ny_grid2-1),  dtype   =   float)
    x2 = np.zeros((nx_grid2),  dtype   =   float)
    y2 = np.zeros((ny_grid2),  dtype   =   float)
    x2[0] = xs
    y2[0] = ys
    for i in range(1, hx0.size+1):
        i2 = i*2
        for j in range(1, hy0.size+1):              #赋值中间的网格
            j2 = j*2
            grid2[i2,j2] = grid0[i,j]
            grid2[i2-1,j2] = grid0[i,j]*0.5+grid0[i-1,j]*0.5
            grid2[i2,j2-1] = grid0[i,j]*0.5+grid0[i,j-1]*0.5
            grid2[i2-1,j2-1] = grid0[i,j]*0.25+grid0[i,j-1]*0.25+grid0[i-1,j]*0.25+grid0[i-1,j-1]*0.25
    for i in range(0, hy0.size):#赋值边界
        grid2[0,2*i]  =  grid0[0,i]
        grid2[0,2*i+1]  =  0.5*grid0[0,i]+0.5*grid0[0,i+1] 
        grid2[-1,2*i]  =  grid0[-1,i]
        grid2[-1,2*i+1]  = 0.5* grid0[-1,i]+0.5*grid0[-1,i+1] 
    grid2[0,int(hy0.size*2)] =  grid0[0,hy0.size]  #边界最后一个点
    grid2[-1,int(hy0.size*2)] =  grid0[-1,hy0.size]  #边界最后一个点

    for j in range(0, hx0.size):#赋值边界
        grid2[2*j,0]  =  grid0[j,0]
        grid2[2*j+1,0]  =  0.5*grid0[j,0]+0.5*grid0[j+1,0] 
        grid2[2*j,-1]  =  grid0[j,-1]
        grid2[2*j+1,-1]  = 0.5* grid0[j,-1]+0.5*grid0[j+1,-1] 
    grid2[int(hx0.size*2),0] =  grid0[hx0.size,0]  #边界最后一个点
    grid2[int(hx0.size*2),-1] =  grid0[hx0.size,-1]  #边界最后一个点

    return grid2
    

def hx_level(xsize):  #根据数组的大小判断是第几重网格，返回这一重的步长，x，y坐标
    level = int(math.log2((nx_grid-1)/xsize))#     计算粗化了几次
    #level_total  =  int(math.log2(phi.size-1))-5
    dstep  =  int(2**(level))
    hx2_grid  =  int((hx.size)/dstep)
    print(level,dstep,hx.size,hx2_grid)
    hy2_grid  =  int((hy.size)/dstep)
    hx2  =  np.zeros(hx2_grid,  dtype   =   float)
    hy2  =  np.zeros(hy2_grid,  dtype   =   float)
    for  i in range(0, hx2.size):                            
        j0  =  dstep*i
        h2_step  =  0.
        for j in range(dstep):    
            h2_step  =  h2_step + hx[j0 + j]
        hx2[i]  =  h2_step
    print(hx2)

    for  i in range(0, hy2.size):                            
        j0  =  dstep*i
        h2_step  =  0.
        for j in range(dstep):    
            h2_step  =  h2_step + hy[j0 + j]
        hy2[i]  =  h2_step
    x2 = np.zeros((hx2.size+1),  dtype   =   float)
    y2 = np.zeros((hy2.size+1),  dtype   =   float)
    x2[0] = xs
    y2[0] = ys
    for i in range(1, x2.size):# 赋值x2,默认不等距网格
        x2[i] = x2[i-1]+hx2[i-1]
    for j in range(1, y2.size):# 赋值y2,默认不等距网格
        y2[j] = y2[j-1]+hy2[j-1]
    return hx2,hy2,x2,y2
    

def coarsen2(grid0,hx0,hy0):
    #dstep  =  int(2**(level-1))
    nx_grid2  =  int((hx0.size)/2 + 1) #x 方向粗化后网格数量
    ny_grid2  =  int((hy0.size)/2 + 1)
    grid2  =  np.zeros((nx_grid2,ny_grid2),  dtype   =   float)
    hx2 = np.zeros((nx_grid2-1),  dtype   =   float)
    hy2 = np.zeros((ny_grid2-1),  dtype   =   float)
    x2 = np.zeros((nx_grid2),  dtype   =   float)
    y2 = np.zeros((ny_grid2),  dtype   =   float)
    x2[0] = xs
    y2[0] = ys
    for i in range(1, nx_grid2-1):
        for j in range(1, ny_grid2-1):              #赋值中间的网格
           i2  =  2*i
           j2  =  2*j
           weight_total = 1./hx0[i2-1]+1./hx0[i2]+1./hy0[j2-1]+1./hy0[j2]
           weight1 = (1./hx0[i2-1])/weight_total     # 周边四个网格的权重
           weight2 = (1./hx0[i2])/weight_total
           weight3 = (1./hy0[j2-1])/weight_total
           weight4 = (1./hy0[j2])/weight_total
           grid2[i,j]  =  0.4*grid0[i2,j2] +  0.6*weight1*grid0[i2-1,j2] + 0.6*weight2*grid0[i2+1,j2]+0.6*weight3*grid0[i2,j2-1]+0.6*weight4*grid0[i2,j2+1]
        print(weight1,weight2,weight3,weight4)
    for i in range(0, ny_grid2):#赋值边界
        grid2[0,i]  =  grid0[0,2*i]
        grid2[-1,i]  =  grid0[-1,2*i]
    for j in range(0, nx_grid2):
        grid2[j,0]  =  grid0[2*j,0]
        grid2[j,-1]  =  grid0[2*j,-1]

    return grid2#,hx2,hy2,x2,y2      

def Relax2(b2, phi0,  h,x1,x2):#注意，这个子程序还没来得及写成非均匀网格的，先用正方形网格凑合下
    omig = 1.85 #松弛银子
    leveli = int(math.log2((nx_grid-1)/h.size))
    print('level',leveli)
    ite  =  20*((leveli)*leveli)
    k = -4            #k 取负)值可保证至少迭代一定次数
    while(k <ite):
        print(k)
        for j in range(0,h.size+1):  #向上
            for i in range(1,h.size):       #沿着横坐标。      非均匀网格要改这里
                
                if j  =  =  0:                     #最低下一行
                    phi0[j,i] = np.copy(phi0[j+1,i])
                elif j =  = h.size:
                    phi0[j,i] = np.copy(phi0[j-1,i])#最上一行
                else:#五点差分法
                    phi0[j,i] = (1-omig)*np.copy(phi0[j,i])+omig*(np.copy(phi0[j,i-1])+np.copy(phi0[j,i+1])+np.copy(phi0[j-1,i])+np.copy(phi0[j+1,i])+h[i]*h[i]* b2[i,j])/4.  
        k = k+1
    return phi0

def MG(b_mg,phi0, x1, x2, h):#多重网格法的子程序
    vh = Relax2(b_mg,phi0,  h, x1, x2)
    rh = residual(b_mg,vh, h, x1, x2)

    level_total  =  4#int(math.log2(phi.size-1))-4          #计算层数，V型，先粗化，再细化
    for i in range (1, level_total):   #coarse and iterate
        print(i)
        hx,hy,x11,x21 = hx_level(vh.shape[0]-1)
        hx2,hy2,x12,x22 = hx_level(int(hx.size/2))
        v2h0   =   coarsen2(vh,hx,hx)
        b2h  =  coarsen2(b_mg,hx,hx) 
        v2h = Relax2(b2h,v2h0,  hx2, x12, x22)
        vh   =   v2h
        b_mg  =  b2h
    for i in range (1, level_total):  # fine and itearate
        print(i)
        hx,hy,x11,x21 = hx_level(vh.shape[0]-1)

        hx2,hy2,x12,x22 = hx_level(int(hx.size*2))       
        b2h  =  fine2(b_mg,hx,hx) 
        v2h0   =   fine2(vh,hx,hx)
        v2h  =  Relax2(b2h,v2h0,  hx2, x12, x22)
        vh   =   v2h
        b_mg  =  b2h 
      

    return vh,x12,x22 





def init(nx_grid,ny_grid,xs,xe,ys,ye,xleft_boundary_value,xright_boundary_value,dyleft_boundary_value,dyright_boundary_value):
    hx  =  (xe-xs)/(nx_grid-1)* np.ones(nx_grid-1,  dtype   =   float)   #x 方向步长，可以是非等距
    hy  =  (ye-ys)/(ny_grid-1)* np.ones(ny_grid-1,  dtype   =   float)   #y 方向步长
    x  =  np.zeros(nx_grid,  dtype   =   float)
    x[0]  =  xs
    x[-1]  =  xe
    for i in range(1, nx_grid-1):
        x[i]  =  x[i-1] + hx[i-1]
    y  =  np.zeros(ny_grid,  dtype   =   float)
    y[0]  =  ys
    y[-1]  =  ye
    for i in range(1, ny_grid-1):             # initialize x and y
        y[i]  =  y[i-1] + hy[i-1]

    phi   =   np.ones((x.size,y.size),  dtype   =   float)*0.1   # initialize results
    phi[:,0]  =  xleft_boundary_value
    phi[:,-1]  =  xright_boundary_value

    b  =  np.zeros((x.size,y.size),  dtype   =   float)#        initialize b
    for i in range(0, nx_grid):
        for j in range(0, ny_grid):
            b[i,j] = bxy(x[i],y[j])
    return x,y,hx,hy,phi,b

def bxy(x0,y0):#方程右边的函数
    return -2.*math.pi*math.pi*math.sin(math.pi*x0)*math.cos(math.pi*y0)

nx_grid  =  257 #at least 512 grids to reach enough depth
xs  =  -1.
xe  =  1.

ny_grid  =  257  #at least 512 grids to reach enough depth
ys  =  -1.
ye  =  1.

xleft_boundary_value = 0.     # u(x = -1) = u(x = 1) = 0
xright_boundary_value = 0.
dyleft_boundary_value = 0.    # du(y = -1) = u(y = 1) = 0
dyright_boundary_value = 0.

x,y,hx,hy,phi,b =  init(nx_grid,ny_grid,xs,xe,ys,ye,xleft_boundary_value,xright_boundary_value,dyleft_boundary_value,dyright_boundary_value)

#grid2,hx2,hy2,x2,y2  = coarsen2(b,hx,hy)
#grid3 = fine2(grid2,hx2,hy2)

#print('hx',hx)
#hxlevel2,hylevel2,x22,y22 = hx_level(hx2.size)

phi2 = Relax2(b,phi,  hx,x,y)
retult,x121,y121 = MG(b,phi, x, y, hx)

plt.figure(1)
print(retult.max(),retult.min())
contourf(x121,  y121,  retult, 80,  cmap  =  'seismic')  
plt.colorbar()
'''
plt.figure(2)
contourf(y,  x,  phi2, 80,  cmap  =  'seismic')  

plt.colorbar()
'''

plt.show()
plt.close()

结果如图
评论区有个网友指出了，程序里缺少残差和逐层修正，实际上多重网格法有两种，修正法和完全逼近方式，完全逼近方式传递的是未知量本身，适用于非线性方程组。而修正法是在最细网格上求解未知量 X 自身，其它网格上求解未知量的迭代误差量，适用于线性方程组