numlab4

import numpy as np, math as math, matplotlib.pyplot as plt, copy

def plotErrorsByT(time, errors):
    errors.reverse()
    plt.figure(figsize=(12, 8))
    plt.xlabel('time')
    plt.ylabel('error')
    plt.title('Errors by time')
    plt.plot(time, errors, color = 'b')

class Data:

    def __init__(self, ia, ib, iff, iax0, ibx0, ifx0, iaxl, ibxl, ifxl, iay0, iby0, ify0, iayl, ibyl, ifyl, ifpsi, iintervalOnX, iintervalOnY, inx, iny, itime, irealFunction):
        self.a = ia
        self.b = ib
        self.f = iff

        self.ax0 = iax0
        self.bx0 = ibx0
        self.fx0 = ifx0

        self.axl = iaxl
        self.bxl = ibxl
        self.fxl = ifxl

        self.ay0 = iay0
        self.by0 = iby0
        self.fy0 = ify0

        self.ayl = iayl
        self.byl = ibyl
        self.fyl = ifyl

        self.fpsi = ifpsi

        self.intervalOnX = iintervalOnX
        self.intervalOnY = iintervalOnY

        self.nx = inx
        self.ny = iny

        self.time = itime

        self.realFunction = irealFunction

class Solution:

    @classmethod
    def process(cls, data : Data, method):
        stepX = (data.intervalOnX[1] - data.intervalOnX[0]) / (data.nx - 1)
        stepY = (data.intervalOnY[1] - data.intervalOnY[0]) / (data.ny - 1)
        tau = cls.__sigma / (data.a / stepX ** 2 + data.b / stepY ** 2)
        time_cnt = int(data.time / tau)

        x = np.array([stepX * i for i in range(data.nx)])
        y = np.array([stepY * i for i in range(data.ny)])
        t = np.array([tau * k for k in range(time_cnt)])

        if (method == "Alternative directions"):
            resU = cls.alternativeDirectionsMethod(data, time_cnt, stepX, stepY, tau)

        elif (method == "Fractional steps"):
            resU = cls.fractionalStepsMethod(data, time_cnt, stepX, stepY, tau)

        return x, y, t, resU

    @classmethod
    def alternativeDirectionsMethod(cls, data : Data, time_cnt, stepX, stepY, tau):
        u = np.zeros((time_cnt, data.nx, data.ny))

        a_1 = np.array([0.0] + [-data.a * tau * stepY ** 2 for _ in range(1, data.nx - 1)] + [0.0])
        b_1 = np.array([0.0] + [2 * stepX ** 2 * stepY ** 2 + 2 * data.a * stepY ** 2 * tau for _ in range(1, data.nx - 1)] + [0.0])
        c_1 = np.array([0.0] + [-data.a * tau * stepY ** 2 for _ in range(1, data.nx - 1)] + [0.0])

        b_1[0] = stepX * data.bx0 - data.ax0
        c_1[0] = data.ax0
        a_1[-1] = -data.axl
        b_1[-1] = data.axl + stepX * data.bxl

        a_2 = np.array([0.0] + [-data.b * tau * stepX ** 2 for _ in range(1, data.ny - 1)] + [0.0])
        b_2 = np.array([0.0] + [2 * stepX ** 2 * stepY ** 2 + 2 * data.b * stepX ** 2 * tau for _ in range(1, data.ny - 1)] + [0.0])
        c_2 = np.array([0.0] + [-data.b * tau * stepX ** 2 for _ in range(1, data.ny - 1)] + [0.0])

        b_2[0] = stepY * data.by0 - data.ay0
        c_2[0] = data.ay0
        a_2[-1] = -data.ayl
        b_2[-1] = data.ayl + stepY * data.byl

        for i in range(data.nx):
            for j in range(data.ny):
                u[0][i][j] = data.fpsi(stepX * i, stepY * j)

        for k in range(1, time_cnt):
            u_1 = np.zeros((data.ny, data.nx))
            d_1 = np.zeros(data.nx)
            u_2 = np.zeros((data.nx, data.ny))
            d_2 = np.zeros(data.ny)
            for j in range (1, data.ny - 1):
                for i in range(1, data.nx - 1):
                    d_1[i] = (u[k - 1][i][j - 1] - 2 * u[k - 1][i][j] + u[k - 1][i][j + 1]) * data.b * stepX ** 2 * tau + \
                              2 * stepX ** 2 * stepY ** 2 * u[k - 1][i][j] + \
                              stepX ** 2 * stepY ** 2 * tau * data.f(stepX * i, stepY * j, tau * (k - 0.5))
                d_1[0] = stepX * data.fx0(stepY * j ,tau * (k - 0.5))
                d_1[-1] = stepX * data.fxl(stepY * j ,tau * (k - 0.5))

                u_1_j = cls.__tma(a_1, b_1, c_1, d_1, data.nx)
                u_1[j] = copy.deepcopy(u_1_j)
            for i in range(data.nx):
                u_1[0][i] = (stepY * data.fy0(i * stepX, tau * (k - 0.5)) - data.ay0 * u_1[1][i]) / (stepY * data.by0 - data.ay0)
                u_1[-1][i] = (stepY * data.fyl(i * stepX, tau * (k - 0.5)) + data.ayl * u_1[-2][i]) / (stepY * data.byl + data.ayl)
            u_1 = u_1.transpose()

            for i in range(1, data.nx - 1):
                for j in range(1, data.ny - 1):
                    d_2[j] = (u_1[i - 1][j] - 2 * u_1[i][j] + u_1[i + 1][j]) * data.a * stepY ** 2 * tau + \
                              2 * stepX ** 2 * stepY ** 2 * u_1[i][j] + \
                              stepX ** 2 * stepY ** 2 * tau * data.f(stepX * i, stepY * j, tau * k)
                d_2[0] = stepY * data.fy0(stepX * i ,tau * k)
                d_2[-1] = stepY * data.fyl(stepX * i ,tau * k)

                u_2_i = cls.__tma(a_2, b_2, c_2, d_2, data.ny)
                u_2[i] = copy.deepcopy(u_2_i)
            for j in range(data.ny):
                u_2[0][j] = (stepX * data.fx0(j * stepY, tau * k) - data.ax0 * u_2[1][j]) / (stepX * data.bx0 - data.ax0)
                u_2[-1][j] = (stepX * data.fxl(j * stepY, tau * k) + data.axl * u_2[-2][j]) / (stepX * data.bxl + data.axl)
            u[k] = copy.deepcopy(u_2)
        return u

    @classmethod
    def fractionalStepsMethod(cls, data : Data, time_cnt, stepX, stepY, tau):
        u = np.zeros((time_cnt, data.nx, data.ny))

        a_1 = np.array([0.0] + [-2 * data.a * tau for _ in range(1, data.nx - 1)] + [0.0])
        b_1 = np.array([0.0] + [2 * stepX ** 2 + 4 * tau * data.a for _ in range(1, data.nx - 1)] + [0.0])
        c_1 = np.array([0.0] + [-2 * data.a * tau for _ in range(1, data.nx - 1)] + [0.0])

        b_1[0] = stepX * data.bx0 - data.ax0
        c_1[0] = data.ax0
        a_1[-1] = -data.axl
        b_1[-1] = data.axl + stepX * data.bxl

        a_2 = np.array([0.0] + [-2 * data.b * tau for _ in range(1, data.ny - 1)] + [0.0])
        b_2 = np.array([0.0] + [2 * stepY ** 2 + 4 * tau * data.b for _ in range(1, data.ny - 1)] + [0.0])
        c_2 = np.array([0.0] + [-2 * data.b * tau for _ in range(1, data.ny - 1)] + [0.0])

        b_2[0] = stepY * data.by0 - data.ay0
        c_2[0] = data.ay0
        a_2[-1] = -data.ayl
        b_2[-1] = data.ayl + stepY * data.byl

        for i in range(data.nx):
            for j in range(data.ny):
                u[0][i][j] = data.fpsi(stepX * i, stepY * j)

        for k in range(1, time_cnt):
            u_1 = np.zeros((data.ny, data.nx))
            d_1 = np.zeros(data.nx)
            u_2 = np.zeros((data.nx, data.ny))
            d_2 = np.zeros(data.ny)
            for j in range (1, data.ny - 1):
                for i in range(1, data.nx - 1):
                    d_1[i] = 2 * stepX ** 2 * u[k - 1][i][j] + tau * stepX ** 2 * data.f(stepX * i, stepY * j, tau * (k - 0.5))
                d_1[0] = stepX * data.fx0(stepY * j ,tau * (k - 0.5))
                d_1[-1] = stepX * data.fxl(stepY * j ,tau * (k - 0.5))

                u_1_j = cls.__tma(a_1, b_1, c_1, d_1, data.nx)
                u_1[j] = copy.deepcopy(u_1_j)
            for i in range(data.nx):
                u_1[0][i] = (stepY * data.fy0(i * stepX, tau * (k - 0.5)) - data.ay0 * u_1[1][i]) / (stepY * data.by0 - data.ay0)
                u_1[-1][i] = (stepY * data.fyl(i * stepX, tau * (k - 0.5)) + data.ayl * u_1[-2][i]) / (stepY * data.byl + data.ayl)
            u_1 = u_1.transpose()

            for i in range(1, data.nx - 1):
                for j in range(1, data.ny - 1):
                    d_2[j] = 2 * stepY ** 2 * u_1[i][j] + tau * stepY ** 2 * data.f(stepX * i, stepY * j, tau * k)
                d_2[0] = stepY * data.fy0(stepX * i ,tau * k)
                d_2[-1] = stepY * data.fyl(stepX * i ,tau * k)

                u_2_i = cls.__tma(a_2, b_2, c_2, d_2, data.ny)
                u_2[i] = copy.deepcopy(u_2_i)
            for j in range(data.ny):
                u_2[0][j] = (stepX * data.fx0(j * stepY, tau * k) - data.ax0 * u_2[1][j]) / (stepX * data.bx0 - data.ax0)
                u_2[-1][j] = (stepX * data.fxl(j * stepY, tau * k) + data.axl * u_2[-2][j]) / (stepX * data.bxl + data.axl)
            u[k] = copy.deepcopy(u_2)
        return u

    @classmethod
    def realFunctionSolution(cls, data : Data):
        stepX = (data.intervalOnX[1] - data.intervalOnX[0]) / (data.nx - 1)
        stepY = (data.intervalOnY[1] - data.intervalOnY[0]) / (data.ny - 1)
        tau = cls.__sigma / (data.a / stepX ** 2 + data.b / stepY ** 2)
        time_cnt = int(data.time / tau)

        ans = np.zeros((time_cnt, data.nx, data.ny))
        for k in range(time_cnt):
            for i in range(data.nx):
                for j in range(data.ny):
                    ans[k][i][j] = data.realFunction(stepX * i, stepY * j, tau * k)
        return ans

    @classmethod
    def __tma(cls, a, b, c, d, n):
        A = np.zeros(n)
        B = np.zeros(n)
        x = np.zeros(n)
        A[0] = -c[0] / b[0]
        B[0] = d[0] / b[0]

        for j in range(1, n):
            A[j] = -c[j] / (b[j] + a[j] * A[j - 1])
            B[j] = (d[j] - a[j] * B[j - 1]) / (b[j] + a[j] * A[j - 1])
        x[-1] = B[-1]
        for j in range(n - 2, -1, -1):
            x[j] = A[j] * x[j + 1] + B[j]
        return x

    __sigma = 0.4

