numpy.trace – 配列の対角線の要素の合計を取得

numpy.traceは、渡した配列の対角線の要素の合計を返す関数です。

2次元配列を渡した場合は、指定の対角線の要素の合計値を返します。3次元以上の配列を渡した場合は、オプション引数のaxis1とaxis2で指定した軸の2次元配列の対角線要素を格納した配列を返します。この時のshapeは、元の配列から指定の軸を取り除いたものになります。

実際のコードで確認していきましょう。

1. numpy.traceの使い方

それでは、numpy.traceの書き方から確認していきましょう。

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

それでは実際のコードを見ながら、使い方を確認していきましょう。

1.1. 配列の対角線の要素の合計値を返す

numpy.traceは、引数に渡した2次元配列の指定の対角線の要素の合計値を返します。

以下のコードで生成した2次元配列を例に見ていきましょう。

import numpy as np
arr = np.random.randint(1, 10, (5, 5))
print(arr)

[[1 9 1 2 9]
 [8 7 3 2 2]
 [5 3 3 5 3]
 [5 5 4 3 1]
 [4 7 4 8 7]]

この配列を、numpy.traceに渡すと、以下のように対角線の要素[1, 7, 3, 3, 7]の合計値を返します。

#  対角線の要素の合計を返す
np.trace(arr)

21

1.2. 対角線の位置を指定（offset)

オプション引数offsetで、どの対角線の要素を合計するかを指定することができます。

正の値を指定した場合は、上から右に、その値だけずれていきます。負の値を指定した場合は、下から左に、その値だけずれていきます。

以下のコードでは、offset=1を指定しているので[9, 3, 5, 1]の合計値を返しています。

#  offsetで正の値を指定した場合、その値だけ右にずれる
np.trace(arr, offset=1)

18

以下のコードでは、offset=-2を指定しているので[5 5 4]の合計値を返しています。

#  offsetで負の値を指定した場合、その値だけ左にずれる
np.trace(arr, offset=-2)

14

1.3. 3次元配列以上の場合(axis1, axis2)

numpy.traceに3次元以上の配列を渡す場合は、その配列のどの軸とどの軸の対角線の合計を取得するかをオプション引数axis1, axis2で指定します。デフォルトでは、axis1=0, axis2=1になっています。

以下のshape(2, 3, 3)の配列を例に見ていきましょう。

'''  3次元配列の場合  '''
arr = np.random.randint(1, 10, (2, 3, 3))
print(arr)

[[[9 5 6]
  [9 3 8]
  [7 1 8]]

 [[1 9 4]
  [5 2 3]
  [8 6 9]]]

この3次元配列をnumpy.traceに渡すと、次のように返ってきます。

#  デフォルトはaxis0とaxis1の2次元配列として、その対角線を返す。
np.trace(arr)

array([14,  7,  9])

3次元配列の場合、axis=0は「奥行き」、axis=1は「行」、axis=2は「列」です。そのため、デフォルトのaxis1=0, axis2=1では奥行き方向の要素を1つの軸に、行方向の要素を1つの軸にした2次元配列の対角線の要素を合計しています。

また、numpy.traceの戻り値のshapeは、元の配列のshapeから指定の軸を除いたものになります。この場合では、元の配列はshape(2, 3, 3)で、axis1=0, axis2=1なので、戻り値の配列は、それを除いたshape(3, )になります。

次に、axis1=0, axis2=2を指定した場合を見てみましょう。

以下のコードをご覧ください。

#  奥行き方向の要素を列として並べた2次元配列の対角線を返す
np.trace(arr, axis1=0, axis2=2)

array([18, 11, 13])

この場合、奥行き方向の要素を1つの軸に、列方向の要素をもう1つの軸にした2次元配列の対角線の合計を求めています。

次に、axis1=1, axis2=2にした場合は以下のようになります。

#  行方向の要素を列として並べた2次元配列の対角線を返す。
np.trace(arr, axis1=1, axis2=2)

array([20, 12])

1.4. 要素の型を指定(dtype)

オプション引数dtypeで、戻り値の要素の型を指定することができます。

なおdtypeの種類や、指定する際の書き方などは『NumPyのdtype属性の一覧と確認・指定・変更方法』で解説していますので、ご確認ください。

以下のコードでは、float型を指定しています。

'''  float型を指定  '''
#  2次元配列を生成
arr = np.random.randint(1, 10, (3, 3))
np.trace(arr, dtype=float)

12.0

以下のコードでは、complex型でビット数も指定しています。

#  3次元配列の場合
arr = np.random.randint(1, 10, (2, 3, 3))
#  complex256型を指定
np.trace(arr, dtype=np.complex256)

array([17.+0.j, 10.+0.j,  5.+0.j], dtype=complex256)

ただし、Python3標準のint64, float64, complex128以外を使う際は、予期せぬところで数値の丸めのルールが異なるので、本当に必要だという確信がない限りはビット数の変更はおすすめしません。

1.5. 関数の戻り値で既存の配列を上書き(out)

numpy.trace()は、3次元配列以上の配列を渡した場合は、対角線の合計を要素とした新しい配列を生成します。

この新しい配列で既存の配列を上書きしたい場合はオプション引数outを使います。outで指定する配列は、numpy.trace()の戻り値とshapeが同一である必要があります。もし同一でなければエラーになります。

そして、配列のdtypeは、アウトプット先の配列のdtypeになります。

なお、3次元配列以上を渡した場合のnumpy.trace()の戻り値のshapeは、「1.3. 3次元配列以上の場合(axis1, axis2)」で解説していますので頭に入れておきましょう。

それでは以下のコードをご覧ください。

'''  関数の戻り値で既存の配列を上書き  '''
#  3次元配列を生成
arr = np.random.randint(1, 10, (2, 3, 3))
print('arr: \n',arr, '\n')
#  アウトプット先の配列を生成
a_out = np.zeros(3)
print('a_out(numpy.trace()実行前): \n', a_out, '\n')
#  numpy.trace()を実行
np.trace(arr, out=a_out)
print('a_out(numpy.trace()実行後): \n', a_out)

arr: 
 [[[4 3 8]
  [1 6 2]
  [5 4 4]]

 [[5 8 3]
  [8 6 7]
  [1 1 9]]] 

a_out(numpy.trace()実行前): 
 [0. 0. 0.] 

a_out(numpy.trace()実行後): 
 [12.  9. 15.]

関数実行前と実行後を見ると、配列が上書きされていることがわかりますね。

2. まとめ

以上が、numpy.traceの使い方です。NumPy配列の対角線の操作については、以下の関数も併せて確認しておくと良いでしょう。

• eye: 対角線の要素が1で他は0埋めの配列を生成
• identity: 対角線の要素が1で他は0埋めの正方形の配列を生成
• diag: 2次元配列から対角線の要素を取得または1次元配列を対角線に配置
• diagonal: 既存の配列の対角線を要素とした配列を生成
• diagflat: 対角線が指定の要素、他が0埋めの配列を生成