how to set diagonal of volume of matrices in numpy?-CodePudding

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?

CodePudding user response：

You can do this:

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

CodePudding user response：

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