numlab3

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

class Equation:
    def __init__(self, params):
        self.bx = params['bx']
        self.by = params['by']
        self.c = params['c']

        self.alpha_x0 = params['alpha_x0']
        self.beta_x0 = params['beta_x0']
        self.f_phi_x0 = params['f_phi_x0']

        self.alpha_xl = params['alpha_xl']
        self.beta_xl = params['beta_xl']
        self.f_phi_xl = params['f_phi_xl']

        self.alpha_y0 = params['alpha_y0']
        self.beta_y0 = params['beta_y0']
        self.f_phi_y0 = params['f_phi_y0']

        self.alpha_yl = params['alpha_yl']
        self.beta_yl = params['beta_yl']
        self.f_phi_yl = params['f_phi_yl']

        self.interval_x = params['interval_x']
        self.interval_y = params['interval_y']
        self.analytical_solution = params['analytical_solution']

    def solve(self, nx, ny, method, omega=1.0, eps=1e-8, quiet=False):
        hx, hy = self.__prepare(self.interval_x[0], self.interval_x[1], nx,
                                self.interval_y[0], self.interval_y[1], ny)
        if method == 'simple_iterations':
            u = self.__simple_iterations_method(nx, ny, hx, hy, eps, quiet)
        elif method == 'seidel':
            u = self.__seidel_method(nx, ny, hx, hy, 1.0, eps, quiet)
        elif method == 'top_relaxation':
            u = self.__seidel_method(nx, ny, hx, hy, omega, eps, quiet)
        return u

    def __simple_iterations_method(self, nx, ny, hx, hy, eps, quiet):
        u_new = np.ones((nx, ny))
        u = np.ones((nx, ny))
        iter = 0
        while True:
            for i in range(1, nx - 1):
                for j in range(1, ny - 1):
                    u_new[i][j] = ((u[i + 1][j] - u[i - 1][j]) * self.bx * hx * hy ** 2 + \
                                   (u[i][j + 1] - u[i][j - 1]) * self.by * hy * hx ** 2 - \
                                   2 * hy ** 2 * (u[i - 1][j] + u[i + 1][j]) - \
                                   2 * hx ** 2 * (u[i][j - 1] + u[i][j + 1])) / \
                                  (-4 * hy ** 2 - 4 * hx ** 2 - 2 * hx ** 2 * hy ** 2 * self.c)
            for j in range(ny):
                u_new[0][j] = (hx * self.f_phi_x0(j * hy) - self.alpha_x0 * u_new[1][j]) / (
                            hx * self.beta_x0 - self.alpha_x0)
                u_new[-1][j] = (hx * self.f_phi_xl(j * hy) + self.alpha_xl * u_new[-2][j]) / (
                            hx * self.beta_xl + self.alpha_xl)
            for i in range(nx):
                u_new[i][0] = (hy * self.f_phi_y0(i * hx) - self.alpha_y0 * u_new[i][1]) / (
                            hy * self.beta_y0 - self.alpha_y0)
                u_new[i][-1] = (hy * self.f_phi_yl(i * hx) + self.alpha_yl * u_new[i][-2]) / (
                            hy * self.beta_yl + self.alpha_yl)
            eps_k = max([max(np.absolute(u_new[i] - u[i])) for i in range(nx)])
            iter += 1
            if eps_k < eps:
                if not quiet:
                    print('Итераций (метод простых итераций): ', iter)
                return u_new
            u = copy.deepcopy(u_new)

    def __seidel_method(self, nx, ny, hx, hy, omega, eps, quiet):
        u_prev = np.ones((nx, ny))
        u = np.ones((nx, ny))
        iter = 0
        while True:
            for i in range(1, nx - 1):
                for j in range(1, ny - 1):
                    u[i][j] = ((u[i + 1][j] - u[i - 1][j]) * self.bx * hx * hy ** 2 + \
                               (u[i][j + 1] - u[i][j - 1]) * self.by * hy * hx ** 2 - \
                               2 * hy ** 2 * (u[i - 1][j] + u[i + 1][j]) - \
                               2 * hx ** 2 * (u[i][j - 1] + u[i][j + 1])) / \
                              (-4 * hy ** 2 - 4 * hx ** 2 - 2 * hx ** 2 * hy ** 2 * self.c)
            for j in range(ny):
                u[0][j] = (hx * self.f_phi_x0(j * hy) - self.alpha_x0 * u[1][j]) / (hx * self.beta_x0 - self.alpha_x0)
                u[-1][j] = (hx * self.f_phi_xl(j * hy) + self.alpha_xl * u[-2][j]) / (hx * self.beta_xl + self.alpha_xl)
            for i in range(nx):
                u[i][0] = (hy * self.f_phi_y0(i * hx) - self.alpha_y0 * u[i][1]) / (hy * self.beta_y0 - self.alpha_y0)
                u[i][-1] = (hy * self.f_phi_yl(i * hx) + self.alpha_yl * u[i][-2]) / (hy * self.beta_yl + self.alpha_yl)
            for i in range(nx):
                for j in range(ny):
                    u[i][j] = omega * u[i][j] + (1.0 - omega) * u_prev[i][j]
            eps_k = max([max(np.absolute(u[i] - u_prev[i])) for i in range(nx)])
            iter += 1
            if eps_k < eps:
                if not quiet:
                    print('Итераций (метод Зейделя): ', iter)
                return u
            u_prev = copy.deepcopy(u)

    def solution(self, nx, ny, u):
        x = np.linspace(self.interval_x[0], self.interval_x[1], u.shape[1])
        y = np.linspace(self.interval_y[0], self.interval_y[1], u.shape[0])
        X, Y = np.meshgrid(x, y)
        return self.analytical_solution(Y, X)

    def __prepare(self, begin_x, end_x, nx, begin_y, end_y, ny):
        hx = (end_x - begin_x) / (nx - 1)
        hy = (end_y - begin_y) / (ny - 1)
        return hx, hy

