numpy-masking.py
import numpy as np
# create array with 16 integers
b = np.arange(16)
b.shape = (4, 4)
# [[ 0 1 2 3]
# [ 4 5 6 7]
# [ 8 9 10 11]
# [12 13 14 15]]
# create mask with True
mask = np.ones(4, dtype=bool)
# [ True True True True]
# get diagonal
b[mask, mask]
# [ 0 5 10 15]
# create mask with both True and False
another_mask = np.array([[1, 0, 1, 1], [False, True, False, False], [0, 0, 1, 0], [0, 1, 1, 1]], dtype=bool)
# [[ True False True True]
# [False True False False]
# [False False True False]
# [False True True True]]
b[another_mask]
# [ 0 2 3 5 10 13 14 15]
# create mask with conditional
conditional_mask = (b >= 9)
b[conditional_mask]
# [ 9 10 11 12 13 14 15]
# get all values smaller than 9
b[b < 9]
# [0 1 2 3 4 5 6 7 8]
# get all values smaller than 5 or larger than 10
b[(b < 5) | (b > 10)]
# [ 0 1 2 3 4 11 12 13 14 15]
# where
b[np.where(b > 10)]
# [11 12 13 14 15]
# get all columns with a value less than 3
b[:, np.where(b < 3)[1]]
# [[ 0 1 2]
# [ 4 5 6]
# [ 8 9 10]
# [12 13 14]]