Subsetting a 2D numpy array

You've gotten a handful of nice examples of how to do what you want. However, it's also useful to understand the what's happening and why things work the way they do. There are a few simple rules that will help you in the future.

There's a big difference between "fancy" indexing (i.e. using a list/sequence) and "normal" indexing (using a slice). The underlying reason has to do with whether or not the array can be "regularly strided", and therefore whether or not a copy needs to be made. Arbitrary sequences therefore have to be treated differently, if we want to be able to create "views" without making copies.

In your case:

import numpy as npa = np.arange(100).reshape(10,10)n1, n2 = np.arange(5), np.arange(5)# Not what you wantb = a[n1, n2]  # array([ 0, 11, 22, 33, 44])# What you want, but only for simple sequences# Note that no copy of *a* is made!! This is a view.b = a[:5, :5]# What you want, but probably confusing at first. (Also, makes a copy.)# np.meshgrid and np.ix_ are basically equivalent to this.b = a[n1[:,None], n2[None,:]]

Fancy indexing with 1D sequences is basically equivalent to zipping them together and indexing with the result.

print "Fancy Indexing:"print a[n1, n2]print "Manual indexing:"for i, j in zip(n1, n2):    print a[i, j]

However, if the sequences you're indexing with match the dimensionality of the array you're indexing (2D, in this case), The indexing is treated differently. Instead of "zipping the two together", numpy uses the indices like a mask.

In other words, a[[[1, 2, 3]], [[1],[2],[3]]] is treated completely differently than a[[1, 2, 3], [1, 2, 3]], because the sequences/arrays that you're passing in are two-dimensional.

In [4]: a[[[1, 2, 3]], [[1],[2],[3]]]Out[4]:array([[11, 21, 31],       [12, 22, 32],       [13, 23, 33]])In [5]: a[[1, 2, 3], [1, 2, 3]]Out[5]: array([11, 22, 33])

To be a bit more precise,

a[[[1, 2, 3]], [[1],[2],[3]]]

is treated exactly like:

i = [[1, 1, 1],     [2, 2, 2],     [3, 3, 3]])j = [[1, 2, 3],     [1, 2, 3],     [1, 2, 3]]a[i, j]

In other words, whether the input is a row/column vector is a shorthand for how the indices should repeat in the indexing.

np.meshgrid and np.ix_ are just convienent ways to turn your 1D sequences into their 2D versions for indexing:

In [6]: np.ix_([1, 2, 3], [1, 2, 3])Out[6]:(array([[1],       [2],       [3]]), array([[1, 2, 3]]))

Similarly (the sparse argument would make it identical to ix_ above):

In [7]: np.meshgrid([1, 2, 3], [1, 2, 3], indexing='ij')Out[7]:[array([[1, 1, 1],       [2, 2, 2],       [3, 3, 3]]), array([[1, 2, 3],       [1, 2, 3],       [1, 2, 3]])]

Another quick way to build the desired index is to use the np.ix_ function:

>>> a[np.ix_(n1, n2)]array([[ 0,  1,  2,  3,  4],       [10, 11, 12, 13, 14],       [20, 21, 22, 23, 24],       [30, 31, 32, 33, 34],       [40, 41, 42, 43, 44]])

This provides a convenient way to construct an open mesh from sequences of indices.

You could use np.meshgrid to give the n1, n2 arrays the proper shape to perform the desired indexing:

In [104]: a[np.meshgrid(n1,n2, sparse=True, indexing='ij')]Out[104]: array([[ 0,  1,  2,  3,  4],       [10, 11, 12, 13, 14],       [20, 21, 22, 23, 24],       [30, 31, 32, 33, 34],       [40, 41, 42, 43, 44]])

Or, without meshgrid:

In [117]: a[np.array(n1)[:,np.newaxis], np.array(n2)[np.newaxis,:]]Out[117]: array([[ 0,  1,  2,  3,  4],       [10, 11, 12, 13, 14],       [20, 21, 22, 23, 24],       [30, 31, 32, 33, 34],       [40, 41, 42, 43, 44]])

There is a similar example with an explanation of how this integer array indexing works in the docs.

See also the Cookbook recipe Picking out rows and columns.