python - 用于从 3-D 数组中选择具有不同起始索引的相同长度子数组的纯 numpy 表达式

Question

我有一个形状为 numpy 的 3 维 a 数组(我们称之为 (74, 74, 4563) )，我想从前两个维度的每个位置提取一个长度 n 子数组。但是，每个子数组都从不同的位置开始，具体取决于前两个维度 i 和 j 中的索引。

例如，如果 n=1000 ，我可能需要 a[0, 0, 0:1000] ，但也需要 a[0, 1, 2:1002] 等...我有一个二维数组(称为 ix0 )，它是一个二维数组，它告诉我每个子数组的起始位置每个 i/ j 位置。最后，我保证不会出现任何“溢出”——也就是说， ix0 + n 中的所有值都小于 a 的 2 维长度(因此我们不需要担心要求超出范围的索引)存在的范围)。

例如...

a = np.arange(74*74*4563).reshape(74, 74, 4563)
ix0 = np.arange(74*74).reshape(74,74)/2 + 50
a[:, :, ix0:ix0+n]

产生

IndexError: failed to coerce slice entry of type numpy.ndarray to integer

有没有办法做到这一点，而无需循环遍历所有 i/j 索引组合或创建大掩码数组？

Answer 1

以前已经问过类似的问题，但是是针对 2d 的。我可能会尝试查找一下。

但这里是 2d 情况下发生的情况的简单示例

In [1463]: x=np.arange(12).reshape(3,4)
In [1464]: ix0=np.array([0,2,1])
In [1465]: N=2

我们可以迭代 x 的每一行，收集所需的 N 长度切片，然后将它们连接到列表或数组中。一个更普遍的问题是切片的长度不同，在这种情况下它们无法重新组装成数组。

In [1466]: [x[i,ix0[i]:ix0[i]+N] for i in range(3)]
Out[1466]: [array([0, 1]), array([6, 7]), array([ 9, 10])]

然后将该列表包装在 np.array 中。

另一种方法是首先连接索引:

In [1467]: x[np.arange(3)[:,None], np.array([np.r_[ix0[i]:ix0[i]+N] for i in range(3)])]
Out[1467]: 
array([[ 0,  1],
       [ 6,  7],
       [ 9, 10]])

最后一个索引数组是:

In [1468]: np.array([np.r_[ix0[i]:ix0[i]+N] for i in range(3)])
Out[1468]: 
array([[0, 1],
       [2, 3],
       [1, 2]])

要应用于 3d 情况，我们有两种选择。一种是将其 reshape 为二维，应用其中一种策略，然后再 reshape 回来。另一个是概括我为创建这些而采取的行动 - 这应该不会太难，但需要进行一些实验。

通过广播创建最后一个数组应该不难。

In [1469]: ix0[:,None]+np.arange(N)
Out[1469]: 
array([[0, 1],
       [2, 3],
       [1, 2]])

In [1470]: x[np.arange(3)[:,None], ix0[:,None]+np.arange(N)]
Out[1470]: 
array([[ 0,  1],
       [ 6,  7],
       [ 9, 10]])

现在应该更容易推广到 3d

In [1487]: X=np.arange(2*3*10).reshape(2,3,10)

In [1488]: ix0=np.arange(2*3).reshape(2,3)

In [1489]: ix0[...,None]+np.arange(N)
Out[1489]: 
array([[[0, 1],
        [1, 2],
        [2, 3]],

       [[3, 4],
        [4, 5],
        [5, 6]]])

In [1490]: I,J,_=np.ix_(range(2),range(3),range(N))

In [1491]: I.shape
Out[1491]: (2, 1, 1)

In [1492]: J.shape
Out[1492]: (1, 3, 1)

In [1493]: X[I, J, ix0[...,None]+np.arange(N)]
Out[1493]: 
array([[[ 0,  1],
        [11, 12],
        [22, 23]],

       [[33, 34],
        [44, 45],
        [55, 56]]])

我应该确保值是正确的，但形状匹配，在这种情况下这是 80% 的工作。