Multiplying matrices (dot product) along an axis with Numpy (without loops)
Asked Answered
R

1

7

I am working with Numpy on an image processing problem, and I am trying to avoid loops and do the following:

I have a matrix M of dims NxNxKxK (which is a matrix NxN of matrices KxK), and for each row, I wish to multiply (dot product) all the N matrices (KxK) in the row. So that if I do this on the full M (all of the rows) I get a vector V (Nx1) of matrices (KxK) where V[i] holds the dot product of M[i,0]xM[i,1]x...xM[i,N-1].

I implemented a solution to this problem using loops, and I can't figure out a way to do this without loops.

Implementation (with loops):

    a = np.array([[1,1,1], [1,1,1], [1,1,1]])
    mat = np.array([[a,a,a,a], [a*2,a*2,a*2,a*2], [a*3,a*3,a*3,a*3],
                    [a*4,a*4,a*4,a*4]])  # the original matrix
    N, N, k, k = mat.shape
    result = np.ones((N, k, k))  # resulting matrix
    for i in range(N):
        k = functools.reduce(np.dot, mat[i,:])
        result[i,:] = k
    print(result)
Riddle answered 13/12, 2017 at 11:47 Comment(12)
In your example, N is 4. Is there an upper bound on how large N will be in your actual application?Corbeil
Are your rows always composed of identical matrices, i.e. [a,a,a,a], [a*2,a*2,a*2,a*2], etc. ?Springtime
@WarrenWeckesser - You are right that in our example N=4, but in general NO - there is no limit.Riddle
@ThomasKühn - No, it was just for the simplicity of the example. Each matrix (3x3 in our example, kxk in general) can be any matrix.Riddle
I would suggest keeping it as is.Hekate
@Hekate Do you have any experience with einsum? I don't have much experience with it, so I don't have a way of predicting whether or not it can be used to solve this...Riddle
@daniella I have used it yes. And doesn't seem like it could be used here. The accumulation needed here doesn't translate into einsum territory.Hekate
In principle, you could do something like np.einsum('ab,bc,cd,de,ef,fg,gh,hi', *factors[:8]) but it turns out to be much slower than the reduce method.Spermogonium
@PaulPanzer What do the factors represent?Riddle
Just a sequence of KxK matricesSpermogonium
It doesn't seem to work. you realize N=4 (in our example) right? what does the 8 stand for?Riddle
It is really just an example. I hardcoded 8 because generating the specification string in the general case would have incvolved some actual work. For four factors it would be 'ab,bc,cd,de'. There are other options for specifying which may be more suitable when the number of factors is not fixed. But as I said in this particular use case einsum doesn't seem to be a good solution anyway.Spermogonium
C
2

The following uses reduce but not a loop over N:

mat = mat.swapaxes(0, 1)
result = functools.reduce(lambda a, b: np.einsum('ijk,ikl->ijl', a, b), mat[:])

The einsum notation 'jk,kl->jl' expresses matrix multiplication, and the index i indicates it should be done over each value of 1st index. The first index of mat[0] or mat[1] is actually the second index of mat (column index), so as written, multiplication takes place in each column of mat. You wanted it to be done in each row, hence the use of swapaxes.

Whether this is faster or slower than for-loop version depends on the relative size of N and k. The np.dot method is highly optimized, but if the loop over N is very long, einsum might win. Some %timeit results:

  • With N=100, k=2, for-loop version takes 7.5 ms, einsum version takes 4.31 ms.
  • With N=100, k=20, for-loop version takes 27.3 ms, einsum version takes 153 ms.

So, there is a modest gain in specific cases, with major losses in most other cases. But you did not ask for an "efficient" solution, you asked for one "without loops", so here it is ("without loops" != "faster"). As Divakar suggested in a comment, you are probably better off keeping the code as is.

Centillion answered 13/12, 2017 at 15:5 Comment(0)

© 2022 - 2024 — McMap. All rights reserved.