2

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

假设我有一个array1shape数组(M, N, P)。对于我的具体情况，M = 10, N = 5, P = 2000. 我有另一个array2shape数组，其中包含沿最后一个轴(M, N, 1)的有趣数据的起点。array1我想从 给出的索引开始提取这些数据的 50 个点array2，有点像这样：

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

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

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

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

4

2 回答 2

2

您可以使用 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]]]
于 2021-12-02T19:03:31.367 回答
1

由于您在每个位置的索引未对齐，因此您可以创建一个掩码或花哨的索引来提取所需的元素。由于提取的值将是平坦的，因此您必须重新调整它们。

以下是创建蒙版的方法：

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)

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

于 2021-12-02T17:32:17.900 回答