Сравнение итерационного метода Якоби и прямого метода Гаусса#

Решаемая система#

Будем решать систему

\[ Ax=f \]

где

\[\begin{split} \begin{array}{cl} a_{i i}=2, & a_{i, i+1}=-1-\alpha, \quad a_{i, i-1}=-1+\alpha \\ f_1=1-\alpha, & f_i=0, \quad i=2,3, \ldots, n-1, \quad f_n=1+\alpha \end{array} \end{split}\]

Составим функцию, возвращающую нам матрицу A и столбец f по заданному n и alpha.

import numpy as np

def system(n: int, alpha: float):

    A = np.zeros((n, n))
    f = np.zeros((n, 1))


    for i in range(n):
        A[i,i] = 2.
        if i > 0:
            A[i,i-1] = -1 + alpha
        if i < n - 1:
            A[i,i+1] = -1 - alpha

    f = np.zeros(n)
    f[0] = 1 - alpha
    f[n-1] = 1 + alpha

    return (A, f)

system(4, 1)
(array([[ 2., -2.,  0.,  0.],
        [ 0.,  2., -2.,  0.],
        [ 0.,  0.,  2., -2.],
        [ 0.,  0.,  0.,  2.]]),
 array([0., 0., 0., 2.]))

Реализуем метод решения этой системы прямым методом Гаусса#

def gauss_solve(A,f):

    x = np.copy(f)

    n = len(f)
    # Elimination phase
    for k in range(0,n-1):
        for i in range(k+1,n):
            if A[i,k] != 0.0:
                #if not null define λ
                lam = A[i,k]/A[k,k]
                #we calculate the new row of the matrix
                A[i,k+1:n] = A[i,k+1:n] - lam*A[k,k+1:n]
                #we update vector b
                x[i] = x[i] - lam*x[k]
                # backward substitution
    for k in range(n-1,-1,-1):
        x[k] = (x[k] - np.dot(A[k,k+1:n],x[k+1:n]))/A[k,k]

    return x

A, f = system(4, 1)

gauss_solve(A, f)
array([1., 1., 1., 1.])
# Магия питона - мотайте на ус
gauss_solve(*system(4, 1))
array([1., 1., 1., 1.])
A @ gauss_solve(A, f) - f
array([0., 0., 0., 0.])

Реализуем метод Якоби#

\[ x_i^{(k+1)}=\frac{1}{a_{i i}}\left(f_i-\sum_{j \neq i} a_{i j} x_j^{(k)}\right), \quad i=1,2, \ldots, n \]
def jacobi_solve(A, f, eps = 1e-6, MAX_ITER = 100000, x_prev = None):

    n = f.shape[0]

    if x_prev == None:
        x_prev = np.ones(f.shape)*2 # нулевое начальное приближение - плохо!
    x_curr = np.copy(x_prev)

    k = 0

    for k in range(1, MAX_ITER + 1):

        x_prev[:] = x_curr[:] # Это важно, да

        for i in range(n):
            x_curr[i] = 1./A[i, i]*(f[i] - A[i, 0:i] @ x_prev[0:i] - A[i, i+1:] @ x_prev[i+1:])

        if np.sum((x_curr - x_prev)**2) < eps**2:
            break


    if k == MAX_ITER:
        print("Didn't converge!Q!11!!1")

    return x_curr
jacobi_solve(*system(4, 1))
array([1., 1., 1., 1.])

Сравнение#

Необходимо построить 3 графика для трёх разных alpha (маленького, небольшого, очень большого), на каждом мз которых будут кривые времени выполнения для каждого алгоритма в зависимости от n.

Сначала реализуем функцию, которая по заданному alpha, диапазону n и заданной функции будет выдавать нам список с временами выволнения

import time
def series_timed(alpha, ns, func):

    Ts = []

    for n in ns:

        A, f = system(n, alpha)

        start = time.time()
        _ = func(A, f)
        seconds = time.time() - start

        Ts.append(seconds)

    return Ts
series_timed(0.5, [1, 2, 3, 200], jacobi_solve)
[0.002310037612915039,
 0.0005118846893310547,
 0.0013003349304199219,
 0.3280794620513916]

Теперь построим графики сравнения

%%time
import matplotlib.pyplot as plt

alphas = [0.2, 0.5, 1.]
ns = np.arange(10, 400, 20)
funcs = [gauss_solve, jacobi_solve]

fig, axs = plt.subplots(len(alphas))
fig.set_size_inches(10, 10)
fig.suptitle('Time (n)')

for i, alpha in enumerate(alphas):
    print(i)

    for func in funcs:
        axs[i].semilogy(ns, series_timed(alpha, ns, func), label = func.__name__)

    axs[i].title.set_text(f"alpha = {alpha}")
    axs[i].legend()
0
1
2
CPU times: user 31.4 s, sys: 250 ms, total: 31.6 s
Wall time: 33.9 s
../../_images/e7bf1ae6fa0bf0e5781e55fbeec48530941d21b0e3674b7261055ff2c53f8408.png

Ускорение кода python с помощью numba#

Даже используя numpy питоновский код на порядки медленее c или c++. Есть ли способ ускорить код? Оказывается - библиотека numba.

Эта библиотека позволяет с минимальными усилиями достингуть скорости работы c++ с помощью just in time компиляции питоновского кода в сишный. Поддерживает большинство функций numpy, но не все. Также может многое распараллелить за нас для ещё бОльшего ускорения! Однако многие фишки питона не поддерживаются или работают очень медленно (добавление элемента в список и т.д.).

Для пользования совет - стараться писать в сишном стиле без питоновских наворотов.

