Fast compose slices of 3D numpy array as 3D subarray

Question

I have a 3D numpy array x. I want to take a subset of each slice on axis 0 (each subset is the same shape, but with start and end indices that may be different for each slice) and compose these into a separate 3D numpy array. I can achieve this with

import numpy as np

x = np.arange(24).reshape((3, 4, 2))
starts = [0, 2, 1]
ends = [2, 4, 3]

np.stack([x[i, starts[i]:ends[i]] for i in range(3)])

but 1) is there any way to do this in a single operation using fancy indexing, and 2) will doing so speed things up?

Would the difference between starts and ends be always same across all elements, like its 2 here? — Divakar
– Divakar, Commented Sep 30, 2018 at 17:17

Divakar · Accepted Answer · 2018-09-30 17:39:00Z

1

We can leverage np.lib.stride_tricks.as_strided based scikit-image's view_as_windows to get sliding windows. More info on use of as_strided based view_as_windows.

from skimage.util.shape import view_as_windows

L = 2 # ends[0]-starts[0]
w = view_as_windows(x,(1,L,1))[...,0,:,0]
out = w[np.arange(len(starts)), starts].swapaxes(1,2)

Alternatively, a compact version leveraging broadcasting that generates all the required indices and then indexing into the input array, would be -

x[np.arange(len(starts))[:,None],np.asarray(starts)[:,None] + np.arange(L)]

answered Sep 30, 2018 at 17:39

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Fast compose slices of 3D numpy array as 3D subarray

1 Answer 1

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related