Numpy 3次元配列のインデックスを変数に格納する

結論

slice オブジェクトというものがあります。これをタプルで複数まとめて扱えます。

A[(slice(3, 6), slice(5, 8))] == A[3:6, 5:8]

以下、2次元配列の例。

>>> A = np.zeros([10, 10])
>>> A[(slice(3, 6), slice(5, 8))] = 1
>>> A
array([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 1., 1., 1., 0., 0.],
       [0., 0., 0., 0., 0., 1., 1., 1., 0., 0.],
       [0., 0., 0., 0., 0., 1., 1., 1., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])

具体例

以下のような関数を作るとします。

def partial_copy(source, destination):
    destination = destination.copy()
    destination[34:194, 10:150] = source[34:194, 10:150]
    return destination

やはりコピー範囲は引数で与えたいので次のように変更します。

def partial_copy(source, destination, y_start, y_end, x_start, x_end):
    destination = destination.copy()
    destination[y_start:y_end, x_start:x_end] = source[y_start:y_end, x_start:x_end]
    return destination

result = partial_copy(A, B, 34, 194, 10, 150)

だいぶ長い一行になりました。改行や変数を y_s 等にすると短くなりますが、可読性が落ちます。

そこで、tuple of slices を使って、次のように書くとスッキリします。

def partial_copy(source, destination, copy_area):
    destination = destination.copy()
    destination[copy_area] = source[copy_area]
    return destination

result = partial_copy(A, B, (slice(34, 194), slice(10, 150))

書いてて、呼び出し側が長くなって可読性が落ちていることに気づきました。

こんな書き方があるんだと、参考までに。