CODEWITHSUNDEEP
pythonarraysnumpyvectorizationslice
CalvintheDog
CalvintheDog

Reputation: 11

Slicing a 2D numpy array using vectors for start-stop indices

First post here, so please go easy on me. :)

I want to vectorize the following:

rowStart=array of length N        
rowStop=rowStart+4        
colStart=array of length N    
colStop=colStart+4    

x=np.random.rand(512,512) #dummy test array   

output=np.zeros([N,4,4])
for i in range(N):
   output[i,:,:]=x[ rowStart[i]:rowStop[i], colStart[i]:colStop[i] ]

What I'd like to be able to do is something like:

output=x[rowStart:rowStop, colStart:colStop ]

where numpy recognizes that the slicing indices are vectors and broadcasts the slicing. I understand that this probably doesn't work because while I know that my slice output is always the same size, numpy doesn't.

I've looked at various approaches, including "fancy" or "advanced" indexing (which seems to work for indexing, not slicing), massive boolean indexing using meshgrids (not practical from a memory standpoint, as my N can get to 50k-100k), and np.take, which just seems to be another way of doing fancy/advanced indexing.

I could see how I could potentially use fancy/advanced indexing if I could get an array that looks like:

[np.arange(rowStart[0],rowStop[0]),
 np.arange(rowStart[1],rowStop[1]), 
...,  
np.arange(rowStart[N],rowStop[N])]

and a similar one for columns, but I'm also having trouble figuring out a vectorized approach for creating that.

I'd appreciate any advice you can provide. Thanks!

Upvotes: 1

Views: 206

Answers (1)

Divakar
Divakar

Reputation: 221514

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

from skimage.util.shape import view_as_windows

BSZ = (4, 4) # block size
w = view_as_windows(x, BSZ)
out = w[rowStart, colStart]

Upvotes: 1

Related Questions