how to set diagonal of volume of matrices in numpy?

Question

I have a volume of matrices, say a = np.zeros((5, 7, 4, 4))

I want to set the diagonal to some value, across all matrices.

Just like I can get the diagonal with diag = np.diagonal(a, axis1=2, axis2=3), I would like to set the diagonal.

However, np.fill_diagonal does not support axis=.

How can I still achieve this?

Can you do it with loops, or are you looking for a built-in solution? — ofey, Jan 04 '22 at 13:12
`fill_diagonal` works on a 4d array. Transpose your axes first. — hpaulj, Jan 04 '22 at 16:27

ofey · Answer 1 · 2022-01-04T17:35:31.827

0

import numpy as np

a = np.zeros((5, 7, 4, 4))
s0, s1, s2, s3 = a.shape

for index in range(s0):
    a[index].reshape(s1,-1)[:,::s3+1] = 3 # assign the diagonal values to 3

Taken from fast way to set diagonals of an M x N x N matrix with a little change. I believe this does what you want.

Edit: Aha. As @hpaulj's suggested, we can do it without a loop as you've desired.

a = np.zeros((5, 7, 4, 4))
s0, s1, s2, s3 = a.shape
a.reshape(s0 * s1, -1)[:,::s3+1] = 3 # assign the diagonal values to 3

edited Jan 04 '22 at 17:35

answered Jan 04 '22 at 13:36

ofey

135
8

1

you should be able to apply the linked answer to a 4d without the loop. – hpaulj Jan 04 '22 at 16:32
@hpaulj Thanks for the idea. Just 'flattening' it even further works better :) – ofey Jan 04 '22 at 17:41

score 0 · Accepted Answer · answered Jan 04 '22 at 16:40

0

You can do this:

s1, s2, s3, s4 = a.shape
i, j = np.diag_indices(s3)  # s3 == s4 == 4 here.
a[..., i, j] = 3

answered Jan 04 '22 at 16:40

swag2198

2,546
1
7
18

how to set diagonal of volume of matrices in numpy?

2 Answers2