3

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?

PhysLQ
  • 149
  • 2
  • 9

2 Answers2

3

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
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • Thanks. This works well. Is there any further way to speed this up, to avoid doing numerous reshapes? – PhysLQ May 16 '17 at 08:15
  • @PhysLQ Well, reshaping to create a view is the best way to have a viewed access to an array and reshape is virtually free. I don't see any other usage of reshape with the proposed method, or do we need reshapes elsewhere? – Divakar May 16 '17 at 10:06
2

If you do np.source(np.fill_diagonal) you'll see that in the 2d case it uses a 'strided' approach

    if a.ndim == 2:
        step = a.shape[1] + 1
        end = a.shape[1] * a.shape[1]
    a.flat[:end:step] = val

@Divakar's solution applies this to your 3d case by 'flattening' on 2 dimensions.

You could sum the columns with M.sum(axis=1). Though I vaguely recall some timings that found that einsum was actually a bit faster. sum is a little more conventional.

Someone has has asked for an ability to expand dimensions in einsum, but I don't think that will happen.

hpaulj
  • 221,503
  • 14
  • 230
  • 353