mask = np.tril(np.ones(3, dtype=bool)
print mask
[[ True False False]
[ True True False]
[ True True True]]
B = np.zeros(9)
B.shape = (3,3)
print B
[[ 0 0 0 ]
[ 0 0 0 ]
[ 0 0 0 ]]
B[mask]
array([0,0,0,0,0,0])
C = np.array([[1],[0],[0],[1],[0],[1]])
B[mask] = C
ValueError: boolean index array should have 1 dimension
I tried to apply .flatten():
B[mask] = C.flatten()
print B
array([[1, 0, 0],
[0, 0, 0],
[1, 0, 1]])
But my intended result is a diagonal matrix.
array([[1, 0, 0],
[0, 1, 0],
[0, 0, 1]])
What am I doing wrong?
The issue is you are assuming column-major ordered values, which is not the case with NumPy/Python. So, for a general case of using a mask to assign values with the column-major ordered values, we need to transpose both the input array and mask and assign those, like so -
B.T[mask.T] = C.flatten()
Sample run for explanation on getting the right order and assigning -
In [36]: B = np.arange(1,10).reshape(3,3)
In [37]: B
Out[37]:
array([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
In [38]: mask = np.tril(np.ones(3, dtype=bool))
In [39]: mask
Out[39]:
array([[ True, False, False],
[ True, True, False],
[ True, True, True]])
In [40]: B.T[mask.T]
Out[40]: array([1, 4, 7, 5, 8, 9]) # right order (col major) obtained
# Finally assign into masked positions
In [41]: B.T[mask.T] = C.flatten()
In [42]: B
Out[42]:
array([[1, 2, 3],
[0, 1, 6],
[0, 0, 1]])