Трассировка каждой матрицы в массиве NumPy

Question

Учитывая список нескольких матриц, я хочу вычислить след каждой матрицы.

np.asarray([[[1,2,3],
             [4,5,6],
             [7,8,9]],
           [[10,11,12],
            [13,14,15],
            [16,17,18]]])

Результат должен быть:

np.asarray([15, 42])

Answer 1

Используйте axis1=1 и axis2=2 аргументы в np.trace.

np.trace(x, axis1=1, axis2=2)

Answer 2

Вот способ использования расширенной индексации и sum вдоль второй оси:

a = np.asarray([[[1,2,3],
                 [4,5,6],
                 [7,8,9]],
               [[10,11,12],
                [13,14,15],
                [16,17,18]]])

a[:, range(a.shape[1]), range(a.shape[2])].sum(1)
# array([15, 42])

Answer 3

Используйте np.trace функцию с аргументами axis1=1 и axis2=2

mat = np.asarray([[[1,2,3],
             [4,5,6],
             [7,8,9]],
           [[10,11,12],
            [13,14,15],
            [16,17,18]]])

np.trace(mat, axis1=1, axis2=2)

вывод:

array([15, 42])