Steven C. Howell
Steven C. Howell

Reputation: 18654

Matrix multiply a numpy array of matrices

I am expanding code designed to perform a function on 2 vectors so that it instead handles 2 arrays of vectors. I was using numpy.dot to take the product of two 3x3 matrices. Now I want to do this with an array of 3x3 matrices. I was not able to figure out how to do this with numpy.einsum but I think that is what I need, I'm just struggling to understand how it works.

Here is an example of what I want using a loop. Is there a way to do this without the loop?

>>> import numpy as np
>>> m = np.arange(27).reshape(3,3,3)
>>> print m
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]],

       [[18, 19, 20],
        [21, 22, 23],
        [24, 25, 26]]])
>>> m2 = np.zeros(m.shape)
>>> for i in length(m):
        m2[i] = np.dot(m[i],m[i])
>>> print m2
    array([[[   15.,    18.,    21.],
            [   42.,    54.,    66.],
            [   69.,    90.,   111.]],

           [[  366.,   396.,   426.],
            [  474.,   513.,   552.],
            [  582.,   630.,   678.]],

           [[ 1203.,  1260.,  1317.],
            [ 1392.,  1458.,  1524.],
            [ 1581.,  1656.,  1731.]]])

Upvotes: 3

Views: 1088

Answers (3)

Luca Citi
Luca Citi

Reputation: 1430

I know it's an old post but in case someone ends up here in 2023 (or later)... matmul does what you want:

>>> import numpy as np
>>> m = np.arange(27).reshape(3,3,3)
>>> print(np.matmul(m, m))
[[[  15   18   21]
  [  42   54   66]
  [  69   90  111]]

 [[ 366  396  426]
  [ 474  513  552]
  [ 582  630  678]]

 [[1203 1260 1317]
  [1392 1458 1524]
  [1581 1656 1731]]]

If the two operands passed to matmul are N-dimensional arrays (with N>2), they are treated as stacks of matrices stored in the last two dimensions. For example np.matmul(a, b) where a is 4×6×5×2 and b is 4×6×2×3 (i.e. a and b are 4×6 stacks of 5×2 and 2×3 matrices, respectively) gives a 4×6 stack of 5×3 matrices stored in a 4×6×5×3 ndarray.

Upvotes: 0

Alexander
Alexander

Reputation: 109706

You can also use Pandas. In the example below, 'p' is equivalent to your 'm' and is a 3D representation of the data. Using list comprehension, p2 calculates the dot product of each matrix. For comparison purposes, the results are then converted back to a list of numpy arrays.

import pandas as pd

%%timeit
m = np.arange(27).reshape(3,3,3)
p = pd.Panel(m)
p2 = [p[i].dot(p[i]) for i in p.items]

1000 loops, best of 3: 846 µs per loop

m2 = [p2[i].values for i in p2.items]
print(m2)

[array([[ 15,  18,  21],
       [ 42,  54,  66],
       [ 69,  90, 111]]), 
array([[366, 396, 426],
       [474, 513, 552],
       [582, 630, 678]]), 
array([[1203, 1260, 1317],
       [1392, 1458, 1524],
       [1581, 1656, 1731]])]

Numpy is much faster, however.

%%timeit
np.einsum('fij,fjk->fik', m, m)

100000 loops, best of 3: 5.01 µs per loop

Comparing it straight to np.dot:

%%timeit
[np.dot(m[i], m[i]) for i in range(len(m))]

100000 loops, best of 3: 6.78 µs per loop

Upvotes: 1

Steven C. Howell
Steven C. Howell

Reputation: 18654

I found a numpy.einsum syntax in this post Python, numpy, einsum multiply a stack of matrices which does what I want. I am not clear on why that works and want to understand how to construct the indexing strings for future uses.

>>> print np.einsum('fij,fjk->fik', V, V)
    [[[  15   18   21]
      [  42   54   66]
      [  69   90  111]]

     [[ 366  396  426]
      [ 474  513  552]
      [ 582  630  678]]

     [[1203 1260 1317]
      [1392 1458 1524]
      [1581 1656 1731]]]

Upvotes: 3

Related Questions