Реализация алгоритма трехдиагональной матрицы (TDMA) с NumPy

Question

Я использую TDMA в Python, используя NumPy. Трехдиагональная матрица хранится в трех массивах:

a = array([...])
b = array([...])
c = array([...])

Я хотел бы эффективно рассчитать alpha коэффициенты. Алгоритм выглядит следующим образом:

# n = size of the given matrix - 1
alpha = zeros(n)
alpha[0] = b[0] / c[0]
for i in range(n-1):
    alpha[i+1] = b[i] / (c[i] - a[i] * alpha[i])

Однако это неэффективно из-за цикла for в Python. Хочу, хочу, что-то вроде этого подхода:

# n = size of the given matrix - 1
alpha = zeros(n)
alpha[0] = b[0] / c[0]
alpha[1:] = b[1:] / (c[1:] - a[1:] * alpha[:-1])

В этом последнем случае результат неверен, поскольку NumPy сохраняет правую часть последнего выражения во временном массиве и затем присваивает ссылки на его элементы alpha[1:]. Следовательно, a[1:] * alpha[:-1] - это просто массив нулей.

Есть ли способ указать NumPy использовать значения alpha, рассчитанные на предыдущих шагах внутри его внутреннего цикла?

Спасибо.

Answer 1

Очевидно, в Python нет способа сделать это без использования C или его pythonic вариаций.

Answer 2

Если его трехдиагональные системы, которые вы хотите решить, есть solve_banded() в numpy.linalg. Не уверен, что это то, что вы ищете.