python - Filling the diagonals of square matrices inside a 3D ndarray with values given by a 2D ndarray

Question

Given a 3D ndarray z with shape (k,n,n), is it possbile without using iteration to fill the diagonals of the k nxn matrices with values given by a 2D ndarray v with shape (k,n)?

For example, the result of the operation should be the same as looping over k matrices:

z = np.zeros((3,10,10))
v = np.arange(30).reshape((3,10))

for i in range(len(z)): 
    np.fill_diagonal(z[i], v[i])

Is there an way to do this without repeatedly calling np.fill_diagonal inside a loop? If possible, I would prefer a solution that can be applied to arrays of higher dimensions as well, where z.shape == (a,b,c,...,k,n,n) and v.shape = (a,b,c,...,k,n)

Answer 1

Here's for generic n-dim arrays -

diag_view = np.einsum('...ii->...i',z)
diag_view[:] = v

Another with reshaping -

n = v.shape[-1] 
z.reshape(-1,n**2)[:,::n+1] = v.reshape(-1,n)
# or z.reshape(z.shape[:-2]+(-1,))[...,::n+1] = v

Another with masking -

m = np.eye(n, dtype=bool) # n = v.shape[-1] from earlier
z[...,m] = v

Initializing output z

If we need to initialize the output array z and one that will cover for generic n-dim cases, it would be :

z = np.zeros(v.shape + (v.shape[-1],), dtype=v.dtype)

Then, we proceed with the earlier listed approaches.