I have an ndarray:
ndaM = np.array([[1,2,3],[4,5,6],[7,8,9]]);
and a vector of 0s and 1s:
vSet = [0,1,1];
How do I use vSet to extract the submatrix with only rows and columns whose index is nonzero in vSet? In this case, the matrix
[[5,6],[8,9]]
Here is a solution that slices in one go using np.ix_
ndaM[np.ix_(*2*(np.array(vSet, bool),))]
# array([[5, 6],
# [8, 9]])
Or in more readable two lines
mask = np.array(vSet, bool)
ndaM[np.ix_(mask, mask)]
np.ix_ applies np.nonzero (np.where) to a boolean input to turn it into a numeric index. (array([[1], [2]]), array([[1, 2]])). `
You can go with a two-step slicing:
# convert vSet to a boolean array
bSet = np.array(vSet).astype(bool)
# slice in two steps, rows first, then columns
ndaM[bSet][:, bSet]
#array([[5, 6],
# [8, 9]])
Or use np.ix_ to create an index mesh grid and then use it for indexing:
ndaM[np.ix_(bSet, bSet)]
#array([[5, 6],
# [8, 9]])