Реализация алгоритма трехдиагональных матриц (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 сохраняет правую часть последнего выражения в temprorary массиве, а затем присваивает ссылки на его элементы в alpha[1:]. Поэтому a[1:] * alpha[:-1] - это всего лишь массив нулей.

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

Спасибо.

Answer 1

По-видимому, нет возможности сделать это в Python без использования C или его pythonic вариантов.

Answer 2

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