问题描述
我想根据使用 3D 数组的索引对 4D numpy 数组中条件选择的元素求平均值。
换句话说，我的 4D 数组 DATA 具有以下维度：[ntime,nz,ny,nx]
因为我用来有条件地采样的 3D 数组 COND 只是 [ntime,nx] 的函数（时间片的数量，x 和 y 点相同）>
我想做广播，因此使用类似 DATA[COND[None,...]] 的东西，但问题是“缺失”的垂直维度不在右侧，而是在时间与 x 和 y 空间之间的中间。我可以在垂直水平上循环，但我认为这会很慢。有没有办法以某种方式将数据索引为
DATA[cond[times],:,COND[ys],COND[xs]]?
设置一些虚拟数组：
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,3)) # 2 time levels,2 Z levels,and 3 x and y points
给予：
COND
array([[[1,1,0],[1,[0,1]],[[1,1],0]]])
DATA
array([[[[26,58,92],[69,80,73],[47,50,76]],[[37,34,38],[67,11,[75,3]]],给予：
[[[ 2,19,12],[65,75,81],[14,71,60]],[[46,28,[87,13,96],[12,69,95]]]])
我可以使用 argwhere 找到参数：
idx=np.argwhere(COND==1)
array([[0,2],2]])
现在我想做类似的事情
np.mean(DATA[idx[...,None,...]])
或
np.mean(DATA[idx[0],idx[1],idx[2])
这应该给我一个答案，其中有 2 个数字对应于当时的平均数据值，当 COND=1 时 x 和 y 点
这个问题与此有关： filtering a 3D numpy array according to 2D numpy array
但是我的 klev 索引在中间而不是左边或右边，所以我不能使用 [...,None] 解决方案
解决方法
使用 zip 获取沿每个轴的索引
IIUC，你已经完成了大部分工作，即idx
>>> [*zip(*idx)]
[(0,1,1),(0,2,2)]
>>> t,y,x = zip(*idx)
>>> DATA[t,:,x]
array([[26,37],[58,34],[69,67],[50,80],[76,3],[ 2,46],[19,28],[12,81],[81,96]])
>>> DATA[t,x].mean(0)
array([43.66666667,52.44444444])
使用 np.where 获取索引
获得 numpy.where 的更简单方法：
>>> np.where(COND)
(array([0,1],dtype=int64),array([0,2],dtype=int64))
使用 np.nonzero 获取索引
或者，numpy.nonzero，可能是最露骨的：
>>> np.nonzero(COND)
(array([0,dtype=int64))
直接使用条件数组
值得注意的是，处理 ndarray 时的一个方便的技巧是 numpy.transpose，正如您在链接的帖子中看到的那样，在您的问题中，在编制索引时，维度是左对齐的，但您的数组在当前的形式不适合那种索引，所以如果你的聚合维度在最右边，而索引维度在左边，那就行了。
因此，如果您的数据可以重新排序：
Instead of:
dim = (2,3,3)
axis-> 0,3
It were:
dim = (2,2)
axis-> 0,1
本来可以的。
使用 np.transpose 重新排列坐标轴
您可以使用 numpy.transpose：
>>> np.transpose(DATA,axes=(0,1))[COND==1].mean(axis=0)
array([43.66666667,52.44444444])
使用 np.roll 滚动轴
你也可以roll你的轴（==1）到最后（即第四维），使用numpy.rollaxis：
>>> np.rollaxis(DATA,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])