#!pip install numba
from numba import njit, objmode, prange

@njit()
def system_numba(n: int, alpha: float):

    A = np.zeros((n, n))
    f = np.zeros((n, 1))


    for i in range(n):
        A[i,i] = 2.
        if i > 0:
            A[i,i-1] = -1 + alpha
        if i < n - 1:
            A[i,i+1] = -1 - alpha

    f = np.zeros(n)
    f[0] = 1 - alpha
    f[n-1] = 1 + alpha

    return (A, f)

@njit()
def gauss_solve_numba(A,f):

    x = np.copy(f)

    n = len(f)
    # Elimination phase
    for k in range(0,n-1):
        for i in range(k+1,n):
            if A[i,k] != 0.0:
                #if not null define λ
                lam = A[i,k]/A[k,k]
                #we calculate the new row of the matrix
                A[i,k+1:n] = A[i,k+1:n] - lam*A[k,k+1:n]
                #we update vector b
                x[i] = x[i] - lam*x[k]
                # backward substitution
    for k in range(n-1,-1,-1):
        x[k] = (x[k] - np.dot(A[k,k+1:n],x[k+1:n]))/A[k,k]
    return x

@njit()
def jacobi_solve_numba(A, f, eps = 1e-6, MAX_ITER = 100000, x_prev = None):

    n = f.shape[0]
    if x_prev == None:
        x_prev = np.ones(f.shape)*2 # нулевое начальное приближение - плохо!
    x_curr = x_prev*1
    k = 0
    for k in range(1, MAX_ITER + 1):
        x_prev[:] = x_curr[:] # Это важно, да

        for i in range(n):
            x_curr[i] = 1./A[i, i]*(f[i] - A[i, 0:i] @ x_prev[0:i] - A[i, i+1:] @ x_prev[i+1:])
        if np.sum((x_curr - x_prev)**2) < eps**2:
            break
    if k == MAX_ITER:
        print("Didn't converge!Q!11!!1")
    return x_curr

@njit(parallel=True)
def series_timed_numba(alpha, ns, func):

    Ts = np.zeros(ns.shape)

    for i in prange(len(ns)): # только числовые циклы, prange для распараллеливания НЕЗАВИСИМЫХ циклов
                              # (т.е. шаг цикла не должен опираться на вычисления с предыдущих шагов)

        A, f = system_numba(ns[i], alpha)


        with objmode(start='f8'):
            start = time.perf_counter()

        _ = func(A, f)

        with objmode(seconds='f8'):
            seconds = time.perf_counter() - start

        Ts[i] = seconds # Убрал append

    return Ts
%%time

alphas = [0.2, 0.5, 1.]
ns = np.arange(10, 400, 20)
funcs = [gauss_solve_numba, jacobi_solve_numba]

fig, axs = plt.subplots(len(alphas))
fig.set_size_inches(10, 10)
fig.suptitle('Time (n)')

for i, alpha in enumerate(alphas):
    print(i)

    for func in funcs:
        axs[i].semilogy(ns, series_timed_numba(alpha, ns, func), label = func.__name__)

    axs[i].title.set_text(f"alpha = {alpha}")
    axs[i].legend()
0
1
2
CPU times: user 1.73 s, sys: 0 ns, total: 1.73 s
Wall time: 346 ms
../../_images/5a0856eeb494a451578dc0d8caca71c76a3447cbd0dca7085501e98f6e091497.png
%%time

alphas = [0.2, 0.5, 1.]
ns = np.arange(10, 400, 20)
funcs = [gauss_solve_numba, jacobi_solve_numba]

fig, axs = plt.subplots(len(alphas))
fig.set_size_inches(10, 10)
fig.suptitle('Time (n)')

for i, alpha in enumerate(alphas):
    print(i)

    for func in funcs:
        axs[i].semilogy(ns, series_timed_numba(alpha, ns, func), label = func.__name__)

    axs[i].title.set_text(f"alpha = {alpha}")
    axs[i].legend()
0
1
2
CPU times: user 2.45 s, sys: 15.6 ms, total: 2.47 s
Wall time: 386 ms
../../_images/f5e317c8f5b3cd9a985b13f8f1aa943abba59ba5a618ac78e1ffee0651a93c99.png

Видим ускорение кода в ~70 раз путём изменения пары строчек! Без parallel = True и prange ускорение было бы в 26 раз со второго запуска - т.е. распараллеливание ускорило код ещё в 3 раза.

Стоит помнить, что при первои вызове функция может работать даже дольше, чем обычная питоновская, потому что происходит компиляция. Все последующие вызовы уже будут ускорены.

Попробуем рассмотреть более большие диапазоны n.

%%time

alphas = [0.2, 0.5, 1.]
ns = np.arange(100, 2000, 200)
funcs = [gauss_solve_numba, jacobi_solve_numba]

fig, axs = plt.subplots(len(alphas))
fig.set_size_inches(10, 10)
fig.suptitle('Time (n)')

for i, alpha in enumerate(alphas):
    print(i)

    for func in funcs:
        axs[i].semilogy(ns, series_timed_numba(alpha, ns, func), label = func.__name__)

    axs[i].title.set_text(f"alpha = {alpha}")
    axs[i].legend()
0
1
2
CPU times: user 2min 19s, sys: 125 ms, total: 2min 19s
Wall time: 39.4 s
../../_images/0e5b59f75d95d4dfc3a75bd18dc08f6f2675eb10c2d7aff6c2f0c1d99231aa0e.png

Другие примеры использования numba можете посмотреть в документ FastPython_ipynb.ipynb.