Question

Given a numpy matrix A of shape (m,n) and a list ind of Boolean values with length n, I want to extract the columns of A where the corresponding element of the Boolean list is true.

My first naive try

Asub = A[:,cols]

gives quite weird results, not necessary to cite here…

Following pv.’s answer to this question, I tried with numpy.ix_ as follows:

>>> A = numpy.diag([1,2,3])
>>> ind = [True, True, False]
>>> A[:,numpy.ix_(ind)]
array([[[1, 0]],

       [[0, 2]],

       [[0, 0]]])
>>> A[:,numpy.ix_(ind)].shape
(3, 1, 2)

Now the result is of inappropriate shape: I want to have a (3,2) resulting array. I guess I could collapse the result down to two dimensions, but there surely must be a better way of doing this?

Answer 1

As the docs discuss, boolean indexing like you want “occurs when obj is an array object of Boolean type (such as may be returned from comparison operators).”

IOW, ind needs to be an ndarray of type bool:

In [15]: A = numpy.diag([1,2,3])

In [16]: ind = [True, True, False]

In [17]: A[:,ind]
Out[17]: 
array([[0, 0, 1],
       [2, 2, 0],
       [0, 0, 0]])

which occurs because it’s interpreting the bools as integers, and giving you columns [1, 1, 0].

OTOH:

In [18]: A[:,numpy.array(ind)]
Out[18]: 
array([[1, 0],
       [0, 2],
       [0, 0]])