今回もNumPyの小ネタです。特に難しい処理ではないのですが、NumPyの行列(2次元配列)から三角行列を取り出したいことがあります。
自分の場合は、距離行列とか相関行列などに対して、$x_i$と$x_j$の距離や相関と、$x_j$と$x_i$の距離や相関は等しいから一方だけでいいとかそういう状況が多いです。
まぁ、値を取り出すだけなら、2重のforループなどをちゃんと書けば済む話なのですが、前回の記事で紹介したようなちょっとしたテクニックを使いたい場合など、明示的に三角行列を作りたいことがあります。
参考: NumPyの多次元配列から値が大きいn個のインデックスを取得する
自分で実装しても特に難しくないのですが、NumPyに専用関数が用意されていたのでこの記事ではそれを紹介したいと思います。
とりあえず、適当な行列を作って自分でやってみましょう。(実は実務では要素数が数万*数万みたいな巨大行列扱うことがあり、2重for文回すのはちょっと待つので嫌だなと思うのですが、小さい行列ならこれで十分です)
対角成分も0にする、狭義の下三角行列を作る場合は以下のようになるでしょうか。
import numpy as np
# 適当にデータを生成する
np.random.seed(4)
ary = np.random.randint(10, 100, size=(5, 5))
print(ary)
"""
[[56 65 79 11 97]
[82 60 19 68 65]
[65 67 46 60 54]
[48 62 13 10 65]
[31 31 83 48 66]]
"""
# 上三角成分を0にすることで下三角行列を作る
for i in range(ary.shape[0]):
for j in range(i, ary.shape[1]):
ary[i, j] = 0
print(ary)
"""
[[ 0 0 0 0 0]
[82 0 0 0 0]
[65 67 0 0 0]
[48 62 13 0 0]
[31 31 83 48 0]]
"""はい、何も難しくなくできましたね。
これを2重のfor文を使わずにやる場合、NumPyには triuとtril というそれぞれ上三角行列と下三角行列を取り出すメソッドが用意されています。
参考:
numpy.triu — NumPy v1.22 Manual
numpy.tril — NumPy v1.22 Manual
とりあえず、動かしてみましょうか。２個目にkっていうオプションの引数(デフォルトは0)を取り、これを調整することで対角成分を残すかどうか、また対角成分に限らずどの斜めラインまで成分を残すかを調整できます。
# もう一回データを生成する
np.random.seed(4)
ary = np.random.randint(10, 100, size=(5, 5))
# 上三角行列
print(np.triu(ary))
"""
[[56 65 79 11 97]
[ 0 60 19 68 65]
[ 0 0 46 60 54]
[ 0 0 0 10 65]
[ 0 0 0 0 66]]
"""
# k=1とすると、対角成分も消える。(消える行が右上に広がる)
print(np.triu(ary, k=1))
"""
[[ 0 65 79 11 97]
[ 0 0 19 68 65]
[ 0 0 0 60 54]
[ 0 0 0 0 65]
[ 0 0 0 0 0]]
"""
# k=2とすると、さらに消える(消える行が右上に広がる)
print(np.triu(ary, k=2))
"""
[[ 0 0 79 11 97]
[ 0 0 0 68 65]
[ 0 0 0 0 54]
[ 0 0 0 0 0]
[ 0 0 0 0 0]]
"""
# k=-1とすると、逆に残す範囲が左下に広がる。より小さい負の数も同様。
print(np.triu(ary, k=-1))
"""
[[56 65 79 11 97]
[82 60 19 68 65]
[ 0 67 46 60 54]
[ 0 0 13 10 65]
[ 0 0 0 48 66]]
"""
# trilは下三角行列
print(np.tril(ary))
"""
[[56 0 0 0 0]
[82 60 0 0 0]
[65 67 46 0 0]
[48 62 13 10 0]
[31 31 83 48 66]]
"""
# tril の k=1 は残す範囲が広がる。境界線が右上にスライドするという意味ではtriuと同じ。
print(np.tril(ary, k=1))
"""
[[56 65 0 0 0]
[82 60 19 0 0]
[65 67 46 60 0]
[48 62 13 10 65]
[31 31 83 48 66]]
"""
# tril で対角成分も消したい場合はk=-1
print(np.tril(ary, k=-1))
"""
[[ 0 0 0 0 0]
[82 0 0 0 0]
[65 67 0 0 0]
[48 62 13 0 0]
[31 31 83 48 0]]
"""狭義の三角行列(要するに対角成分も0)を取り出したい時、上三角行列(triu)の時はk=1で、下三角行列(tril)の時はk=-1 ってのがちょっと厄介ですね。まぁ、境界線が上に移動するか下に移動するかと考えるのが誤解が少ないかと思います。
三角行列といえば、NumPyには、tri という 三角行列を生成するメソッドもあります。(これもkという引数を取ります。)
参考: numpy.tri — NumPy v1.22 Manual
このtriで生成した三角行列を使ってマスクすることで、三角行列を作ることも可能です。というより、triuやtril の実装をみるとそういう作りになっています。
参考: triuのソースコード(GitHub)
NumPyのwhereメソッドとか使っていい感じに実装されていいますが、ぶっちゃけ掛け算(要素積)するだけで良いでしょう。
# tri で三角成分が1の三角行列を作れる
print(np.tri(5, k=-1))
"""
[[0. 0. 0. 0. 0.]
[1. 0. 0. 0. 0.]
[1. 1. 0. 0. 0.]
[1. 1. 1. 0. 0.]
[1. 1. 1. 1. 0.]]
"""
# これを要素積すると、元の行列は下三角行列になる。
print(ary * np.tri(5, k=-1))
"""
[[ 0. 0. 0. 0. 0.]
[82. 0. 0. 0. 0.]
[65. 67. 0. 0. 0.]
[48. 62. 13. 0. 0.]
[31. 31. 83. 48. 0.]]
"""
# 上三角行列が欲しいなら転置してから掛ける
print(ary * np.tri(5, k=-1).T)
"""
[[ 0. 65. 79. 11. 97.]
[ 0. 0. 19. 68. 65.]
[ 0. 0. 0. 60. 54.]
[ 0. 0. 0. 0. 65.]
[ 0. 0. 0. 0. 0.]]
"""さて、最後に処理時間を見ておきましょう。最初に少しぼやいていましたが、2重for文のデメリットは巨大な行列を処理する場合の処理時間です。2万行*2万行(要素数4億)の行列で見てみましょう。
データ生成。
big_ary = np.random.randint(100, size=(20000, 20000))まず、tirlです。こちら3秒ちょっとで終わりました。
% % time
big_ary_tril = np.tril(big_ary, k=-1)
"""
CPU times: user 1.18 s, sys: 1.33 s, total: 2.5 s
Wall time: 3.25 s
"""次にfor文を回した場合です。こちらは50秒近くかかっています。もう少しで1分超えますね。
% % time
for i in range(big_ary.shape[0]):
for j in range(i, big_ary.shape[1]):
big_ary[i, j] = 0
"""
CPU times: user 45.3 s, sys: 1.59 s, total: 46.8 s
Wall time: 49.5 s
"""速度、計算効率的にはtirl/triuを使った方が良いということが確認できました。
triu/trilは実装上は全ての要素に対して残すか0埋めするかの判定をかけているのに対し、for文の方は0で埋める要素だけにアクセスしているので、計算量というか演算回数だけならfor文の方が半分以下のはずですが、for文の方が圧倒的に遅いのが不思議です。NumPyのもっとコアな部分で並列処理がかなりしっかりと作り込まれているのでしょう。