柱坐标下多重网格法解泊松方程-python

金刚归来

已于 2022-10-23 23:52:25 修改

阅读量825

点赞数

文章标签： python numpy 开发语言

于 2022-10-03 17:25:41 首次发布

本文链接：https://blog.csdn.net/zry1318/article/details/127154005

版权

写了一个柱坐标下多重网格法解泊松方程的code，外边界采用的是第一类边界条件，直接用解析解赋值，内边界需要注意的是在x=0的转轴上面有奇点。

柱坐标下泊松方程形式为

$\frac{\partial^2 u }{\partial x^2}+\frac{\partial^2 u }{\partial z^2}+\frac{1}{x}\frac{\partial u}{\partial x}=b(x,z)$

当x趋向于0时， $lim\frac{1}{x}\frac{\partial u}{\partial x}=\frac{\partial^2 u }{\partial x^2}$

将方程离散化，其中第三项采用中心差分格式。

开始采用了有残差的多重网格形式，一直不收敛，后来索性残差去掉了，收敛的还可以。

import math
import numpy as np  
import matplotlib.pyplot as plt
from scipy import interpolate
import time
def coarsen(x_fine, y_fine, phi_fine):# Lagrange type interpolation, two order
 
    x2 = np.linspace(x_fine[0],x_fine[-1],int(x_fine.size/2))
    y2 = np.linspace(y_fine[0],y_fine[-1],int(y_fine.size/2))
    f = interpolate.interp2d(x_fine, y_fine, phi_fine, kind='cubic')#{‘linear’, ‘cubic’, ‘quintic’}
    phi2 = f(x2,y2)
 
    return phi2, x2, y2
 
def fine(x_coarse, y_coarse, phi_coarse):# Lagrange type interpolation, two order
    x2 = np.linspace(x_coarse[0],x_coarse[-1],int(x_coarse.size*2))
    y2 = np.linspace(y_coarse[0],y_coarse[-1],int(y_coarse.size*2))
    f = interpolate.interp2d(x_coarse, y_coarse, phi_coarse, kind='cubic')#{‘linear’, ‘cubic’, ‘quintic’}
    phi2 = f(x2,y2)
 
    return phi2, x2, y2
def Relax2( b,  phi, x, z,leveli):
    omega = 1.95
    ite  = 10*3**leveli   # the iteration time, you can set by other condition
    k  =  0
    h_r = x[1]-x[0]
    h_z = z[1]-z[0]
    
    while(k<ite): #控制迭代次
        print('relax,k and leveli=',k,leveli)
        for i in range(0,x.size-1):
            for j in range(0, z.size-1):
                
                phi_ji =  np.copy(phi[j,i])
                dr_2 = h_r*h_r # dr^2
                dz_2 = h_z*h_z
                
                if (i!=0 and j!=0):# the inner boundary
 
                    a1 = -2/dr_2-2/dz_2
                    a2 = 1/dr_2+1/x[i]/2/h_r
                    a3 = 1/dr_2-1/x[i]/2/h_r
                    a4 = 1/dz_2
                    a5 = 1/dz_2
                    phi[j,i] = (1.-omega)*phi_ji-omega/a1*(a2* (phi[j,i+1])+a3* (phi[j,i-1])+a4* (phi[j+1,i])+a5* (phi[j-1,i])-b[j,i])
 
                    continue
                elif (i!=0 and j==0):
                    a1 = -2/dr_2-2/dz_2 #factor to u[j,i]
                    a2 = 1/dr_2+1/x[i]/2/h_r#factor to u[j,i+1]
                    a3 = 1/dr_2-1/x[i]/2/h_r                #factor to u[j,i-1]
                    a4 = 2/dz_2                #factor to u[j+1,i]
                    phi[j,i] = (1.-omega)*phi_ji-omega/a1*(a2* (phi[j,i+1])+a3* (phi[j,i-1])+a4* (phi[j+1,i])-b[j,i])
                    continue
                elif (i==0 and j!=0):#z axis
                    a1 = -4/dr_2-2/dz_2
                    a2 = 4/dr_2
                    a3=1/dz_2
                    a4=1/dz_2
                    phi[j,i] = (1.-omega)*phi_ji-omega/a1*(a2* (phi[j,i+1])+a3* (phi[j-1,i])+a4* (phi[j+1,i])-b[j,i])
                    continue
                else:
 
                    a2 = 4/dr_2
                    a3 = 2/dz_2
 
                    a1 = -1*a2-a3
                    phi[j,i] = (1.-omega)*phi_ji-omega/a1*(a2* (phi[j,i+1])+a3* (phi[j+1,i])-b[j,i])
 
        k = k+1
    return phi
