如何根据 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])