Shorter version of this numpy array indexing

Question

I have the following code in python (numpy array or scipy.sparse.matrices), it works:

X[a,:][:,b]

But it doesn't look elegant. 'a' and 'b' are 1-D boolean mask.

'a' has the same length as X.shape[0] and 'b' has the same length as X.shape[1]

I tried X[a,b] but it doesn't work.

What I am trying to accomplish is to select particular rows and columns at the same time. For example, select row 0,7,8 then from that result select all rows from column 2,3,4

How would you make this shorter and more elegant?

X is a scipy sparse matrix but I think the indexing is the same as in the dense numpy array matrix. — off99555, Jun 18 '16 at 14:22
Well NumPy arrays are a different thing than NumPy matrices. See here : http://stackoverflow.com/questions/4151128/what-are-the-differences-between-numpy-arrays-and-matrices-which-one-should-i-u — Divakar, Jun 18 '16 at 14:25
With a sparse matrix it may be faster to 'index' with matrix multiplication, e.g. `ix*M*jx`. The trick is to construct the indexing matrices fast enough. — hpaulj, Jun 18 '16 at 17:14

Divakar · Accepted Answer · 2016-06-18T15:10:40.987

You could use np.ix_ for such a broadcasted indexing, like so -

X[np.ix_(a,b)]

Though this won't be any shorter than the original code, but hopefully should be faster. This is because we are avoiding the intermediate output as with the original code that created X[a,:] with one slicing and then another slicing X[a,:][:,b] to give us the final output.

Also, this method would work for a and b as both int and boolean arrays.

Sample run

In [141]: X = np.random.randint(0,99,(6,5))

In [142]: m,n = X.shape

In [143]: a = np.in1d(np.arange(m),np.random.randint(0,m,(m)))

In [144]: b = np.in1d(np.arange(n),np.random.randint(0,n,(n)))

In [145]: X[a,:][:,b]
Out[145]: 
array([[17, 81, 64],
       [87, 16, 54],
       [98, 22, 11],
       [26, 54, 64]])

In [146]: X[np.ix_(a,b)]
Out[146]: 
array([[17, 81, 64],
       [87, 16, 54],
       [98, 22, 11],
       [26, 54, 64]])

Runtime test

In [147]: X = np.random.randint(0,99,(600,500))

In [148]: m,n = X.shape

In [149]: a = np.in1d(np.arange(m),np.random.randint(0,m,(m)))

In [150]: b = np.in1d(np.arange(n),np.random.randint(0,n,(n)))

In [151]: %timeit X[a,:][:,b]
1000 loops, best of 3: 1.74 ms per loop

In [152]: %timeit X[np.ix_(a,b)]
1000 loops, best of 3: 1.24 ms per loop

Look at the output of `ix_` to see how it can be done without it. — hpaulj, Jun 18 '16 at 15:32
`np.ix_` applies `np.nonzero` to boolean arguments, converting them to indices. — hpaulj, Jun 18 '16 at 16:49

Shorter version of this numpy array indexing

1 Answers1