What's the best way to set the diagonals of a particular dimension to 1?

user_222151 0
Viewed 34

Let's say I have a tensor:

q = np.arange(5*3*3).reshape(5,3,3)

I want to set the 3x3 diagonals to be 1, across axis 0 (i.e. where j=k).

I thought this should do it:

np.apply_along_axis(lambda x: np.fill_diagonal(x,1), 0, q)

but it doesn't seem to work.

What's the cleanest way to do this?

numpy python
2 Answers

Using the same index array in both dimensions selects a diagonal:

In [13]: q = np.arange(5*3*3).reshape(5,3,3)
In [14]: i=np.arange(3)
In [15]: q[:,i,i]
Out[15]: 
array([[ 0,  4,  8],
       [ 9, 13, 17],
       [18, 22, 26],
       [27, 31, 35],
       [36, 40, 44]])
In [16]: q[:,i,i]=1
In [17]: q
Out[17]: 
array([[[ 1,  1,  2],
        [ 3,  1,  5],
        [ 6,  7,  1]],

       [[ 1, 10, 11],
        [12,  1, 14],
        [15, 16,  1]],

       [[ 1, 19, 20],
        [21,  1, 23],
        [24, 25,  1]],

       [[ 1, 28, 29],
        [30,  1, 32],
        [33, 34,  1]],

       [[ 1, 37, 38],
        [39,  1, 41],
        [42, 43,  1]]])
user_9019250

this is not the cleanest way but this is one way:

q[np.full((5,3,3), np.identity(3)) == 1] = 1
user_82304040
Related