[SOLVED] Fast way to set diagonals of an (M x N x N) matrix? Einsum / n-dimensional fill

Fast way to set diagonals of an (M x N x N) matrix? Einsum / n-dimensional fill_diagonal?

I'm trying to write fast, optimized code based on matrices, and have recently discovered einsum as a tool for achieving significant speed-up.

Is it possible to use this to set the diagonals of a multidimensional array efficiently, or can it only return data?

In my problem, I'm trying to set the diagonals for an array of square matrices (shape: M x N x N) by summing the columns in each square (N x N) matrix.

My current (slow, loop-based) solution is:

# Build dummy array
dimx = 2  # Dimension x (likely to be < 100)
dimy = 3  # Dimension y (likely to be between 2 and 10)
M = np.random.randint(low=1, high=9, size=[dimx, dimy, dimy])

# Blank the diagonals so we can see the intended effect
np.fill_diagonal(M[0], 0)
np.fill_diagonal(M[1], 0)

# Compute diagonals based on summing columns
diags = np.einsum('ijk->ik', M)

# Set the diagonal for each matrix 
# THIS IS LOW. CAN IT BE IMPROVED?
for i in range(len(M)):
    np.fill_diagonal(M[i], diags[i])

# Print result   
M

Can this be improved at all please? It seems np.fill_diagonal doesn't accepted non-square matrices (hence forcing my loop based solution). Perhaps einsum can help here too?

Solution

One approach would be to reshape to 2D, set the columns at steps of ncols+1 with the diagonal values. Reshaping creates a view and as such allows us to directly access those diagonal positions. Thus, the implementation would be -

s0,s1,s2 = M.shape
M.reshape(s0,-1)[:,::s2+1] = diags