俵言

しがない社会人が書く、勉強とかのこと。最近は機械学習や kaggle 関連がメイン。

小ネタ:numpy.repeat での繰り返しパターン生成と advanced indexing

日常的にちょくちょく numpy 芸・ pandas 芸をするのですが、そういうのを備忘録的に書いていこうかなと*1
今回は numpy.repeat + α のお話です。

目次

やりたいこと

例えば以下のような一次元配列があったときに

>>> import numpy as np 
>>> a = np.arange(4)
>>> a
array([0, 1, 2, 3])

こういう風に縦に並べたいな、ってことがあります。

array([[0, 1, 2, 3],
       [0, 1, 2, 3],
       [0, 1, 2, 3]])

ひとつのやり方は、並べたい数だけリストに入れた上で numpy.stack を使うこと

>>> np.stack([a for _ in range(3)], axis=0)
array([[0, 1, 2, 3],
       [0, 1, 2, 3],
       [0, 1, 2, 3]])

なのですが、もっとシンプルに書きたくなることありませんか?(少なくとも僕にはあります。)

こういう場合とりあえず numpy の公式 docs を漁るのですが、numpy.repeat とかいう如何にもな名前のやつがあったので「これだ...!」と試そうとしたわけです。

素朴な失敗例

先程の a は一次元配列です。こう、画面上の見た目としては縦方向に伸びて欲しいなーとやってみます。(つまり np.stack 同様に新しい軸が生成されて欲しい。)

>>> np.repeat(a, 3, axis=0)
array([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3])

違う、そうじゃない(画像略)。これはちょっと予想外で、軸が増えないにしても、

array([0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3])

にはならないんですね🤔*2 よくよく見ると docs*3 での説明文は

Repeat elements of an array.

となっていて、要素ごとに繰り返しを作ります。要素ごとに判定が行われるわけですから、np.repeat(a, 3, axis=0) は実際は np.repeat(a, [3, 3, 3, 3], axis=0) と同じです。こう考えると以下のような動きも納得がいきますね😀*4

>>> np.repeat(a, [5, 2, 0, 3], axis=0)
array([0, 0, 0, 0, 0, 1, 1, 3, 3, 3])

解決策:新しい軸を作る

本題に戻りますが、実はやり方は非常に簡単です。軸を新しく作ってその方向に繰り返せばいいだけ。

>>> np.repeat(a[None, :], 3, axis=0)
array([[0, 1, 2, 3],
       [0, 1, 2, 3],
       [0, 1, 2, 3]])

縦に並べて横に繰り返したいならこう。

>>> np.repeat(a[:, None], 3, axis=1)
array([[0, 0, 0],
       [1, 1, 1],
       [2, 2, 2],
       [3, 3, 3]])

二次元配列に対して適用したい場合も同様。

>>> A =  np.arange(2 * 3).reshape(2, 3)
>>> A
array([[0, 1, 2],
       [3, 4, 5]])
>>> np.repeat(A[None, :], 3, axis=0)
array([[[0, 1, 2],
        [3, 4, 5]],

       [[0, 1, 2],
        [3, 4, 5]],

       [[0, 1, 2],
        [3, 4, 5]]])

やったぜ。

でも実は、もっと簡単な方法があります。それは numpy.tile*5 を使うことです。

>>> np.tile(A, (3, 1, 1))
array([[[0, 1, 2],
        [3, 4, 5]],

       [[0, 1, 2],
        [3, 4, 5]],

       [[0, 1, 2],
        [3, 4, 5]]])

「要素ごとに繰り返し数指定しないしこれでよくね?」と思わなくもないですが、numpy.tile は各軸に対して数を指定する(上の例では (3, 1, 1)) ので*6少し面倒な気もします。
ただ、複数の軸方向に増やしたい場合は非常に便利ですね(np.repeat は一つの軸についてしか処理できないため)。

応用:advanced indexing

「いやこれ何に使うの?」って方も居ると思うので例を上げます。実は直近で使う機会*7があったため記事を書きました。

与えられたデータが一見テーブル形式、つまり事例ごとにベクトルが対応している (shape = (n_example, n_feature)) が、実際は事例ごとに対応する行列(shape = (n_example, h_feature, w_feature))が含まれているケースを考えます*8
単純に行列を flatten した形になっているなら reshape だけで済みますが、たまに厄介なケースがあります。以下のような例を考えてみましょう(ちょっと微妙な例ですが許して...)。

>>> D_given = np.array([
...   [111, 112, 121, 122, 113, 123, 131, 132, 133, 11, 12, 13, 14], 
...   [211, 212, 221, 222, 213, 223, 231, 232, 233, 21, 22, 23, 24], 
... ])
>>> D_given
array([[111, 112, 121, 122, 113, 123, 131, 132, 133,  11,  12,  13,  14],
       [211, 212, 221, 222, 213, 223, 231, 232, 233,  21,  22,  23,  24]])

