Python: Non diagonal elements of a matrix to 0

Question:

What is the quickest way to convert the non-diagonal elements of a square symmetrical numpy ndarray to 0?

Asked By: Prgmr

||

Answers:

I’d check out the speed of saving the diagonal away, then zap the matrix, then restore the diagonal:

n = len(mat)
d = mat.ravel()[::n+1]
values = d.copy()
mat[:,:] = 0
d[:] = values

if the matrix is not huge may be however that just allocating a new one is faster

mat = numpy.diag(numpy.diag(mat))
Answered By: 6502

Here is a solution that also works on non-contiguous arrays:

a = np.arange(110).reshape(10, 11)[:, :10]

diag = np.einsum('ii->i', a)
# or if a is not guaranteed to be square
# mn = min(a.shape)
# diag = np.einsum('ii->i', a[:mn, :mn])
save = diag.copy()
a[...] = 0
diag[...] = save

a

# array([[  0,   0,   0,   0,   0,   0,   0,   0,   0,   0],
#        [  0,  12,   0,   0,   0,   0,   0,   0,   0,   0],
#        [  0,   0,  24,   0,   0,   0,   0,   0,   0,   0],
#        [  0,   0,   0,  36,   0,   0,   0,   0,   0,   0],
#        [  0,   0,   0,   0,  48,   0,   0,   0,   0,   0],
#        [  0,   0,   0,   0,   0,  60,   0,   0,   0,   0],
#        [  0,   0,   0,   0,   0,   0,  72,   0,   0,   0],
#        [  0,   0,   0,   0,   0,   0,   0,  84,   0,   0],
#        [  0,   0,   0,   0,   0,   0,   0,   0,  96,   0],
#        [  0,   0,   0,   0,   0,   0,   0,   0,   0, 108]])
Answered By: Paul Panzer

There is a numpy method identity that puts ones on a matrix main diagonal. See the docs.

Input:

1 2
3 4

Example code:

P1 = np.array([[1,2],[3, 4]])

temp = np.identity(len(P1))

Ans = P1*temp

Output:

1 0
0 4
Answered By: Dev Chandan