indexing area in numpy 2D array, using a 2D array providing the indices

Question

Suppose I have a numpy 2D array and I want to index or extract a region of it using another 2D array which provides the index pairs, similar as in https://stackoverflow.com/a/14999290/17173241::

a = np.arange(36).reshape(6,6) 
>>> a
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]])

b = np.array([[0,1],[1,2],[2,1],[3,3],[4,4],[5,3]]) 
>>> b
array([[0, 1],
       [1, 2],
       [2, 1],
       [3, 3],
       [4, 4],
       [5, 3]])


>>> a[b[:,0],b[:,1]]
array([ 1,  8, 13, 21, 28, 33])

this works as intended. Now I want to use the second index of b as e.g. the mid point of a range, to extract not only single points in a for each row, but lets say 3 points together per row. How do I do it?

I tried

>>> a[b[:,0],b[:,1]-1:b[:,1]+1] 
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: only integer scalar arrays can be converted to a scalar index

but this gives a TypeError, which I cannot think through..

I came up with this attempt, because I thought it is similar to

>>> a[b[:,0],1:3] 
array([[ 1,  2],
       [ 7,  8],
       [13, 14],
       [19, 20],
       [25, 26],
       [31, 32]])

although of course the last example uses always the same indices as the second index.

I also tried this:

>>> a[b[:,0],(b[:,1]-1,b[:,1]+1)]
array([[ 0,  7, 12, 20, 27, 32],        
       [ 2,  9, 14, 22, 29, 34]])

yields the lower and upper bound of my range as it is read as a tuple, I guess. I want to achieve something like

array([[ 0,  7, 12, 20, 27, 32],
       [ 1,  8, 13, 21, 28, 33]
       [ 2,  9, 14, 22, 29, 34]])

or even better

array([[ 0, 1, 2],
       [ 7, 8, 9],
       [12, 13, 14],
       [20, 21, 22],
       [27, 28, 29],
       [32, 33, 34]])

BONUS: for a range that wide, that indices become negative i.e., b[:,1]-2:b[:,1]+1], how would I avoid getting values from the end of each row, but instead getting zeros.

Answer 1

I'll show how to index 2 items from each row of a (3,3) array. Hopefully it gives an idea that can be generalized. I need to clean up the copy-n-paste from the ipython qtconsole session, but here's a start

The wource array

x=np.arange(9).reshape(3,3)*10.0    
x
Out[3]: 
array([[ 0., 10., 20.],
       [30., 40., 50.],
       [60., 70., 80.]])

row index (3,1) shape

b1=np.arange(3)[:,None]

b1
Out[6]: 
array([[0],
       [1],
       [2]])

column index (3,2). Note that (3,1) broadcasts with (3,2):

b2=np.array([[0,1],[1,2],[0,2]])

b2
Out[9]: 
array([[0, 1],
       [1, 2],
       [0, 2]])

In [10]: x[b1,b2]
Out[10]: 
array([[ 0., 10.],
       [40., 50.],
       [60., 80.]])

Using linspace to generate b2 from ranges

In [14]: b2=np.linspace([0,1,1],[1,2,2],2).astype(int)

In [15]: b2
Out[15]: 
array([[0, 1, 1],
       [1, 2, 2]])

In [16]: x[b1,b2.T]
Out[16]: 
array([[ 0., 10.],
       [40., 50.],
       [70., 80.]])

The key is that b1 and b2 need to broadcast.

np.take_along_axis is a new function that can take care of creating the b1 array. Otherwise it functions the same way.

Answer 2

Your initial question can be solved with an inline for loop:

c = [a[i[0], (i[1]-1):(i[1]+2)] for i in b]

However, you can't access elements of this list using the c[a, b] method, you have to use c[a][b]