def residual(b, phi, x, z): 
    res  =  np.zeros((z.size,x.size),dtype=float)
    for i in range(0,x.size-1):
        for j in range(0,z.size-1):
                dr_2 = h_r*h_r # dr^2
                dz_2 = h_z*h_z
                if (i==0 and j==0):
                    res[j,i] =b[0,0]-4*(phi[0,1]-phi[0,0])/dr_2-2*(phi[1,0]-phi[0,0])/dz_2
                elif (i!=0 and j==0):
                    a1 = -2/dr_2-2/dz_2 #factor to u[j,i]
                    a2 = 1/dr_2+1/x[i]/2/h_r#factor to u[j,i+1]
                    a3 = 1/dr_2-1/x[i]/2/h_r                #factor to u[j,i-1]
                    a4 = 2/dz_2                #factor to u[j+1,i]
                    res[j,i] = b[j,i]-a1*phi[j,i]-a2*phi[j,i+1]-a3*phi[j,i+1]-a4*phi[j+1,i]
                elif (i==0 and j!=0):#z axis
                    a1 = -4/dr_2-2/dz_2
                    a2 = 4/dr_2
                    a3=1/dz_2
                    a4=1/dz_2
                    res[j,i]= b[j,i]-a1*phi[j,i]-a2*phi[j,i+1]-a3*phi[j+1,i]-a4*phi[j-1,i]
                else:
                    a1 = -2/dr_2-2/dz_2
                    a2 = 1/dr_2+1/x[i]/2/h_r
                    a3 = 1/dr_2-1/x[i]/2/h_r
                    a4 = 1/dz_2
                    a5 = 1/dz_2
                    res[j,i] = b[j,i]-a1*phi[j,i]-a2*phi[j,i+1]-a3*phi[j,i-1]-a4*phi[j+1,i]-a5*phi[j-1,i]
    
    return res
 
def bij(x,z):
    b0 = density_nfw(x,z)
    return b0
def density_nfw(x,z):                #density of the NFW dark matter potential
    rho0_nfw =1# 0.00854*1e9#M0/Kpc^3
    rh = 19.6 #Kpc
    rho=np.zeros((z.size,x.size))
    for j in range(0,z.size):
        for i in range(0,x.size):
            r = math.sqrt(x[i]*x[i]+z[j]*z[j])
            rrh = r/rh
            if rrh<0.02:
                rrh=0.02
            rho[j,i] = rho0_nfw/rrh/(1+rrh)/(1+rrh)
    return rho
def phi_nfw(x,z):       # the analysical result
    rho0_nfw=1#assum 4pi g rho=1
    rh = 19.6
    phi0=np.zeros((z.size,x.size),dtype=float)
    for j in range(0,z.size):
        for i in range(x.size):
            r = math.sqrt(x[i]*x[i]+z[j]*z[j])
            rrh = r/rh
            if rrh<0.02:
                rrh=0.02
            phi0[j,i] =  -1*rh* rh*math.log(1+rrh)/rrh
    return phi0
 
def MG2( b,phi, x,z):
 
    h_r0  =  x[1]-x[0] 
    h_z0  =  z[1]-z[0]               
    vh  =  Relax2(b,  phi, x, z,1)
 
    total_level=int(math.log(x.size, 2)-1)
    t = 0
 
    while (t <1): # number of V circles, here is 1
        for i in range(1, total_level): 
            k = i
 
            v2h,x2,z2 = coarsen(x,z,vh)
            b2 = bij(x2,z2)
            vh = Relax2(b2,  v2h, x2, z2, k)
            x = x2
            z = z2
 
        for i in range(1,total_level): 
 
            v2h,x2,z2 = fine(x,z,vh)
            b2=bij(x2,z2)
            vh0 = Relax2(b2,  v2h, x2, z2, k)
            vh = vh0
            x = x2
            z = z2
            k=k-1
 
        t = t+1
        print('Vtime=',t)
    return vh
 
