如何在numpy中获取子数组

How to get a subarray in numpy

我有一个 3d 数组,我想得到一个以索引 indx 为中心的大小为 (2n+1) 的子数组。使用切片我可以使用

y[slice(indx[0]-n,indx[0]+n+1),slice(indx[1]-n,indx[1]+n+1),slice(indx[2]-n,indx[2]+n+1)]

如果我想为每个维度设置不同的尺寸,这只会变得更丑。有没有更好的方法来做到这一点。

您不需要使用 slice 构造函数,除非您想要存储切片对象供以后使用。相反,你可以简单地做:

y[indx[0]-n:indx[0]+n+1, indx[1]-n:indx[1]+n+1, indx[2]-n:indx[2]+n+1]

如果您想在不单独指定每个索引的情况下执行此操作,您可以使用列表理解:

y[[slice(i-n, i+n+1) for i in indx]]

您可以创建 numpy 数组以索引到 3D array 的不同维度,然后使用 ix_ function to create indexing map and thus get the sliced output. The benefit with ix_ is that it allows for broadcasted indexing maps. More info on this could be found here。然后,您可以为通用解决方案的每个维度指定不同的 window 大小。这是示例输入数据的实现 -

import numpy as np

A = np.random.randint(0,9,(17,18,16))  # Input array
indx = np.array([5,10,8])              # Pivot indices for each dim
N = [4,3,2]                            # Window sizes

# Arrays of start & stop indices
start = indx - N
stop = indx + N + 1

# Create indexing arrays for each dimension
xc = np.arange(start[0],stop[0])
yc = np.arange(start[1],stop[1])
zc = np.arange(start[2],stop[2])

# Create mesh from multiple arrays for use as indexing map 
# and thus get desired sliced output
Aout = A[np.ix_(xc,yc,zc)]

因此,对于具有 window 大小数组的给定数据,N = [4,3,2]whos 信息显示 -

In [318]: whos
Variable   Type       Data/Info
-------------------------------
A          ndarray    17x18x16: 4896 elems, type `int32`, 19584 bytes
Aout       ndarray    9x7x5: 315 elems, type `int32`, 1260 bytes

输出的 whos 信息 Aout 似乎与预期的输出形状一致,它必须是 2N+1