Быстрое (самое быстрое) возведение в степень трехмерной матрицы NumPy ⇐ Python

[Anonymous]

Код: Выделить всё

Q

представляет собой трехмерную матрицу и может иметь, например, следующую форму:

(4000, 25, 25)

Я хочу возвести Q в степень n для {0, 1, ..., k и просуммировать все это.
По сути, я хочу посчитать

\sum_{i=0}^{k-1}Q^n

У меня есть следующая функция работает должным образом:

Код: Выделить всё

def sum_of_powers(Q: np.ndarray, k: int) -> np.ndarray:
Qs = np.sum([
np.linalg.matrix_power(Q, n) for n in range(k)
], axis=0)

return Qs

Можно ли ускорить мою функцию или есть более быстрый способ получить тот же результат?

Подробнее здесь: https://stackoverflow.com/questions/793 ... -3d-matrix