numpy.trace

numpy.trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None)

Функция numpy.trace() вычисляет cумму диагональных элементов (след) матрицы.

Параметры:
a - массив NumPy или подобный массиву объект.
Двумерный или многомерный массив.
offset - целое число (необязательный параметр).
Определяет смещение диагонали по которой вычисляется сумма элементов, положительные числа смещают ее вверх, отрицательные - вниз.
axis1, axis2 - целые числа (необязательный параметр).
Определяют массива, относительно которых рассматриваются его подматрицы. По умолчанию axis1 = 0 и axis2 = 1, что соответствует двум первым осям, а не двум последним как принято в пакете linalg.
dtype - тип данных NumPy (необязательный параметр).
Определяет тип данных результата. По умолчанию dtype=None, т.е. тип данных не преобразуется.
out - массив NumPy (необязательный параметр).
Массив в который можно поместить результат функции. Данный массив должен соответствовать форме и типу данных результирующего массива функции (зачастую, тип данных может быть преобразован автоматически). Указание данного параметра, позволяет избежать лишней операции присваивания тем самым немного ускоряя работу вашего кода. Полезный параметр если вы очень часто обращаетесь к функции в цикле.
Возвращает:
результат - массив NumPy или число
Сумму диагональных элементов в виде числа если на вход подана одна матрица, в виде массива, если указан массив матриц.

Замечание

Данная функция может вычислять сумму диагональных элементов не только квадратных матриц, однако, такую операцию уже нельзя назвать следом матрицы.

Смотрите так же:
diag, diagflat

Примеры

>>> import numpy as np
>>> 
>>> a = np.eye(3)
>>> a
array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]])
>>> 
>>> np.trace(a)
3.0
>>> 
>>> 
>>> b = np.arange(1, 10).reshape(3, 3)
>>> b
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])
>>> 
>>> np.trace(b, offset = 0)    #  по умолчанию offset = 0
15
>>> 
>>> np.trace(b, offset = 1)
8
>>> np.trace(b, offset = -1)
12
>>> 
>>> 
>>> #  Массивы не обязаны быть квадратными:
... c = np.arange(1, 13).reshape(3, 4)
>>> c
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]])
>>> 
>>> np.trace(c)
18

По умолчанию axis1 = 0 и axis2 = 1, что немного нарушает правила модуля linalg, в котором принято, что подматрицы определяются на основании последних двух осей многомерного массива. Это приводит к немного странным результатам:

>>> d = np.arange(12).reshape(3, 2, 2)
>>> d
array([[[ 0,  1],
        [ 2,  3]],

       [[ 4,  5],
        [ 6,  7]],

       [[ 8,  9],
        [10, 11]]])
>>> 
>>> np.trace(d)
array([6, 8])

Дело в том, что из-за установленных по умолчанию параметров axis1 = 0 и axis2 = 1 обрабатываются не те массивы, которые мы ожидаем:

>>> d[:, :, 0]
array([[ 0,  2],
       [ 4,  6],
       [ 8, 10]])
>>> 
>>> np.trace(d[:, :, 0])
6
>>> 
>>> 
>>> d[:, :, 1]
array([[ 1,  3],
       [ 5,  7],
       [ 9, 11]])
>>> 
>>> np.trace(d[:, :, 1])
8

Если внутри вашего многомерного массива, по последним двум осям расположены матрицы, то достаточно указать именно две последние оси

>>> d
array([[[ 0,  1],
        [ 2,  3]],

       [[ 4,  5],
        [ 6,  7]],

       [[ 8,  9],
        [10, 11]]])
>>> 
>>> np.trace(d, axis1 = 1, axis2 = 2)
array([ 3, 11, 19])
>>> 
>>> np.trace(d, axis1 = -2, axis2 = -1)
array([ 3, 11, 19])

Параметр dtype позволяет указать тип данных результирующего массива:

>>> np.trace(d, axis1 = -2, axis2 = -1, dtype = np.float)
array([ 3., 11., 19.])

Параметр out позволяет указать массив в который необходимо поместить результат вычислений:

>>> trace_d = np.ones(3)
>>> trace_d
array([1., 1., 1.])
>>> 
>>> np.trace(d, axis1 = -2, axis2 = -1, out = trace_d)
array([ 3., 11., 19.])
>>> 
>>> trace_d
array([ 3., 11., 19.])