'''
def MG2( b,phi, x,z):
    h_r0  =  x[1]-x[0] 
    h_z0  =  z[1]-z[0]               
    vh0  =  Relax2(b,  phi, x, z,1)
    rh  =  residual(b, phi, x, z)# residual
    eh0 = np.zeros(rh.shape,dtype = float)
    eh =  Relax2(rh, eh0,  x,z,1)
    vh = vh0+eh
    rh  =  residual(b, vh,  x, z)
    t = 0
    while ((np.max(rh) > 1e-2) and (t <1)): 
        for i in range(1,5): 
            k = i
#            r2h,x2,z2 =  coarsen(x,z,rh)   # residual to fine grid
            v2h,x2,z2 = coarsen(x,z,vh)
            b2=bij(x2,z2)
            vh0 = Relax2(b2,  v2h, x2, z2, k)
            rh = residual(b2, vh0, x2, z2)
            eh =  Relax2(rh, eh0,  x,y,5)
            vh = vh0+eh
            eh0 = np.zeros(rh.shape,dtype = float)
            e2h0 = np.zeros(r2h.shape,dtype = float)
            
            e2h  =  Relax2(r2h, e2h0,x2,z2,k)
            phi = v2h+e2h
            b = bij(x2,z2)
            vh =  Relax2(b, phi, x2,z2,k)
            rh = residual(b, vh, x2,z2)
            x = x2
            z = z2
        for i in range(1,5): 
            k = 6-i
            r2h,x2,z2 =  fine(x,z,rh)   # residual to fine grid
            e2h0 = np.zeros(r2h.shape,dtype = float)# initiate the err
            e2h  =  Relax2(r2h, e2h0, x2,z2,k)     #calculate err
            v2h,x2,z2= fine(x,z,vh)                  #to fine grid
            phi = v2h+e2h                        # update phi
            b = bij(x2,z2)                         # calculate b on this level
            vh =  Relax2(b, phi, x2,z2,k)          # calculate  vh again with updated phi
            rh =  residual(b, vh, x2,z2)            #calculate residual
            x = x2
            z = z2
        t = t+1
        print('Vtime=',t)
    return vh  
'''
 
'''
def MG2(b,  phi, x, z):
    h_r  =  x[1]-x[0]
    y_r  = z[1]-z[0]
    vh  =  Relax2(b,  phi, x, z)
    rh  =  residual(b,  phi, x, z)# residual
    i = 0
    while (i<10): 
            #print(phi)
        r2h,x2 ,z2=  coarsen(x,z,rh)   # residual to coarse grid
        h_r2=x2[1]-x2[0]
        h_z2=z2[1]-z2[0]
        e2h0 = np.zeros((z2.size,x2.size),dtype = float)
        e2h  =  Relax2(r2h, e2h0,x2,z2)
        print(e2h.shape)
        eh,x, z = fine(x2,z2,e2h)
        print(vh.shape,eh.shape)
        v1h = vh+eh
        h_r = x[2]-x[1]
        h_z = z[1]-z[0]
        phi = Relax2(b, v1h, x,z )
        vh = v1h
        rh = residual(b, vh, x,z)
        i = i+1
        print(i)
    return phi
'''
def initiate():
    time_start = time.time()
 
    n_x_grid =128
    n_z_grid=128
    xs = 0.
    xe =200 # in kpc
    zs=0.
 
    ze=200
    x=np.linspace(xs,xe,n_x_grid)
    z=np.linspace(zs,ze,n_z_grid)
    phi  =  np.ones((z.size,x.size), dtype=float)
 
    phi2 = phi_nfw(x,z)
    phi[-1,:] =  (phi2[-1,:])
    phi[:,-1] =  (phi2[:,-1])
    b = density_nfw(x,z)
    h_r=x[1]-x[0]
    h_z=z[1]-z[0]
    '''
    vh0  =  Relax2(b,  phi, x, z,5)
    rh  =  residual(b, phi, x, z)# residual
    eh0 = np.zeros(rh.shape,dtype = float)
    eh =0#  Relax2(rh, eh0,  x,z,5)
    vh = vh0+eh
    #rh  =  residual(b, vh,  x,z)
    '''
    result =  MG2( b,  phi, x, z) # solve the equation
    grav=6.67e-8 #cgs
    m0=1.989e33
    kpc=3.086e21
    rho_halo0=0.00854*1e9*m0/kpc#a^2里面有个kpc的平方
    potential = 4*math.pi*grav*rho_halo0*result
    phi_error = 4*math.pi*grav*rho_halo0*(phi2-result)
    time_end = time.time() 
    #print(result.shape)
    time_c= time_end - time_start   #运行所花时间
    print('time cost', time_c, 's')
    return potential,phi_error, x, z
 
def plot(x, z, result,name):
    print(name)
    plt.figure(1)
    plt.contourf(x,z,result ,100,cmap='seismic')
 
    plt.colorbar()
    plt.title(name)
    plt.xlabel('x(kpc)')
    plt.ylabel('z(kpc)')
    plt.savefig(name+'.png')
    plt.show()
    plt.close()
 
    return 0
 
def plot1d(x,phi,name):
    print(name)
    plt.figure(1)
    plt.plot(x,phi)
 
    plt.title(name)
    plt.xlabel('x(kpc)')
    plt.ylabel('phi')
    plt.savefig(name+'.png')
    plt.show()
    plt.close()
 
    return 0
potent, err, x, z = initiate()
plot(x,z, potent,'Numerical results(erg)')
 
plot(x,z,err,'error')
plot1d(x,err[0,:],'err1d')