def draw(x_bounds, t_bounds, n, max_time, res, orig):
    x = np.linspace(x_bounds[0], x_bounds[1], n)
    y = np.linspace(t_bounds[0], t_bounds[1], max_time)
    X, Y = np.meshgrid(x, y)
    ax = plt.axes(projection='3d')
    ax.plot_surface(X, Y, res, color='blue')
    ax.plot_wireframe(X, Y, orig, color='red')
    #ax.contour3D(X, Y, res, X, Y, 50, cmap='binary')
    plt.show()

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

def main():
    params = {
        'bx': -2.0,
        'by': 0,
        'c': -3.0,

        'alpha_x0': 0.0,
        'beta_x0': 1.0,
        'f_phi_x0': lambda y: np.cos(y),

        'alpha_xl': 0.0,
        'beta_xl': 1.0,
        'f_phi_xl': lambda y: 0.0,

        'alpha_y0': 0.0,
        'beta_y0': 1.0,
        'f_phi_y0': lambda x: np.exp(-x) * np.cos(x),

        'alpha_yl': 0.0,
        'beta_yl': 1.0,
        'f_phi_yl': lambda x: 0.0,

        'interval_x': (0.0, np.pi / 2),
        'interval_y': (0.0, np.pi / 2),
        'analytical_solution': lambda x, y: np.exp(-x) * np.cos(x) * np.cos(y)
    }

    equation = Equation(params)

    # ['simple_iterations', 'seidel', 'top_relaxation']
    method = 'top_relaxation'
    nx = 25
    ny = 25
    omega = 1.7

    u = equation.solve(nx=nx, ny=ny, method=method, omega=omega, eps=1e-8)
    u_orig = equation.solution(nx=nx, ny=ny, u=u)

    #draw(equation.interval_x, equation.interval_y, u_orig.shape[1], u_orig.shape[0], u_orig)
    draw(equation.interval_x, equation.interval_y, u.shape[1], u.shape[0], u, u_orig)

    plt.show()
    er_l = []
    h = []
    for i in range(25, 5, -1):
        nx = i
        ny = i
        h.append(np.pi / (2 * nx))

        u = equation.solve(nx=nx, ny=ny, method=method, omega=omega, eps=1e-8)
        u_orig = equation.solution(nx=nx, ny=ny, u=u)
        er_l.append(norm(u_orig, u))
    plt.title('Ошибка метода простых итераций')
    plt.xlabel('h')
    plt.ylabel('error')
    plt.plot(h, er_l)
    plt.show()

if __name__ == '__main__':
    main()