def plotErrors(data : Data, method, start, where, amount):
    errors = []
    dataTmp = data
    if (where == 'x'):
        dataTmp.nx = start
    elif (where == 'y'):
        dataTmp.ny = start

    if (method == "Alternative directions"):
        if (where == 'x'):
            for i in range(amount):
                dataTmp.nx += 1
                realSolution = Solution.realFunctionSolution(data)
                x, y, t, utmp = Solution.process(dataTmp, method)
                errorTmp = max([max([max(np.absolute(utmp[k][i] - realSolution[k][i])) for i in range(len(x))]) for k in range(len(t))])
                errors.append(errorTmp)
        elif (where == 'y'):
            for i in range(amount):
                dataTmp.ny += 1
                realSolution = Solution.realFunctionSolution(data)
                x, y, t, utmp = Solution.process(dataTmp, method)
                errorTmp = max([max([max(np.absolute(utmp[k][i] - realSolution[k][i])) for i in range(len(x))]) for k in range(len(t))])
                errors.append(errorTmp)

    elif (method == "Fractional steps"):
        if (where == 'x'):
            for i in range(amount):
                dataTmp.nx += 1
                realSolution = Solution.realFunctionSolution(data)
                x, y, t, utmp = Solution.process(data, method)
                errorTmp = max([max([max(np.absolute(utmp[k][i] - realSolution[k][i])) for i in range(len(x))]) for k in range(len(t))])
                errors.append(errorTmp)
        elif (where == 'y'):
            for i in range(amount):
                dataTmp.ny += 1
                realSolution = Solution.realFunctionSolution(data)
                x, y, t, utmp = Solution.process(data, method)
                errorTmp = max([max([max(np.absolute(utmp[k][i] - realSolution[k][i])) for i in range(len(x))]) for k in range(len(t))])
                errors.append(errorTmp)

    nums = [i for i in range(start, start + amount)]
    if (where == 'x'):
        errors.reverse()

    plt.figure(figsize=(12, 8))
    plt.plot(nums, errors, color = 'r')
    plt.title('Errors with different ' + where)

