将低维numpy数组的部分提取到更高维数组的最终轴

时间:2016-07-21 08:15:08

标签: python arrays numpy

我有一个静态形状 - (l,l)数组C。我想将它的一部分提取到其他数组K中,其形状为(m,m,n,n)。我要从C中提取的内容的起始索引在数组i0中给出,其形状为(m,m)

K的某些元素将由K[i,j,:,:] = C[i0[i,j]:i0[i,j]+n, i0[i,j]:i0[i,j]+n]给出。所以关闭其他一些类似的问题似乎可以做到这一点......

C[i0[None, None, ...] + np.arange(n)[..., None, None], 
  i0[None, None, ...] + np.arange(n)[..., None, None], I, J]

引发IndexError。我想这是因为C只是2D,并且维度不能增加。虽然通过平铺C可以很容易地解决这个问题,但由于C很大,重新制作m*m次会相当昂贵。

所以我的问题是如何将2D阵列的不同(2D)部分提取到4D阵列的相应部分。

1 个答案:

答案 0 :(得分:1)

一种方法是使用np.meshgrid创建2D(n,n)形状窗口对应的索引网格,添加i0,其中包含两个新轴,其中将进行广播。最后,我们只需索引C即可为我们提供所需的4D输出。因此,一个实现就像这样 -

N = np.arange(n)
X,Y = np.meshgrid(N,N)
out = C[i0[...,None,None] + Y,i0[...,None,None] + X]

示例运行 -

In [153]: C
Out[153]: 
array([[3, 5, 1, 6, 3, 5, 8, 7, 0, 2],
       [8, 4, 6, 8, 7, 2, 6, 2, 5, 0],
       [3, 7, 7, 7, 3, 4, 4, 6, 7, 6],
       [7, 0, 8, 2, 1, 1, 0, 4, 4, 6],
       [2, 4, 6, 0, 0, 5, 6, 8, 0, 0],
       [4, 6, 1, 0, 5, 6, 2, 1, 7, 4],
       [0, 5, 5, 3, 7, 5, 7, 1, 4, 0],
       [6, 4, 4, 7, 2, 4, 6, 6, 6, 5],
       [5, 2, 3, 2, 2, 5, 4, 5, 2, 5],
       [3, 7, 1, 0, 4, 4, 6, 6, 2, 2]])

In [154]: i0
Out[154]: 
array([[1, 0, 4, 4],
       [0, 4, 4, 0],
       [2, 3, 1, 3],
       [2, 2, 0, 4]])

In [155]: n = 3

In [157]: out[0,0,:,:]
Out[157]: 
array([[4, 6, 8],
       [7, 7, 7],
       [0, 8, 2]])

In [158]: C[i0[0,0]:i0[0,0]+n,i0[0,0]:i0[0,0]+n]
Out[158]: 
array([[4, 6, 8],
       [7, 7, 7],
       [0, 8, 2]])