fill lower triangle (including diagonal) of a Numpy Array with -infinity - Python

Question

I have a 2D array as below.

[[1 2 3]
 [4 5 6]
 [7 8 9]]

I need it to be converted into:

[[-inf  2   3] 
 [-inf -inf 6]
 [-inf -inf -inf]]

i.e., fill the lower triangle, including diagonal, to -infinity. How to do it using python ? Please help.

Answer 1

Use np.tril_indices:

m[np.tril_indices(m.shape[0])] = -np.inf
print(m)

array([[-inf,   2.,   3.],
       [-inf, -inf,   6.],
       [-inf, -inf, -inf]])

Suggested by @Kevin, use:

m[np.tril_indices_from(m)] = -np.inf

Note: the dtype of your array must be float because np.inf is a float.