numpyで特定のaxisの値に応じて処理を変えたい


pythonのnumpyを使って、特定のaxisの値に応じて処理を変えたいと思っています。

例えば、下記のような `a[3][3][3]` のndarrayがあった際に

[code]

    [[[ 0  1  2]
      [ 3  4  5]
      [ 6  7  8]]
    
     [[ 9 10 11]
      [12 13 14]
      [15 16 17]]
    
     [[18 19 20]
      [21 22 23]
      [24 25 26]]]
    
[/code]

`a[:, :, :1]` に対しては10以上を0に上書きし、`a[:, :, 1:2]` に対しては15以上を0に上書きし、`a[:, :, 2:]`
に対しては20以上を0に上書きする処理をしたいと思っています。

numpyの関数を調べて、下記の手順で実装して意図通りの処理ができるようになりましたが、スライスを用いている個所を、(スライスを用いない)スマートな処理ができないかを考えております。

  1. スライスを用いてndarrayを分割する
  2. それぞれのndarrayに対して処理を実施する
  3. 再度dstackでndarrayをまとめる。

おそらく何か適切な関数があると考えているのですが、そのような処理に適切か関数はありますでしょうか？

[code]

    import numpy as np
    
    def main():
        a = np.arange(27).reshape(3, 3, 3)
        print(a)
    
        a0 = a[:, :, :1]
        a1 = a[:, :, 1:2]
        a2 = a[:, :, 2:]
    
        a0_update = np.where(a0 < 10, a0, 0)
        a1_update = np.where(a1 < 15, a1, 0)
        a2_update = np.where(a2 < 20, a2, 0)
    
        a_update = np.dstack((a0_update, a1_update, a2_update))
    
        print(a_update.shape)
        print(a_update)
    
    if __name__ == "__main__":
        main()
    
[/code]




[numpy.moveaxis — NumPy v1.23
Manual](https://numpy.org/doc/stable/reference/generated/numpy.moveaxis.html)

> **Returns: result :** _**np.ndarray**_
>
> Array with moved axes. This array is **a view** of the input array.
[code]

    b = np.moveaxis(a, -1, 0)
    for i, n in enumerate([10, 15, 20]):
        b[i][b[i] >= n] = 0
    
    print(a)
    
    #
    [[[ 0  1  2]
      [ 3  4  5]
      [ 6  7  8]]
    
     [[ 9 10 11]
      [ 0 13 14]
      [ 0  0 17]]
    
     [[ 0  0  0]
      [ 0  0  0]
      [ 0  0  0]]]
    
[/code]




ブールインデックスを使って以下のようにすればできますよ

[code]

    a[a >= np.array([10, 15, 20])] = 0
    
[/code]