a = 1
b = 1

f = lambda x, y, t: -x * y * np.sin(t)

alphax0 = 0
alphaxl = -1
alphay0 = 0
alphayl = -1

betax0 = 1
betaxl = 1
betay0 = 1
betayl = 1

fx0 = lambda y, t: 0.0
fxl = lambda y, t: 0.0
fy0 = lambda x, t: 0.0
fyl = lambda x, t: 0.0

fpsi = lambda x, y: x * y

intervalOnX = (0.0, 1.0)
intervalOnY = (0.0, 1.0)

realFunction = lambda x, y, t: x * y * np.cos(t)

nx = 20
ny = 20
time = 3
data = Data(a, b, f, alphax0, betax0, fx0, alphaxl, betaxl, fxl, alphay0, betay0, fy0, alphayl, betayl, fyl, fpsi, intervalOnX, intervalOnY, nx, ny, time, realFunction)

realSolution = Solution.realFunctionSolution(data)

method = "Alternative directions"
nx1 = 20
ny1 = 20
time1 = 3
data1 = Data(a, b, f, alphax0, betax0, fx0, alphaxl, betaxl, fxl, alphay0, betay0, fy0, alphayl, betayl, fyl, fpsi, intervalOnX, intervalOnY, nx1, ny1, time1, realFunction)
xAD, yAD, tAD, uAD = Solution.process(data1, method)

