3_VP_1_crank-nicolson

import pylab
import numpy
from matplotlib import mlab
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

#точное решение
def ft(t,x):
    return - x**4 + x + t*x + t**2 - t*numpy.exp(x)
#начальное условие (x,0)
def fn(x):
    return - x**4 + x
#левое граничное условие (0,t)
def fl(t):
    return t**2 - t
#правое граничное условие  (1,t)
def fr(t):
    return t + t**2 - t*numpy.exp(1);

#функция f
def f(t,x):
    return x + 2*t - numpy.exp(x) - a*(-12*x**2-t*numpy.exp(x))
#используется в вычислении прогоночных коэффициентов  beta
def e(j, k):
    return U[j][k] + c*(U[j][k+1] - 2*U[j][k] + U[j][k-1]) +  f(t*(j+1/2), h*k) * t

#--------------main_part
t_knots = 11
x_knots = 11
a = 0.019

t = 1/(t_knots-1) #шаг по t
h = 1/(x_knots-1) #шаг по х


#начальные условия
U = [[0]*x_knots for i in range(t_knots)]
for k in range (x_knots):
   U[0][k]=fn(k*h)
#метод
   #значения для вычисления прогоночных коэффициентов
c = (a*t)/(2*h**2)
A = -c
B = 1 + 2*c

for j in range(1, t_knots):
    #alpha - коэффициенты
    alpha = []
    alpha.append(0)
    #beta - коэффициенты
    beta = []
    beta.append(fl(j*t))
    for k in range(1, x_knots-1):
        alpha.append(-A/(B+A*(alpha[k-1])))
        beta.append((e(j - 1, k) - A*beta[k - 1])/(B + A*alpha[k - 1]))
    #правое граничное условие
    U[j][x_knots-1]=fr(j*t)

    #обратный ход
    for k in range(x_knots - 2, -1, -1):
        U[j][k] = alpha[k]*U[j][k + 1] + beta[k]


print ('матрица приближенного решения :')
for j in range(t_knots):
    for k in range (x_knots):
        print("%.4f" %U[j][k] ,end = ' ')
    print()


#вычисление погрешности
max = 0
for j in range(t_knots):
    for k in range(x_knots):
        if max < numpy.fabs(U[j][k] - ft(j*t, k*h)):
            max = numpy.fabs(U[j][k] - ft(j*t, k*h))
print()
print("h^2 + t^2:", h**2 + t**2, " максимальная погрешность: ", max)

#------------рисование------------
x = 0.5
a = 0
b = 1
#рисование приближенного решения
tlist = mlab.frange(a, b, t)
zlist = [U[t][round(1/h*x)] for t in range (t_knots)]
pylab.plot(tlist, zlist, color = "red",label = 'Приближенное решение')

#точное решение
tlist = mlab.frange(a,b,t) #массив значений t
flist = [ft(t,x) for t in tlist]  #х=0.5, t идет от 0 до 1 по 0.1
print('Точное решение для х =', x, 'и для значений t от',a ,'до',b,'с шагом', t , ':')
for i in range(0,t_knots):
    print(i/(t_knots-1) , " : ", "%.4f" % flist[i])
pylab.plot(tlist,flist, label = 'Точное решение')
pylab.legend()
pylab.xlabel('t')
pylab.show()

#точное решение 3Д
x = numpy.arange(a, b + h, h)
y = numpy.arange(a, b + t, t)
xgrid, ygrid = numpy.meshgrid(x, y) #двумерная матрица-сетка
zgrid = ft(ygrid, xgrid) #значения в узлах сетки
fig = pylab.figure()
axes = Axes3D(fig)
axes.plot_surface(xgrid, ygrid, zgrid,  cmap=cm.jet)
pylab.xlabel('x')
pylab.ylabel('t')
pylab.show()