Reputation: 45
For example, If I have Numpy arrays which is initialised by:
a = np.arange(12).reshape(6,2)
[out] array([[ 0, 1],
[ 2, 3],
[ 4, 5],
[ 6, 7],
[ 8, 9],
[10, 11]])
and
mask = np.array([0, 2])
My target is to mask array by range in a axis. like this
for i in mask:
target.append(a[i:i+3,:])
So, it should be:
[out] array([[[0, 1],
[2, 3],
[4, 5]],
[[4, 5],
[6, 7],
[8, 9]]])
but that's inefficient. Then, I've tried
a[mask:mask+3,:]
but it said
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
TypeError: only integer scalar arrays can be converted to a scalar index
Upvotes: 2
Views: 193
Reputation: 221774
Approach #1
We could leverage broadcasting
to generate all indices and index -
In [19]: a
Out[19]:
array([[ 0, 1],
[ 2, 3],
[ 4, 5],
[ 6, 7],
[ 8, 9],
[10, 11]])
In [21]: mask
Out[21]: array([0, 2])
In [24]: a[mask[:,None] + np.arange(3)]
Out[24]:
array([[[0, 1],
[2, 3],
[4, 5]],
[[4, 5],
[6, 7],
[8, 9]]])
Approach #2
We can also leverage np.lib.stride_tricks.as_strided
based scikit-image's view_as_windows
for a more efficient solution -
In [43]: from skimage.util.shape import view_as_windows
In [44]: view_as_windows(a,(3,a.shape[1]))[mask][:,0]
Out[44]:
array([[[0, 1],
[2, 3],
[4, 5]],
[[4, 5],
[6, 7],
[8, 9]]])
Upvotes: 2