6

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

Prgmr
  • 129
  • 4
  • 11

3 Answers3

4

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))
6502
  • 112,025
  • 15
  • 165
  • 265
  • Which version of numpy are you using? With 1.13.1, the line `d[:] = values` results in `ValueError: assignment destination is read-only`. – Warren Weckesser Dec 05 '17 at 08:01
  • 1
    @WarrenWeckesser: you're completely right, for strange reasons `diag` is sharing data with the matrix, but is read-only. However it's trivial to use `ravel` to get the diagonal of a matrix with read-write access... – 6502 Dec 05 '17 at 08:18
  • With a 2-D array, `diag` calls [`diagonal`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.diagonal.html); see its docstring for an explanation of the read-only behavior. – Warren Weckesser Dec 05 '17 at 09:12
  • FYI: Your solution will work if `mat` is C-contiguous. It will not work with, for example, `mat = np.random.randint(1, 10, size=(5, 5)).T`, because `mat.ravel()` will be a copy. – Warren Weckesser Dec 05 '17 at 09:15
  • Instead, you could use `d = as_strided(mat, strides=(sum(mat.strides),), shape=(min(mat.shape),))`, where `as_strided` is imported with `from numpy.lib.stride_tricks import as_strided`. – Warren Weckesser Dec 05 '17 at 09:26
  • @WarrenWeckesser the new `einsum('ii->i', mat)` which I suppose does the same under the hood appears to be quite a bit faster. – Paul Panzer Dec 05 '17 at 09:38
  • @PaulPanzer Great! And it's nice to avoid the footgun potential of `as_strided`. – Warren Weckesser Dec 05 '17 at 09:52
3

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]])
Paul Panzer
  • 51,835
  • 3
  • 54
  • 99
  • This is neat! I wish TensorFlow tensors had `.einsum()` method too. – Massoud Nov 21 '19 at 18:05
  • 1
    We could also just do `a = a * np.eye(len(a))`. It should be the same I suppose or could there be any pitfalls for that? – Shaan May 14 '21 at 10:33
1

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
Dmytro Chasovskyi
  • 3,209
  • 4
  • 40
  • 82
Dev Chandan
  • 11
  • 1