从python中的数组中提取数组而不使用循环

Question

我正在尝试从数组中提取数组的一部分。

假设我有一个形状为(M, N, P)的数组array1。对于我的具体案例，M = 10、N = 5、P = 2000。我有另一个数组，形状为(M, N, 1)的array2，它包含array1中沿最后一个轴的感兴趣数据的起点。我想从array2给出的索引开始提取50个点的数据，有点像这样：

array1[:, :, array2:array2 + 50] 

我希望得到(M, N, 50)形状的结果。不幸的是，我得到了错误：

TypeError: only integer scalar arrays can be converted to a scalar index

当然，我也可以通过循环遍历数组来获得结果，但我觉得必须有一种更聪明的方法，因为我经常需要这种方法。

Answer 1

您可以使用array2中的值与最后一个维度的索引范围的比较来构建掩码：

例如：

import numpy as np

M,N,P,k = 4,2,15,3   # yours would be 10,5,2000,50
A1 = np.arange(M*N*P).reshape((M,N,P))
A2 = np.arange(M*N).reshape((M,N,1)) + 1
rP = np.arange(P)[None,None,:]
A3 = A1[(rP>=A2)&(rP<A2+k)].reshape((M,N,k))

输入：

print(A1)
[[[  0   1   2   3   4   5   6   7   8   9  10  11  12  13  14]
[ 15  16  17  18  19  20  21  22  23  24  25  26  27  28  29]]
[[ 30  31  32  33  34  35  36  37  38  39  40  41  42  43  44]
[ 45  46  47  48  49  50  51  52  53  54  55  56  57  58  59]]
[[ 60  61  62  63  64  65  66  67  68  69  70  71  72  73  74]
[ 75  76  77  78  79  80  81  82  83  84  85  86  87  88  89]]
[[ 90  91  92  93  94  95  96  97  98  99 100 101 102 103 104]
[105 106 107 108 109 110 111 112 113 114 115 116 117 118 119]]]
print(A2)
[[[1]
[2]]
[[3]
[4]]
[[5]
[6]]
[[7]
[8]]]

输出：

print(A3)
[[[  1   2   3]
[ 17  18  19]]
[[ 33  34  35]
[ 49  50  51]]
[[ 65  66  67]
[ 81  82  83]]
[[ 97  98  99]
[113 114 115]]]

Answer 2

由于每个位置的索引都不对齐，因此可以创建一个掩码或花式索引来提取所需的元素。由于提取的值将是平面的，因此必须对其进行整形。

以下是如何创建掩码：

K = 50
mask = np.zeros((M, N, P + 1), dtype=np.int8)
np.put_along_axis(mask, array2, 1, axis=-1)
np.put_along_axis(mask, array2 + K, -1, axis=-1)
mask.cumsum(axis=-1, out=mask)
mask = mask[..., :-1].view(bool)

我们使用np.int8和np.bool_具有相同的项目大小，并且np.cumsum在每个轴上将初始遮罩位置传播到最终遮罩位置。

剩下的相当简单：

array3 = array1[mask].reshape(M, N, K)

通过绕过np.put_along_axis并在适当的情况下使用带剪切的直接索引，可以在构建掩码时避免额外的元素：

mask = np.zeros_like(array1, dtype=np.int8)
r = np.tile(np.arange(M)[:, None, None], (1, N, 1))
c = np.tile(np.arange(N)[None, :, None], (M, 1, 1))
clip_mask = array2 + K < P
mask[r, c, array2] = 1
mask[r[clip_mask], c[clip_mask], array2[clip_mask] + K] = -1
mask = np.cumsum(mask, axis=-1, out=mask).view(bool)

这一切都是非常浪费的：要获得形状为(M, N, K)的数组，您需要创建一个大小为(M, N, P)的布尔掩码，以及一些大小为(M, N, 1)的索引数组，另一个大小是(M, N, 1)的掩码，然后创建这些索引数组的一些掩码版本。在这里使用for循环确实没有错，只要编译它们，例如使用cython或numba。