代码之家  ›  专栏  ›  技术社区  ›  ClimateUnboxed

如何基于3D数组中的索引对4D数组中有条件选择的numpy数组项进行平均

  •  0
  • ClimateUnboxed  · 技术社区  · 5 年前

    我想根据使用3D数组的索引,对4D numpy数组中有条件选择的元素进行平均。

    资料

    我的3D阵列在哪里 我用来有条件地采样的只是[ntime,ny,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,(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个数字对应于COND=1时x和y点的平均数据值

    这个问题与此有关: filtering a 3D numpy array according to 2D numpy array

    但我的KLV指数在中间,而不是左边或右边,所以我不能使用 [...,None] 解决方案

    0 回复  |  直到 5 年前
        1
  •  1
  •   Sayandip Dutta    5 年前

    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))
    

    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))
    

    直接使用条件数组

    ndarray s是 numpy.transpose ,正如您在链接文章中所看到的,在您的问题中,索引时,维度是左对齐的,但当前形式的数组不适合这种索引,因此,如果聚合维度位于最右侧,而索引维度位于左侧,那么就可以了。

    因此,如果您的数据可以重新排序:

    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

    roll 您的轴(=1)到末端(即第四维),使用 numpy.rollaxis

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

    使用移动轴 np.转置

    或者,你可以 move 你的轴心来自 source 维度到 destination np.moveaxis

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