errorAD = max(np.array([max([max(np.absolute(uAD[k][i] - realSolution[k][i])) for i in range(len(xAD))]) for k in range(len(tAD))]))
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.set_title("График точного и численного решения задачи")
ax.plot_wireframe(X, Y, U_analytic, color = "red", label = "Точное решение")
ax.plot_wireframe(xAD, yAD, uAD, color = "blue", label = "Численное решение")
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("U")
ax.legend()
plt.show()

method = "Fractional steps"
nx2 = 20
ny2 = 20
time2 = 3
data2 = Data(a, b, f, alphax0, betax0, fx0, alphaxl, betaxl, fxl, alphay0, betay0, fy0, alphayl, betayl, fyl, fpsi, intervalOnX, intervalOnY, nx2, ny2, time2, realFunction)
xFS, yFS, tFS, uFS = Solution.process(data2, method)

errorFS = max(np.array([max([max(np.absolute(uFS[k][i] - realSolution[k][i])) for i in range(len(xFS))]) for k in range(len(tFS))]))

plt.figure(figsize=(14, 14))
t = np.linspace(0, len(tAD) - 1, num=5, endpoint=True, dtype=int)

ax = plt.axes(projection='3d')

X, Y = np.meshgrid(xAD, yAD)

colors = ['b', 'orange','green', 'r', 'purple', 'yellow']
c = 0
for i in t:
    ax.plot_wireframe(X, Y, uAD[i], color = colors[c])
    c += 1

