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.

Answer 3

You can also create an diagonal matrix with np.eye and multiply it with suitable constant c.

np.eye(n, dtype=int) * c   # n x n matrix with diagonal part being c

Then, add it to your matrix a

a =  np.array([[ 0.57733218,  0.09794384,  0.44497735],
               [ 0.87061284,  0.10253493,  0.56643557],
               [ 0.76358739,  0.44902046,  0.86064797]])

a += np.eye(3) * 20 

Or perhaps make a mask with np.eye(3, dtype=bool) such that you can use it to select the diagonal part and add on it like

a[np.eye(3, dtype=bool)] += 20

Thanks for Divakar for the suggestion for using the mask and Willem Van Onsem's advice for making the answer clear.

Edit: both Willem and Divakar's methods are a lot faster than this when the size of data is big.

Answer 4

Noticing that your nD-array is square (or cuboid in general case), you can extract the indices of the diagonal elements np.diag_indices_from, get the original values of diagonal elements by indexing into the array, then add your desired constant; And then update the original diagonal values by indexing into the array using the indices that we get from np.diag_indices like:

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

    In [29]: new_diag_vals = arr[np.diag_indices_from(arr)] + 20

    In [30]: arr[np.diag_indices(arr.shape[0])] = new_diag_vals

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

---

If you want to do it in one line, then:

In [36]: arr[np.diag_indices(arr.shape[0])] =  arr[np.diag_indices_from(arr)] + 20

P.S. Note that this modifies the original array.