2D热传导方程的FDM-非齐次边界

# -*- coding: utf-8 -*-
"""
Created on Wed Apr 29 12:35:31 2020

@author: 
"""
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from  matplotlib import cm
import scipy.linalg as la

def Heat_Equation_1_2D(X,T,N,m,InitialCondition,Boundary_left,Boundary_right,D):
    
    if(D==1):
        n=N
        x1=np.linspace(0,1,50)
        y1=InitialCondition(x1)#map(InitialCondition,x)
        plt.figure(figsize=(10,6))
        plt.plot(x1,y1)
        plt.xlabel('$x$',fontsize=15)
        plt.ylabel('$f(x)$',fontsize=15)
        plt.title('1DHeat Equation Initial Condition')
        print(X)
        x=X
        t=T
        h=(max(x)-min(x))/n  #  步长 x轴方向
        k=(max(t)-min(t))/m # 时间步长
    
        kappa=1
        rho=kappa*k/h/h
        print('rho')
        '''
        显式格式不能任意取时间和空间的网格点数，即M与N不能随意取值。
        我们称显式格式为条件稳定。特别地，需       要满足所谓CFL条件（Courant–Friedrichs–Lewy）：
        '''
        if(rho<=0.5):
            print('显示格式此时满足条件稳定：')
        print('rh0=%f' % (rho))
        U=np.zeros((n+1,m+1)) #
        #print(U)
        
        for i in range(1,n):  #表示列数
            U[i][0]=InitialCondition(x[i])
        for j in range(0,m+1):   #表示 行数
            U[0][j]=Boundary_left(t[j])
            U[n][j]=Boundary_right(t[j])
        #print('初始化的U：')
        #print(U)
    
        for j in range(0,m):#  (1,  3)
            for i in range(1,n):  # (0,  5)
                #rho*U[j-1][k]+(1-2*rho)*U[j][k]+rho*U[j+1][k]
                #U[i+1][j]=rho*U[i][j-1]+(1-2*rho)*U[i][j]+rho*U[i][j+1]
                U[i][j+1]=rho*U[i+1][j]+(1-2*rho)*U[i][j]+rho*U[i-1][j]
        #print(U.shape)
        
        plt.figure(figsize=(10,6))
        plt.plot(x, U.T[0,:])
        plt.plot(x, U.T[int(0.10/ k),:])
        plt.plot(x, U.T[int(0.15/ k),:])
        plt.plot(x, U.T[int(0.20/ k),:])
        plt.xlabel('$x$', fontsize = 15)
        plt.ylabel(r'$U(dot, tau)$', fontsize = 15)
        plt.title(u'1D Heat equation', fontsize = 15)
        plt.legend([r'$t = 0.$', r'$t = 0.10$', r'$t = 0.15$', r'$t = 0.20$'], fontsize = 15)
        
        
        
        fig=plt.figure(figsize=(10,6))
        ax=Axes3D(fig)
        x1,y1=np.meshgrid(t,x)
        surf=ax.plot_surface(x1,y1,U,rstride=1,cstride=1,cmap=cm.coolwarm,alpha=0.9)   #绘面'rainbow'
        ax.set_xlabel('t')
        ax.set_ylabel('x')
        ax.set_zlabel('U')
        ax.set_xticks([0,0.05,0.1,0.15,0.2])
        ax.set_yticks([0,0.2,0.4,0.6,0.8,1])
        ax.set_zticks([0,0.2,0.4,0.6,0.8,1])
        ax.set_title('Hot equation$ u_t=A u_{xx} $')     
        ax.grid(False)
        
        cbar=plt.colorbar(surf,shrink=0.5,aspect=5)
        plt.show()
    #%%
        fig1=plt.figure(figsize=(10,6))
        #ax1=Axes3D(fig1)
        ax1=fig1.add_subplot(1,1,1,projection='3d')
        x1,y1=np.meshgrid(x,t)
        surf1=ax1.plot_surface(x1,y1,U.T,cmap=cm.coolwarm)   #绘面'rainbow' rstride=1,cstride=1,cmap=cm.coolwarm,alpha=0.9
        
        ax1.set_xlabel('x',fontdict={"size":16})
        ax1.set_ylabel('t',fontdict={"size":16})
        ax1.set_zlabel('U',fontdict={"size":16})
        ax1.set_xticks([0,0.2,0.4,0.6,0.8,1])
        ax1.set_yticks([0,0.05,0.1,0.15,0.2])
        ax1.set_xticks([0,0.2,0.4,0.6,0.8,1])
        ax1.set_title('Hot equation$ u_t=A u_{xx} ,$')     
        ax1.grid(False)
        cbar=fig1.colorbar(surf1,shrink=0.75,aspect=5)
        plt.show()
