Buddyshot Buddyshot - 19 days ago 7
Python Question

How to select columns of a numpy matrix based on a 1-D boolean mask?

Let us consider the matrix

A


[[1, 0, 1, 0, 0, 0],
[1, 0, 0, 0, 0, 0],
[1, 0, 1, 0, 0, 0],
[1, 0, 1, 0, 0, 1],
[0, 0, 1, 0, 0, 1],
[0, 0, 1, 0, 0, 1]]


We want to identify and remove all the columns where each element is 0. We can generate a mask such as

mask = np.all(A == 0, axis=0)
# output: [False, True, False, True, True, False]


How do I use
mask
(or
~mask
) to create a copy of
A
where only the non-null columns are kept? That is

[[1, 1, 0],
[1, 0, 0],
[1, 1, 0],
[1, 1, 1],
[0, 1, 1],
[0, 1, 1]]

Answer
>>> import numpy as np
>>> A = np.array([[1, 0, 1, 0, 0, 0],
                  [1, 0, 0, 0, 0, 0],
                  [1, 0, 1, 0, 0, 0],
                  [1, 0, 1, 0, 0, 1],
                  [0, 0, 1, 0, 0, 1],
                  [0, 0, 1, 0, 0, 1]])

>>> mask = np.all(A == 0, axis=0)
>>> mask
array([False,  True, False,  True,  True, False], dtype=bool)

>>> A[:,mask]
array([[0, 0, 0],
       [0, 0, 0],
       [0, 0, 0],
       [0, 0, 0],
       [0, 0, 0],
       [0, 0, 0]])

>>> A[:,~mask]
array([[1, 1, 0],
       [1, 0, 0],
       [1, 1, 0],
       [1, 1, 1],
       [0, 1, 1],
       [0, 1, 1]])