実はこのテーブルデータは以下のベクトルと行列から作られています。

>>> D_org_vecs = np.array([[11, 12, 13, 14], [21, 22, 23, 24]])
>>> D_org_mats = np.array([
...   [[111, 112, 113], [121, 122, 123], [131, 132, 133]],
...   [[211, 212, 213], [221, 222, 223], [231, 232, 233]],
... ])
>>> D_org_vecs
array([[11, 12, 13, 14],
       [21, 22, 23, 24]])
>>> D_org_mats
array([[[111, 112, 113],
        [121, 122, 123],
        [131, 132, 133]],

       [[211, 212, 213],
        [221, 222, 223],
        [231, 232, 233]]])

ベクトルの方はスライスすれば良いだけですが、行列の方は処理するのに少し工夫が必要です。

与えられたデータの一行目(一事例目)のうち行列を構成するものだけ見てみましょう。

>>> D_given[0, -4:]
array([111, 112, 121, 122, 113, 123, 131, 132, 133])

左から 0 - 8 の index を振ると、実は元データの行列上の配置は以下です。

>>> idx_arr = np.array([[0, 1, 4], [2, 3, 5], [6, 7, 8]])
>>> idx_arr
array([[0, 1, 4],
       [2, 3, 5],
       [6, 7, 8]])

numpy は賢いので、元となる配列の要素の位置を値として格納した配列(上記の idx_arr) を渡すと変換してくれます。

>>>D_given[0][idx_arr]
array([[111, 112, 113],
       [121, 122, 123],
       [131, 132, 133]])

さて、ここからが numpy.repeatnumpy.tile が関わってくる部分です。一行一行処理するのであれば上に書いたものを利用して、

>>> D_convert = np.zeros((2, 3, 3), dtype=int) 
>>> for i in range(2):
...   D_convert[i] = D_given[i][idx_arr]
>>> D_convert
array([[[111, 112, 113],
        [121, 122, 123],
        [131, 132, 133]],

       [[211, 212, 213],
        [221, 222, 223],
        [231, 232, 233]]])

としてしまえば良いです。ただ、一気に処理したくないですか?

Advanced Indexing*9 を利用することで、この処理は一発で書けます。

>>> D_given[
...   # #  D_given において何行目かを表現
...   np.tile(np.arange(2)[:, None, None], (1, 3, 3)), 
...   # # D_given の各行において何列目かを表現
...   np.repeat(idx_arr[None, ...], 2, axis=0)
... ]
array([[[111, 112, 113],
        [121, 122, 123],
        [131, 132, 133]],

       [[211, 212, 213],
        [221, 222, 223],
        [231, 232, 233]]])

簡単に説明すると、一行一行処理する場合はベクトル(一次元配列)に対しての処理なので idx_arr を渡すだけで済むのですが、全体に一気に渡す場合はどの行かを示す必要があるわけです。 D_given に渡している要素の一つ目は対応する行を敷き詰めたものになっており、

>>> np.tile(np.arange(2)[:, None, None], (1, 3, 3))
array([[[0, 0, 0],
        [0, 0, 0],
        [0, 0, 0]],

       [[1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]]])

二つ目の要素は idx_arr を複製したものになります。

>>> np.repeat(idx_arr[None, ...], 2, axis=0) 
array([[[0, 1, 4],
        [2, 3, 5],
        [6, 7, 8]],

       [[0, 1, 4],
        [2, 3, 5],
        [6, 7, 8]]])

もちろん行ごとに別の idx_arr を渡すことも可のですが、そういうデータは色々やばそうな気がする...

本当は Advanced Indexing についてもっとちゃんと説明した方がいいのですが、記事の長さが2-3倍になりそうなので割愛させてください*10

おわりに

こういうの考えてる間に愚直に処理書いて回した方が早くね? ケースバイケースだとは思いますが、元の行列に変換する処理が非効率な場合は一旦 idx_arr を作ってしまって一気に処理した方がめちゃくちゃ速くなると思います。上述のややこしい advanced indexing を無理に使わなくても、一行一行 idx_arr で処理するだけでもマシになる気がする。

パズル的に楽しいので、今後もちょくちょく書こうかなと思います。

*1:「こんなん知ってるわ」って方はどうか許して...

*2:後述する numpy.tile で実現可能

*3:https://numpy.org/doc/stable/reference/generated/numpy.repeat.html

*4:多次元配列になると動きが複雑になりますが、ここでは割愛します

*5:https://numpy.org/doc/stable/reference/generated/numpy.tile.html

*6:一部を省略できる場合もある

*7:https://www.kaggle.com/ttahara/technique-for-converting-fnc-to-matrix

*8:行単位で事例を管理したいという場合もあるのであり得ない話ではない

*9:https://numpy.org/doc/stable/reference/arrays.indexing.html#advanced-indexing

*10:前々から自分なりの理解をまとめたいなーとは思っているのですが中々書けない