Numpy indexing by range of arrays

Question

Say I have an array myarr such that myarr.shape = (2,64,64,2). Now if I define myarr2 = myarr[[0,1,0,0,1],...], then the following is true

myarr2.shape #(5,64,64,2)
myarr2[0,...] == myarr[0,...] # = True
myarr2[1,...] == myarr[1,...] # = True
myarr2[2,...] == myarr[0,...] # = True
...

Can this be generalized so the slices are arrays? That is, is there a way to make the following hypothetical code work?

myarr2 = myarr[...,[20,30,40]:[30,40,50],[15,25,35]:[25,35,45],..]
myarr2[0,] == myarr[...,20:30,15:25,...] # = True
myarr2[1,] == myarr[...,30:40,25:35,...] # = True
myarr2[2,] == myarr[...,40:50,35:45,...] # = True

Assuming the slices are all the same size, making a new array by `concatenating` (with `np.array` or `np.stack`) the individual slices is the most straight forward approach. I can also picture using `np.ix_` and `np.r_` to make indexing arrays to do the same thing, but speed would be similar. — hpaulj, Jan 25 '20 at 18:10
See my answer to https://stackoverflow.com/questions/59914256/slicing-a-different-range-at-each-index-of-a-multidimensional-numpy-array — hpaulj, Jan 26 '20 at 00:14

score 0 · Answer 1 · answered Jan 25 '20 at 10:17

you may feed the coordinates of subarrays to the cycle which cuts subarrays from myarray. I don't know hoe you store the indices of subarrays so I put them into nested list idx_list:

idx_list = [[[20,30,40],[30,40,50]],[[15,25,35]:[25,35,45]]]  # assuming 2D cutouts
idx_array = np.array([k for i in idx_list for j in i for k in j]) # unpack
idx_array = idx_array .reshape(-1,2).T  # reshape
myarray2 = np.array([myarray[a:b,c:d] for a,b,c,d in i2])  # cut and combine

hpaulj · Answer 2 · 2020-01-26T03:50:20.423

Let's simplify the problem a bit; first by removing the two outer dimensions that don't affect the core indexing issue; and by reducing the size so we can see and understand the results.

The setup

In [540]: arr = np.arange(7*7).reshape(7,7)                                                      
In [541]: arr                                                                                    
Out[541]: 
array([[ 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]])
In [542]: idx =np.array([[0,2,4,6],[1,3,5,7]])

Now a straightforward iteration approach:

In [543]: alist = [] 
     ...: for i in range(idx.shape[1]-1): 
     ...:     j,k = idx[:,i] 
     ...:     sub = arr[j:j+2, k:k+2] 
     ...:     alist.append(sub) 
     ...:                                                                                        
In [544]: np.array(alist)                                                                        
Out[544]: 
array([[[ 1,  2],
        [ 8,  9]],

       [[17, 18],
        [24, 25]],

       [[33, 34],
        [40, 41]]])
In [545]: _.shape                                                                                
Out[545]: (3, 2, 2)

I simplified the iteration from:

 ...: for i in range(idx.shape[1]-1): 
 ...:     sub = arr[idx[0,i]:idx[0,i+1],idx[1,i]:idx[1,i+1]] 
 ...:     alist.append(sub)

to highlight the fact that we are generating ranges of a consistent size, and make the next transformation more obvious.

So I start with a (7,7) array, and create 3 (2,2) slices.

As I demonstrated in Slicing a different range at each index of a multidimensional numpy array, we can use linspace to expand a set of slices, or ranges.

In [567]: ranges = np.linspace(idx[:,:3],idx[:,:3]+1,2).astype(int)                              
In [568]: ranges                                                                                 
Out[568]: 
array([[[0, 2, 4],
        [1, 3, 5]],

       [[1, 3, 5],
        [2, 4, 6]]])

So ranges[0] expands on the idx[0] slices, etc. But if I simply index with these I get 'diagonal' values from Out[554]:

In [569]: arr[ranges[0], ranges[1]]                                                              
Out[569]: 
array([[ 1, 17, 33],
       [ 9, 25, 41]])

to get blocks I have to add a dimension to the first indices:

In [570]: arr[ranges[0,:,None], ranges[1]]                                                       
Out[570]: 
array([[[ 1, 17, 33],
        [ 2, 18, 34]],

       [[ 8, 24, 40],
        [ 9, 25, 41]]])

these are the same values as in Out[554], but need to be transposed:

In [571]: _.transpose(2,0,1)                                                                     
Out[571]: 
array([[[ 1,  2],
        [ 8,  9]],

       [[17, 18],
        [24, 25]],

       [[33, 34],
        [40, 41]]])

The code's a bit clunky and needs to get generalized, but gives the general idea of how one can substitute one indexing for the iterative one, provide the slices are regular enough. For this small example it probably isn't faster, but it probably will come ahead as the problem size gets larger.

Numpy indexing by range of arrays

2 Answers2