numpy.ufunc.reduceatの意味

100 Numpy Exercisesの78問目より。この関数意味わからなかった。ちなみにこんな問題。ブロック単位での和を計算することを想定している。

78.Consider a 16x16 array, how to get the block-sum (block size is 4x4)? (★★★)

# Author: Robert Kern

Z = np.ones(16,16)
k = 4
S = np.add.reduceat(np.add.reduceat(Z, np.arange(0, Z.shape[0], k), axis=0),
                                       np.arange(0, Z.shape[1], k), axis=1)

試してみる

リファレンスからどんな挙動をするのか試してみる。

numpy.ufunc.reduceat — NumPy v1.14 Manual

>>> import numpy as np
>>> x = np.arange(16).reshape(4,4)
>>> x
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])

とりあえず第2引数のベクトルの要素を1つにしてadd.reduceしてみる（ufuncのところは関数が入るのでaddでもmultiplyでも何でも良い）。

>>> np.add.reduceat(x,[0])
array([[24, 28, 32, 36]], dtype=int32)
>>> np.add.reduceat(x,[1])
array([[24, 27, 30, 33]], dtype=int32)
>>> np.add.reduceat(x,[2])
array([[20, 22, 24, 26]], dtype=int32)
>>> np.add.reduceat(x,[3])
array([[12, 13, 14, 15]], dtype=int32)

なるほど、1番上から、

[0]の場合：4行全ての列で集計した合計
[1]の場合：2行目～4行目の列で集計した合計
[2]の場合：3行目～4行目の列で集計した合計
[3]の場合：4行目の列で集計した合計＝4行目の値

やっていることはわかった。要素を2つに増やしてみる。以下、形式的にn行目の列単位の集計した合計をrownと書く。[2]の場合はrow3+row4、[3]の場合はrow4となる。

>>> np.add.reduceat(x,[0,1])
array([[ 0,  1,  2,  3],
       [24, 27, 30, 33]], dtype=int32)
>>> np.add.reduceat(x,[0,2])
array([[ 4,  6,  8, 10],
       [20, 22, 24, 26]], dtype=int32)
>>> np.add.reduceat(x,[0,3])
array([[12, 15, 18, 21],
       [12, 13, 14, 15]], dtype=int32)
>>> np.add.reduceat(x,[1,3])
array([[12, 14, 16, 18],
       [12, 13, 14, 15]], dtype=int32)
>>> np.add.reduceat(x,[2,3])
array([[ 8,  9, 10, 11],
       [12, 13, 14, 15]], dtype=int32)
>>> np.add.reduceat(x,[1,2])
array([[ 4,  5,  6,  7],
       [20, 22, 24, 26]], dtype=int32)

まとめてみる。

[0,1]→[row1] [row2+row3+row4]
[0,2]→[row1+row2] [row3+row4]
[0,3]→[row1+row2+row3] [row4]
[1,3]→[row2+row3][row4]
[2,3]→[row3][row4]
[1,2]→[row2][row3+row4]

だいたい傾向が読めてきたのではないだろうか？リファレンスを読むと次のようなルールがある。

For i in range(len(indices)), reduceat computes ufunc.reduce(a[indices[i]:indices[i+1]]), which becomes the i-th generalized “row” parallel to axis in the final result (i.e., in a 2-D array, for example, if axis = 0, it becomes the i-th row, but if axis = 1, it becomes the i-th column). There are three exceptions to this:

when i = len(indices) - 1 (so for the last index), indices[i+1] = a.shape[axis].

if indices[i] >= indices[i + 1], the i-th generalized “row” is simply a[indices[i]].

if indices[i] >= len(a) or indices[i] < 0, an error is raised.

もともとreduceという関数を元にしていて、それの区間選択できるバージョンらしい。reduceを実行してみると、

>>> np.add.reduce(x,axis=0)
array([24, 28, 32, 36])

reduceのaxisは列で集計するか行で集計するかというオプションで、reduceatにも存在する。つまりreduceの場合は、reduceatにおける[row1+row2+row3+row4]と同等の結果を返した。

じゃあ区間選択はどうやっているのかというと、2要素の場合を見るにこういうルーチンだろう。こう書けばプログラムに馴染みがある人は「ああ、なるほど」となるはずだ。

for i in range(len(indices)):
    if i == len(indices)-1:
        j = len(indices)
    else:
        j = indices[i+1]
    result = 0
    for p in range(i, j):
        result += p+1行目の列単位で集計した合計（＝row[p+1]）
    yield result

ただしこれは不完全で、indices[i+1]<indices[i]の場合は関数適用するのではなく、単純にrow[indices[i]]を返すという例外がある。他の例外として範囲外のインデックスを与えるとエラーを返すとかあるが、当たり前なので省略。

この例外を使った例はこちら。何を返し、どういう計算が行われているかわかりますか？

>>> np.add.reduceat(x,[0,2,0])
array([[ 4,  6,  8, 10],
       [ 8,  9, 10, 11],
       [24, 28, 32, 36]], dtype=int32)

結果の1行目は[row1+row2]、2行目[row3]、3行目は[row1+row2+row3+row4]。indices[1]>indices[2]なので例外が適用され、2行目は[row3+row4]ではなく[row3]なのがポイント。さらに3行目でインデックスが0に戻り、row1～row4までの合計を計算している。

一周するとインデックスが戻るのはreduceat(x,[0,2,1,3])のように実行するとより詳しく確認できる。

>>> np.add.reduceat(x,[0,2,1,3])
array([[ 4,  6,  8, 10],
       [ 8,  9, 10, 11],
       [12, 14, 16, 18],
       [12, 13, 14, 15]], dtype=int32)

1行目はrange(0,2)なので、[row1+row2]。2行目は例外適用で[row3]。3行目は一周してrange(1,3)で[row2+row3]、4行目はrange(3,4)なので[row4]が返されている。なので、一周するというよりかは、2行目の例外が一周のときの挙動でその例外適用＝インデックスを0に戻すなのかもしれない（もうちょっと詳しく確認する必要はあり）。

かなりとっつきにくい関数ではあったが、理解を深めることはできた。確かに、冒頭の100 Numpy Exercisesに出ているように、ブロック単位での合計を計算するのに便利な関数ではあると思う。

shikoan’s memo

プログラミング初心者のチラ裏

ぷろぐらみんぐ帳

numpy.ufunc.reduceatの意味

試してみる