如何根據 3D 數組中的索引對 4D 數組中的有條件選擇的 numpy 數組條目進行平均

Question

我想根據使用 3D 數組的索引對 4D numpy 數組中的條件選擇元素進行平均。

換句話說，我的 4D 數組DATA具有以下維度： [ntime,nz,ny,nx]

其中我用來有條件地采樣的 3D 數組COND只是 [ntime,ny,nx] 的 function （時間片的數量，x 和 y 點相同）

我想做廣播，因此使用類似DATA[COND[None,...]]但問題是“缺失”的垂直維度不在右側，而是在時間與 x 和 y 空間之間的中間. 我可以遍歷垂直級別，但我認為這會很慢。 有沒有辦法以某種方式將 DATA 索引為

DATA[cond[times],:,COND[ys],COND[xs]]?

設置一些虛擬 arrays：

np.random.seed(1234)
COND=np.random.randint(0,2,(2,3,3))  # 2 time levels, 3 X points and 3 y points
DATA=np.random.randint(0,100,(2,2,3,3)) # 2 time levels, 2 Z levels, and 3 x and y points

給予：

COND
array([[[1, 1, 0],
        [1, 0, 0],
        [0, 1, 1]],

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

DATA
array([[[[26, 58, 92],
         [69, 80, 73],
         [47, 50, 76]],

        [[37, 34, 38],
         [67, 11,  0],
         [75, 80,  3]]],

給予：

   [[[ 2, 19, 12],
     [65, 75, 81],
     [14, 71, 60]],

    [[46, 28, 81],
     [87, 13, 96],
     [12, 69, 95]]]])

我可以使用 argwhere 找到參數：

idx=np.argwhere(COND==1)
array([[0, 0, 0],
       [0, 0, 1],
       [0, 1, 0],
       [0, 2, 1],
       [0, 2, 2],
       [1, 0, 0],
       [1, 0, 1],
       [1, 0, 2],
       [1, 1, 2]])

現在我想做類似的事情

np.mean(DATA[idx[...,None,...]])

或者

np.mean(DATA[idx[0],None,idx[1],idx[2])

這應該給我一個答案，其中 2 個數字對應於時間的平均 DATA 值，當 COND=1 時 x 和 y 點

這個問題與此有關： 根據 2D numpy array filter a 3D numpy array

但我的 klev 索引在中間而不是左或右，所以我不能使用[...,None]解決方案

Answer 1

使用zip獲取沿每個軸的索引

IIUC，你已經完成了大部分工作，即idx

>>> [*zip(*idx)]
[(0, 0, 0, 0, 0, 1, 1, 1, 1),
 (0, 0, 1, 2, 2, 0, 0, 0, 1),
 (0, 1, 0, 1, 2, 0, 1, 2, 2)]

>>> t, y, x = zip(*idx)
>>> DATA[t, :, y, x]

array([[26, 37],
       [58, 34],
       [69, 67],
       [50, 80],
       [76,  3],
       [ 2, 46],
       [19, 28],
       [12, 81],
       [81, 96]])

>>> DATA[t, :, y, x].mean(0)
array([43.66666667, 52.44444444])

使用np.where獲取索引

獲取numpy.where的更簡單方法：

>>> np.where(COND)
(array([0, 0, 0, 0, 0, 1, 1, 1, 1], dtype=int64),
 array([0, 0, 1, 2, 2, 0, 0, 0, 1], dtype=int64),
 array([0, 1, 0, 1, 2, 0, 1, 2, 2], dtype=int64))

使用 np.nonzero 獲取索引

或者， numpy.nonzero ，可能是最明確的：

>>> np.nonzero(COND)
(array([0, 0, 0, 0, 0, 1, 1, 1, 1], dtype=int64),
 array([0, 0, 1, 2, 2, 0, 0, 0, 1], dtype=int64),
 array([0, 1, 0, 1, 2, 0, 1, 2, 2], dtype=int64))

直接使用條件數組

值得注意的是，在處理numpy.transpose ndarray正如您在鏈接的帖子中看到的那樣，在您的問題中，在索引時，尺寸是左對齊的，但是您當前形式的數組不適合那種的索引，所以如果你的聚合維度在最右邊，而索引維度在左邊，那就可以了。

因此，如果您的數據可以重新排序：

Instead of:
dim = (2, 2, 3, 3)
axis-> 0, 1, 2, 3

It were:
dim = (2, 3, 3, 2)
axis-> 0, 2, 3, 1

它會奏效的。

使用np.transpose重新排序軸

您可以為此使用numpy.transpose ：

>>> np.transpose(DATA, axes=(0,2,3,1))[COND==1].mean(axis=0)
array([43.66666667, 52.44444444])

使用np.roll滾動軸

您還可以使用numpy.rollaxis將軸 (==1) roll到末端（即第 4 維）：

>>> np.rollaxis(DATA, 1, 4)[COND==1].mean(0)
array([43.66666667, 52.44444444])

使用np.transpose移動軸

或者，您可以使用move軸從source維度移動到destination維度， np.moveaxis軸 1 移動到軸 3：

>>> np.moveaxis(DATA, source=1, destination=3)[COND==1].mean(0)
array([43.66666667, 52.44444444])