legends = ["t=" + str(tAD[i]) for i in t]
ax.legend(legends)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("U(x, y, t)")

plt.title('Решение при различных t')
plt.show()

errorsAD = [max([max(np.absolute(uAD[k][i] - realSolution[k][i])) for i in range(len(xAD))]) for k in range(len(tAD))]
plotErrorsByT(tAD, errorsAD)
plt.show()

plt.figure(figsize=(14, 14))
t = np.linspace(0, len(tFS) - 1, num=5, endpoint=True, dtype=int)


ax = plt.axes(projection='3d')

X, Y = np.meshgrid(xFS, yFS)

colors = [ 'b', 'orange','green', 'r', 'purple']
c = 0
for i in t:
    ax.plot_wireframe(X, Y, uFS[i], color = colors[c])
    c += 1

legends = ["t=" + str(tFS[i]) for i in t]
ax.legend(legends)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("U(x, y, t)")

plt.show()

errorsFS = [max([max(np.absolute(uAD[k][i] - realSolution[k][i])) for i in range(len(xAD))]) for k in range(len(tAD))]
plotErrorsByT(tFS, errorsFS)
plt.show()

id = 400

plt.figure(figsize=(14, 14))

ax = plt.axes(projection='3d')

X, Y = np.meshgrid(xAD, yAD)


ax.plot_wireframe(X, Y, realSolution[id], color = 'b', label='Real function')
ax.plot_wireframe(X, Y, uAD[id], color = 'g', label='Alternative directions')
ax.plot_wireframe(X, Y, uFS[id], color = 'r', label ='Fractional steps')
plt.legend()
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("U(x, y, t)")

plt.title("Решения методов в t = " + str(tAD[id]))
def norm(cur_u, prev_u):
    max = 0
    for i in range(cur_u.shape[0]):
        for j in range(cur_u.shape[1]):
            if abs(cur_u[i, j] - prev_u[i, j]) > max:
                max = abs(cur_u[i, j] - prev_u[i, j])

    return max
er_l = []
h = []
method = "Alternative directions"
for i in range(20, 5, -1):
    nx1 = i
    ny1 = i
    h.append(1 / (nx1 - 1))
    time1 = 3
    data1 = Data(a, b, f, alphax0, betax0, fx0, alphaxl, betaxl, fxl, alphay0, betay0, fy0, alphayl, betayl, fyl, fpsi, intervalOnX, intervalOnY, nx1, ny1, time1, realFunction)
    xAD, yAD, tAD, uAD = Solution.process(data1, method)

    realSolution = Solution.realFunctionSolution(data1)

    er_l.append(norm(realSolution[10], uAD[10]))
plt.title('Ошибка метода переменных направлений')
plt.xlabel('h')
plt.ylabel('error')
plt.plot(h, er_l)
plt.show()