Векторизованное умножение нескольких последовательностей матриц 2х2

Question

Для невекторизованной версии у меня есть последовательность (2, 2) -образных матриц (то есть ndarray формы (n, 2, 2)), и мне нужно умножать их (умножение матриц) последовательно, что означает n последовательное умножение матриц. Вот как будет выглядеть минимальный пример

def get_matrix_product_eig_val(J):
    # J holds the sequence of matrices to multiply and has shape=(n, 2, 2)
    M = np.identity(2, dtype=np.double)
    for i in range(n_gens):
        M = np.matmul(M, J[i])
    eig_val, eig_vec = np.linalg.eig(M)  # eig_val is what I'm interested in
    return eig_val

Теперь у меня есть массив k таких последовательностей матриц (ndarray формы (k, n, 2, 2)), и мне нужно сделать один и тот же последовательный матричный продукт для каждой из k записей. Наивным подходом было бы сделать

# Now J has shape=(k, n, 2, 2)
for i in range(k):
    eig_vals[i] = get_matrix_product_eig_val(J[i, :, :, :])

Есть ли способ избавиться от цикла и сделать это векторизованным способом?

Примечания:

1) n ожидается порядка порядка 100. k может быть в диапазоне от ~ 100 до ~ 10000

2) Некоторые люди предложили заменить внутренний цикл на np.linalg.multi_dot. Это на самом деле значительно замедляет ход событий

3) Я вижу отдаленное будущее приложение для 3x3 матриц, но конкретное решение для 2x2 вполне подойдет. В любом случае все матрицы будут квадратными и будут иметь одинаковые размеры

Answer 1

for i in range(n - 1):
    J[:, i + 1, :, 0] = np.broadcast_to((J[:, i+1, 0, 0])[..., None], (k, 2)) * J[:, i, :, 0] + \
                        np.broadcast_to((J[:, i+1, 1, 0])[..., None], (k, 2)) * J[:, i, :, 1]
    J[:, i + 1, :, 1] = np.broadcast_to((J[:, i+1, 0, 1])[..., None], (k, 2)) * J[:, i, :, 0] + \
                        np.broadcast_to((J[:, i+1, 1, 1])[..., None], (k, 2)) * J[:, i, :, 1]

eig_vals, _ = np.linalg.eig(J[:, -1, :, :])