#%%        
    if(D==2):
        
        Mx=N[0]
        My=N[1]
        a=0.5
        ox=(max(X[0])-min(X[0]))/Mx    #X轴方向的步长   值为  1 
        x=X[0]     # 6 个点  0 1 2 3 4 5     
        oy=(max(X[1])-min(X[1]))/My    #Y 轴方向的步长   值 为 1                
        y=X[1]  #  6个点 0 1 2 3 4 5         
        ot=T/m  #     #时间 方向的步长 值为 0.04       
        t=np.linspace(0,T,m+1)   #  个数 为6  为 0 0.04 0.08 0.12 0.16 0.2   改变 N 的值  研究
        u=np.zeros((My+1,Mx+1))  #构造以X轴方向 Y方向点数  的矩阵
    #%%    
        #初始化 矩阵  U 的值
        for j in range(0,Mx+1):
            for i in range(0,My+1):
                u[i][j]=InitialCondition(x[j],y[i])
    #%%    
        rx=a*ot/(ox**2)   #结果为 一个数 0.02         
        ry=a*ot/(oy**2)    #结果为 一个数 0.02        
        rx1=1+2*rx       #结果为 一个数 1.04
        rx2=1-2*rx
        ry1=1+2*ry
        ry2=1-2*ry
    #%%    
        # A 为Y方向 隐式 时的系数矩阵
        A=np.zeros((My-1,Mx-1))
        for j in range(0,Mx-1):
            A[j][j]=ry1
            if j>0:
                A[j-1][j]=-ry
                A[j][j-1]=-ry
        #print(A)
    #%%%%
        # B 为Y方向 隐式 时的系数矩阵
        B=np.zeros((My-1,Mx-1))
        for i in range(0,My-1):
            B[i][i]=rx1
            if i>0:
                B[i-1][i]=-rx
                B[i][i-1]=-rx

    #%%
        for k in range(1,m+1):# 时间 T 被分为了 6 个点  5 段  5个间隔   从 第二个时间 节点 开始 循环
            u_1=u     # 把 u 为 一个 6 X 6 的矩阵 复制 给它
            t=k*ot     # ot =T/m      0  0.04  0.08  0.12  0.16  0.2   {0.2/5 = =0.04}
            
            # 时间方向的步长ot   节点 为 t 的值
            # 初值条件 只关于 XY的函数 f(X,Y)
            #边值条件 只关于 X Y t的  二维中
            for i in range(0,My+1):          #边值处理 1   处理 第一列 和最后一列 的边值 
                u[i][0]=Boundary_left(x[0],y[i],t)
                u[i][Mx]=Boundary_left(x[Mx],y[i],t)
            
            for j in range(0,Mx+1):          ##边值处理 2   处理 第一行 和最后一行 的边值
                u[0][j]=Boundary_right(x[j],y[0],t)
                u[My][j]=Boundary_right(x[j],y[My],t)
    
            bx=np.zeros((Mx-1))    # bx  为隐式时Y方向的常数 项
    
            by=np.zeros((My-1))    # by  为隐式时X方向的常数 项
            if((divmod(k,2)[1]==0)):
                for i in range(0,My-1):
                    b2=np.zeros((My-1))
                    b2[0]=ry*u[i+1][0]
                    b2[My-2]=ry*u[i+1][My]
                    b3=b2.T
                    bx=b3+rx*(u_1[i,1:Mx] + u_1[i+2,1:Mx] ) + rx2*u_1[i+1,1:Mx]
                    u[i+1,1:Mx]=la.solve(A,bx.T)
            else:
                for j in range(0,Mx-1):                       
                    c2=np.zeros((Mx-1))
                    c2[0]=rx*u[0][j+1]
                    c2[Mx-2]=rx*u[Mx][j+1]
                    c3=c2.T
                    by=c3+ ry*(u_1[1:My,j] + u_1[1:My,j+2])+ry2*u_1[1:My,j+1]
                    u[1:My,j+1]=la.solve(B,by)

        fig = plt.figure(figsize=(8,4))
        ax = Axes3D(fig)
        x, y = np.meshgrid(x,y)
        surf1=ax.plot_surface(x,y,u.T,rstride=1,cstride=1,cmap=cm.coolwarm,linewidth=0,antialiased=False)  #cmap='viridis'
        
        #ax.set_xlim(0,5)
        #ax.set_ylim(0,5)
        #ax.set_zlim(-10,20)
        ax.set_xlabel('$x$',fontdict={"size":14})
        ax.set_ylabel('$y$',fontdict={"size":14})
        ax.set_zlabel('U',fontdict={"size":14})
        '''
        ax.set_xticks([0,0.4,0.8,1.2,1.6,2])
        ax.set_yticks([0,0.4,0.8,1.2,1.6,2])
        ax.set_zticks([1,1.5,2,2.5])
        '''
        ax.set_title('Hot equation$ u_t=A u_{xx}+A u_{yy} $')     
        ax.grid(False)
        cbar=fig.colorbar(surf1,shrink=0.75,aspect=5)
        plt.show()
        U=u
        return U
#%%
On_Off=2
if On_Off==1:
    N,m=25,1500
    T,x=np.linspace(0,0.2,m+1),np.linspace(0,1,N+1)  #m+1项数列 
    Initia=lambda x :4.0*(1-x)*x                     #初值条件
    Boundary_left=lambda t:0                         #左边值
    Boundary_right=lambda t:0                        #右边值
    Heat_Equation_1_2D(x,T,N,m,Initia,Boundary_left,Boundary_right,D=1)
#%%
if On_Off==2:
    InitialCondition=lambda x,y: 0
    '''
    Boundary_left=lambda x,y,t:x**2*np.sin(y)-y**2*np.sin(x)
    Boundary_right=lambda x,y,t:x**2*np.sin(y)-y**2*np.sin(x)
    '''
    Boundary_left=lambda x,y,t:x*np.sin(np.pi*y)-y*np.sin(np.pi*x)
    Boundary_right=lambda x,y,t:x*np.sin(np.pi*y)-y*np.sin(np.pi*x)
    
    n=50   # X轴方向的网格数目
    o=50   # Y方向的网格数目
    m=1000   # Z时间方向的网格数目   一般 考虑 X 和Y轴网格数目固定 时  时间方向的网格数目变化 的情况
    T=0.2             #取一个 所研究温度的最大值 
    N=np.array([n,o])    #调节 m和 T 来研究 接的变化情况
    X=np.array([np.linspace(0,2,n+1),np.linspace(0,4,o+1)])
    Heat_Equation_1_2D(X,T,N,m,InitialCondition,Boundary_left,Boundary_right,D=2)
#%%    
else:
    print('3D及以上还在努力中,请期待!!!!!!!!!!!!!!!')