pythonnumpynumpy-ndarraydiagonal

Numpy change diagonal values in a 3D array

I have a vector with the following values: dv = array([0., 0., 1.]).

How do I diagonalize this vector into a 3D array with each element in the vector have its own diagonal:

array([[[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]],

       [[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]],

       [[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]]])

So far i have tried:

import numpy as np

a = np.zeros((3,3,3))
di = np.diag_indices(3,3)
a[di] = dv

This is almost correct, but it does not change all of the elements on the diagonal.

Solution

  • Try this -

    a = np.zeros((3,3,3))
dv = np.array([0, 0, 1])

i,j = np.diag_indices(3) #diagonal indices for a 2D matrix (3,3)
a[:,i,j] = dv[:,None]
a

    array([[[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]],

       [[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]],

       [[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]]])

    You basically get 2D matrix diagonals and use those to get a view of the 3 diagonals from the original matrix. You then repeat the diagonal values you want to update to the same shape by broadcasting and map it to the diagonals.

    The above approach is based on fetching a view from the original numpy array and then using assignment to update it. If you are not trying to do an assignment task, you can simply use arr.diagonal with axis1 and axis2 parameters to get a copy of the diagonal values for those 2 axes. Note, the shape of those axes must be equal (square matrix)

    a = np.arange(0,2*3*3).reshape(2,3,3)

# array([[[ 0,  1,  2],
#         [ 3,  4,  5],
#         [ 6,  7,  8]],

#        [[ 9, 10, 11],
#         [12, 13, 14],
#         [15, 16, 17]]])

a.diagonal(axis1=1, axis2=2)

# array([[ 0,  4,  8],
#        [ 9, 13, 17]])