Как написать быстро для циклов с NumPy для неявных методов

Question

Numpy отлично подходит для ускорения кода путем векторизации. Однако, например, при использовании неявных, итерационных методов решения уравнений, таких как метод Гаусса-Зейделя , элемент n+1 массива зависит от первых элементов n, предотвращая векторизацию.

Я реализовал gauss-seidel следующим образом:

def gaussSeidel(A, b, x0, tol = 1e-10):
    n = len(A)
    x = np.copy(x0)
    k = 0

        while (True):
        print(k)
        for i in range(n):
            x[i] = (b[i] - np.delete(A[i], i) @ np.delete(x, i)) / A[i, i]

        if max(abs(x - x0)) < tol:
            return x

        x0 = np.copy(x)
        k += 1

Таким образом, медленнее, чем явная версия, такая как метод Якоби,

def jacobi(A, b, x0, tol = 1e-10):
    n = len(A)
    M = np.copy(A)
    DInv = np.zeros((n, n))
    k = 0
    for i in range(n):
        DInv[i, i] = 1 / M[i, i]
        M[i, i] = 0

    while (True):
        x = DInv @ (b - M @ x0)
        if max(abs(x - x0))  < tol:
            return x
        x0 = x
        k += 1
        print(max(abs(x - x0)))

Насколько я знаю, это не такЗдесь можно устранить цикл for, но есть ли способ ускорить его?

Answer 1

Что-то, что могло бы помочь, перемещает if-тест за пределы следующего цикла

        s = b[i]
        for j in range(n):
            if i == j:
                continue
            s -= A[i, j] * x[j]

, как в

        s = b[i] + A[i, i] * x[i]
        for j in range(n):
            s -= A[i, j] * x[j]

Таким образом, n if-тестов избегают.

Еще одна вещь, которую стоит попробовать, - это пакет numba , который может компилировать ваш код на лету без дополнительных затрат на этап внешней компиляции.