How to add only to diagonals of array in Python?

Question

I have an Array on Python as follows:

array([[ 0.57733218,  0.09794384,  0.44497735],
       [ 0.87061284,  0.10253493,  0.56643557],
       [ 0.76358739,  0.44902046,  0.86064797]])

I want to add a scalar value of 20 to the diagonal of the array such that the output is:

array([[ 20.57733218,  0.09794384,  0.44497735],
       [ 0.87061284,  20.10253493,  0.56643557],
       [ 0.76358739,  0.44902046,  20.86064797]])

Since I may also be dealing with very large matrix arrays, what is the most efficient way to do this diagonal addition by an assignment operation as suggested in the accepted solution of this thread ?

Answer 1

One way would be to assign on flattened slice with an appropriate step-size -

In [233]: a
Out[233]: 
array([[ 0.57733218,  0.09794384,  0.44497735],
       [ 0.87061284,  0.10253493,  0.56643557],
       [ 0.76358739,  0.44902046,  0.86064797]])

In [234]: a.flat[::a.shape[1]+1] += 20

In [235]: a
Out[235]: 
array([[ 20.57733218,   0.09794384,   0.44497735],
       [  0.87061284,  20.10253493,   0.56643557],
       [  0.76358739,   0.44902046,  20.86064797]])

We can also use ndarray.ravel() to get the flattened view and then assign -

a.ravel()[::a.shape[1]+1] += 20

Another approach would be using np.einsum that gives us a view into the diagonal elements -

In [269]: a
Out[269]: 
array([[ 0.57733218,  0.09794384,  0.44497735],
       [ 0.87061284,  0.10253493,  0.56643557],
       [ 0.76358739,  0.44902046,  0.86064797]])

In [270]: d = np.einsum('ii->i', a)

In [271]: d += 20

In [272]: a
Out[272]: 
array([[ 20.57733218,   0.09794384,   0.44497735],
       [  0.87061284,  20.10253493,   0.56643557],
       [  0.76358739,   0.44902046,  20.86064797]])

Benchmarking

In [285]: a = np.random.rand(10000,10000)

# @Willem Van Onsem's soln
In [286]: %timeit np.fill_diagonal(a, a.diagonal() + 20)
10000 loops, best of 3: 159 µs per loop

In [287]: %timeit a.flat[::a.shape[1]+1] += 20
10000 loops, best of 3: 179 µs per loop

In [288]: %timeit a.ravel()[::a.shape[1]+1] += 20
100000 loops, best of 3: 18.2 µs per loop

In [289]: %%timeit
     ...: d = np.einsum('ii->i', a)
     ...: d += 20
100000 loops, best of 3: 18.5 µs per loop

Answer 2

Given that a is the array we want to update, we can make use of .diagonal() and np.fill_diagonal:

np.fill_diagonal(a, a.diagonal() + 20)

We thus first fetch the diagonal of a with a.diagonal(), then add 20 to every element of the diagonal. We use np.fill_diagonal(..) to set the elements of the diagonal.