Основы мат мода, задача 2

main.py

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import numpy as np


'''
инициализируем границы
и другие постоянные
'''
max_x, max_y, t_max = np.pi, 3, 0.1
N, M, J = 50, 50, 100
hx = max_x / (N-1)
hy = max_y / (M-1)
tau = t_max / J

'''узлы сетки'''
X = [i * hx for i in range(N + 1)]
Y = [i * hy for i in range(M + 1)]
T = [i * tau for i in range(J + 1)]


def Fy(w2, w1, w0, n, j):
    '''
    функция f первой системы
    w2=w_{i+1}, w1=w_i, w0=w_{i-1}
    '''
    return 4/hy**2 * (w2-2*w1+w0) + 2*w1/tau + np.cos(hx*n)*np.cos(tau*(j+1/2))


def Fx(w2, w1, w0, n, j):
    '''
    функция f второй системы
    w2=w_{i+1}, w1=w_i, w0=w_{i-1}
    '''
    return 4/hx**2 * (w2-2*w1+w0) + 2*w1/tau + np.cos(hx*n)*np.cos(tau*(j+1/2))


def InitCond(n):
    '''начальное условие'''
    return np.cos(n*hx)


def BuildPlot(X, Y, U):
    '''
    построение графика
    '''
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    X, Y = np.meshgrid(X, Y)
    surf = ax.plot_surface(X, Y, U, cmap='inferno', linewidth=0, antialiased=False)

    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('u')
    fig.colorbar(surf, shrink=0.5, aspect=5)

    plt.show()


def NomSol():
    '''
    численное решение
    '''
    w = np.zeros((N+1, M+1, J+1)) #массив для записи решения

    '''начальное условие'''
    for n in range(N+1):
        for m in range(M+1):
            w[n][m][0] = InitCond(n)

    w05 = np.zeros((N+1, M+1)) #массив для полуцелого слоя
    '''коэффициенты a, b, c'''
    Ax = 4/hx**2
    Ay = 4/hy**2
    Bx = 4/hx**2
    By = 4/hy**2
    Cx = 2/tau + 8/hx**2
    Cy = 2/tau + 8/hy**2

    '''списики для рассчета alpha_i, beta_i'''
    alpha_x = [0 for i in range(N)]
    beta_x = [0 for i in range(N)]
    alpha_y = [0 for i in range(M)]
    beta_y = [0 for i in range(M)]

    '''прогонка'''
    for j in range(J):

        #переход на слой j+1/2
        for m in range(1, M):

            alpha_x[0] = 1
            beta_x[0] = 0

            #прямой ход прогонки
            for n in range(1, N):
                F = Fy(w[n][m+1][j], w[n][m][j], w[n][m-1][j], n, j)
                beta_x[n] = (F+Ax*beta_x[n-1]) / (Cx-Ax*alpha_x[n-1])
                alpha_x[n] = Bx/(Cx-Ax*alpha_x[n-1])

            #обратный ход прогонки
            w05[N][m] = -beta_x[N-1]/(-1+alpha_x[N-1])  #ручной расчет последнего значения
            for n in range(N-1, -1, -1):
                w05[n][m] = alpha_x[n]*w05[n+1][m]+beta_x[n]

        #ГУ
        for m in range(M+1):
            w[N][m][j+1] = w[N-1][m][j+1]
            w[0][m][j+1] = w[1][m][j+1]

        #ГУ
        for n in range(N+1):
            w[n][M][j+1] = w[n][M-1][j+1]
            w[n][0][j+1] = w[n][1][j+1]

        #переход на j+1-ый слой
        for n in range(1, N):

            alpha_y[0] = 1
            beta_y[0] = 0

            #прямой ход прогонки
            for m in range(1, M):
                F = Fx(w05[n+1][m], w05[n][m], w05[n-1][m], n, j)
                beta_y[m] = (F+Ay*beta_y[m-1])/(Cy-Ay*alpha_y[m-1])
                alpha_y[m] = By/(Cy-Ay*alpha_y[m-1])

            #обратный ход прогонки
            w[n][M][j+1] = -beta_y[M-1] / (-1 + alpha_y[M-1]) #ручной расчет последнего значения
            for m in range(M-1, -1, -1):
                w[n][m][j+1] = alpha_y[m]*w[n][m+1][j+1]+beta_y[m]

        #ГУ
        for m in range(M+1):
            w[N][m][j+1] = w[N-1][m][j+1]
            w[0][m][j+1] = w[1][m][j+1]

        #ГУ
        for n in range(N+1):
            w[n][M][j+1] = w[n][M-1][j+1]
            w[n][0][j+1] = w[n][1][j+1]

    #строим график для t = t_max
    BuildPlot(X, Y, w[:, :, J])


def Analitic(x, y, t):
    '''аналитическое решение (формула)'''
    return (1/17) * np.cos(x) * (4*np.sin(t) - np.cos(t) + 18 * np.exp(-4*t))


def Analit():
    '''
    аналитическое решение (заполнение)
    '''
    u = np.zeros((N+1, M+1)) #массив для записи решения

    '''заполнение массива'''
    for i in range(len(X)):
        for j in range(len(Y)):
            u[i][j] = Analitic(X[i], Y[j], t_max)

    # строим график для t = t_max
    BuildPlot(X, Y, u)



def main():
    '''запускаем решение'''
    Analit()
    NomSol()
    

if __name__ == '__main__':
    main()