from numpy import *
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

h = 0.1 # шаг
tau = 0.004 # шаг по времени
begin_t = 0 # начало по времени
end_t = 1 # конец по времени
begin_x = 0 # начало по х
exp = 1 # экспонента

H, T = meshgrid(arange(begin_x, exp, h), arange(begin_t, end_t, tau)) # для построения графика

t = tau
temp = 0

yavn_s = [[i + sin(pi*i) for i in arange(begin_x, exp, h)]]
size = len(yavn_s[0]) # первая строчка массива
sig = tau / (h ** 2) # сигма
grid = math.ceil((end_t - t) / tau) # задание размеров сетки
for i in range(grid):
    yavn_s.append([0] * size)
    for j in range(1, size - 1):
        yavn_s[temp + 1][j] = sig * yavn_s[temp][j - 1] + (1 - 2 * sig) * yavn_s[temp][j] + sig * yavn_s[temp][j + 1]
    yavn_s[temp + 1][0] = 0
    yavn_s[temp + 1][size - 1] = 1
    t += tau
    temp += 1

#Обнуляем переменные для следующей схемы
temp = 0
t = tau

#На каждом шаге
neyavn_s = [[i + sin(pi*i) for i in arange(begin_x, exp, h)]]
for k in range(grid):
    matrx = [[1, 0] + [0 for i in range(size - 2)]]
    b = [0] # правая часть матрицы
    l = 1
    for i in range(size - 2):
        matrx += [[0 for k in range(size)]]
        matrx[l][i] = -sig
        matrx[l][i + 1] = 1 + 2 * sig
        matrx[l][i + 2] = -sig
        b += [neyavn_s[temp][i + 1]]
        l += 1
    matrx += [[0 for i in range(size - 2)] + [0, 1]]
    b += [1]
    neyavn_s.append(list(linalg.solve(matrx, b)))
    temp += 1
    t += tau

#Обнуляем переменные для следующей схемы
temp = 0
t = tau


krank = [[i + sin(pi*i) for i in arange(begin_x, exp, h)]]
for k in range(grid):
    matrx = [[1, 0] + [0 for i in range(size - 2)]]
    b = [0]
    l = 1
    for i in range(size - 2):
        matrx += [[0 for k in range(size)]]
        matrx[l][i] = -sig / 2
        matrx[l][i + 1] = 1 + sig
        matrx[l][i + 2] = -sig / 2
        b += [krank[temp][i + 1]]
        l += 1
    matrx += [[0 for i in range(size - 2)] + [0, 1]]
    b += [1]
    krank.append(list(linalg.solve(matrx, b)))
    temp += 1
    t += tau



#Построение графиков

figure1 = plt.figure("Явная схема")
axes1 = Axes3D(figure1)
axes1.view_init(30, 60)
axes1.plot_surface(H, T, array(yavn_s), cmap=cm.jet)
plt.show()

figure2 = plt.figure("Неявная схема")
axes2 = Axes3D(figure2)
axes2.view_init(30, 60)
axes2.plot_surface(H, T, array(neyavn_s), cmap=cm.jet)
plt.show()

figure3 = plt.figure("Схема Кранка-Николсона")
axes3 = Axes3D(figure3)
axes3.view_init(30, 60)
axes3.plot_surface(H, T, array(krank), cmap=cm.jet)
plt.show()

real = H + exp(-pi**2*T)*sin(pi*H)
figure4 = plt.figure("Настоящая")
axes4 = Axes3D(figure1)
axes4.view_init(30, 60)
axes4.plot_surface(H, T, array(real), cmap=cm.jet)
plt.show()



#Вывод погрешностей

print("Погрешность явной схемы: ")
for i in range(len(real)):
    for j in range(len(real[i])):
        print('{0:.4f}'.format(fabs(real[i][j] - yavn_s[i][j])), end=' ')
    print()

print("Погрешность неявной схемы: ")
for i in range(len(real)):
    for j in range(len(real[i])):
        print('{0:.4f}'.format(fabs(real[i][j] - neyavn_s[i][j])), end=' ')
    print()

print("Погрешность схемы Кранка-Николсона: ")
for i in range(len(real)):
    for j in range(len(real[i])):
        print('{0:.4f}'.format(fabs(real[i][j] - krank[i][j])), end=' ')
    print()