How to cleanly index numpy arrays with arrays (or anything else that supports addition so that it can be offset)

Question

The easiest way to explain my question may be with an example, so let me define some arrays:

>>> test = arange(25).reshape((5,5))
>>> test
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]])
>>> Xinds = array([1,2,3])
>>> Yinds = array([1,2,3])

Now, if I wanted the elements in rows 1, 2, and 3 and in column 0, I could go:

>>> test[Yinds,0]
array([ 5, 10, 15])

If I wanted the items in rows 1, 2, and 3 and all columns, I could go:

>>> test[Yinds, :]
array([[ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19]])

However, if I try to extend this to get the elements in rows 1, 2, and 3 and columns 1, 2, and 3, -- surprise! -- I instead get the elements in (1,1), (2,2), and (3,3)

>>> test[Yinds, Xinds]
array([ 6, 12, 18])

Instead of what I want:

>>> test[Yinds, :][:, Xinds]
array([[ 6,  7,  8],
       [11, 12, 13],
       [16, 17, 18]])
>>> test[1:4,1:4]
array([[ 6,  7,  8],
       [11, 12, 13],
       [16, 17, 18]])

I realize I could define a slice, but I want to be able to add an offset to the indices (e.g. Yinds+offset), and that can't be done with slices.

I could do something like

>>> xStart = 1
>>> xEnd   = 4
>>> yStart = 1
>>> yEnd   = 4

>>> offset = 1
>>> test[yStart+offset:yEnd+offset, xStart+offset:xEnd+offset]
...

or

>>> Xinds = array([1,4])
>>> Yinds = array([1,4])

>>> offset = 1
>>> test[slice(*(Yinds+offset)), slice(*(Xinds+offset))]
...

But neither is particular clean.

Monkey patching the addition operator in to slice doesn't seem to be an option, and inheriting from slice to add the operator doesn't appear to work either; I get the error, "type 'slice' is not an acceptable base type". (*Grumble* This wouldn't be a problem in Ruby *Grumble*)

So, my question is, what's the cleanest way to access a (more than 1 dimensional) sub-array with something that can be stored and offset?

Options so far:

• test[Yinds+offset, :][:, Xinds+offset]
• test[yStart+offset:yEnd+offset, xStart+offset:xEnd+offset]
• test[slice(*(Yinds+offset)), slice(*(Xinds+offset))]

Answer 1

I'm not entirely sure what you want, but maybe ix_ would help? I think I've seen people who know more about numpy than I do use it in similar contexts.

>>> from numpy import array, arange, ix_
>>> a = arange(25).reshape(5,5)
>>> Xinds = array([1,2,3])
>>> Yinds = array([1,2,3])
>>> 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]])
>>> a[ix_(Xinds, Yinds)]
array([[ 6,  7,  8],
       [11, 12, 13],
       [16, 17, 18]])
>>> a[ix_(Xinds+1, Yinds)]
array([[11, 12, 13],
       [16, 17, 18],
       [21, 22, 23]])
>>> Y2inds = array([1,3,4])
>>> a[ix_(Xinds, Y2inds)]
array([[ 6,  8,  9],
       [11, 13, 14],
       [16, 18, 19]])
>>> a[ix_(Xinds, Y2inds-1)]
array([[ 5,  7,  8],
       [10, 12, 13],
       [15, 17, 18]])