Matrix multiplication in Python

Linear algebra is absolutely everywhere, and many algorithms and data structures use this basic mathematical operation. There are several ways we can multiply arrays of data together in Python - some slow, and some extremely efficient/fast.

First and foremost, we have NumPy's dot method, which multiplies two arrays together (I, however, tend to write something along the lines of from numpy import dot as mult, so the fact that it's a matrix multiplication is more obvious). EDIT: It should also be noted that in Python 3.5+, the __matmul__ method defines the @ infix operator, so instead of np.dot(arr1, arr2), we may write arr1 @ arr2 (c.f. PEP 465).

We may also use NumPy's einsum method, which was recently brought to my attention by a friend. For example, given two matrices:
>>> X = np.random.randn(4, 4)
>>> Y = np.random.randn(4, 4)
>>> X
array([[ 1.65507608, -0.39287161, -1.05303575, -0.61274276],
       [-0.88893625, -0.06474647,  0.78705911,  2.00852693],
       [ 0.11752255, -0.91369439, -1.33496289,  0.4302434 ],
       [-0.95187053, -0.53812809,  1.41946549, -0.7446111 ]])
>>> Y
array([[ 0.01541652, -0.60860537,  0.1732339 , -1.19839415],
       [ 0.04467971,  0.32428765,  0.22817222,  0.07159069],
       [ 0.48042088,  0.28170302, -1.09161026,  1.53414709],
       [-0.79462157,  0.81292274,  1.54425452, -0.33880215]])
we may multiply these using np.einsum as follows:
>>> np.einsum('ij,jk->ik', X, Y)
array([[-0.01103961, -1.92944747,  0.40034676, -3.41947261],
       [-1.23449636,  2.37450903,  2.07374774,  1.58763194],
       [-1.02223654, -0.3941331 ,  1.93354371, -2.40004717],
       [ 1.23490695,  0.19936164, -2.98705427,  3.5321358 ]])
It's pretty darn fast. For example, using the following little script, I arrive at the following graph of array scales vs. time using Numba's jit decorator, dot and einsum.
#! /usr/bin/env python3.5
# -*- coding: utf-8 -*-


""" Compare graphically the amount of time it takes to multiply two matrices 
together with different methods.

Note: on my machine, runs with Anaconda Python 3.5 (because of Numba)
"""


import numpy as np
import matplotlib.pyplot as plt
from time import time
from types import BuiltinFunctionType, FunctionType
from typing import List, Union
from functools import partial
from numba import jit


CEIL = 5

def timeFunc(method: Union[BuiltinFunctionType, FunctionType], \
        X: np.ndarray, Y: np.ndarray) -> float:
    """ time a method """
    _start = time()
    method(X, Y)
    _end   = time()
    return _end - _start

@jit
def jittified(X: np.ndarray, Y: np.ndarray) -> np.ndarray:
    """ pure Python implementation of matrix multiplication (mad slow) that's
    been `jittified` with numba
    """
    #assert X.shape[1] == Y.shape[0]

    _X, _Y, _Z = X.shape
    mult = np.empty((X.shape[0], Y.shape[1]))

    for i in range(_X):
        for j in range(_Y):
            s = 0.0
            for k in range(_Z):
                s += X[i][k] * Y[k][j]
            mult[i][j] = s

    return mult

def graph(data: List[List[float]], methods: List[str]) -> None:
    """ graph data """
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    
    for i in range(len(data[0])):
        ax.plot(data[i], marker='.', label=methods[i])

    ax.set_yscale('log')
    ax.set_xlim(-0.1, CEIL - 1.9)

    plt.xticks(range(0, CEIL - 1))
    plt.title('Comparison of different methods', size=12)
    plt.xlabel(r'Length ($\log_{10}(L)$)')
    plt.ylabel(r'Time ($\log_{10}(s)$)')
    plt.legend()
    plt.show()

def main(INIT: int =10) -> None:
    methods = [partial(np.einsum, 'ij, jk->ik'), np.dot, jittified]
    names = []
    for method in methods:
        try:
            names.append(method.__name__)
        except AttributeError:
            # due to functools.partial
            names.append(method.func.__name__)
    times = [[] for _ in range(CEIL)] # watch immutability

    for scale in range(CEIL):
        print("** Manufacturing arrays of dimensions %dx%d" % (INIT ** scale, \
                                                               INIT ** scale))
        X = np.random.randn(INIT ** scale, INIT ** scale)
        Y = np.random.randn(INIT ** scale, INIT ** scale)
        print("** Finished")

        for method in methods:
            res = timeFunc(method, X, Y)
            times[scale].append(res)

    graph(times, names)

if __name__ == '__main__': main()
 

Comparing different methods' times for matrix multiplication in Python.

By the way, this is a great resource for learning a little about np.einsum. (If you'd like a humbling experience, check out the source code on GitHub as well.) The end result is, of course, that this method isn't great for working with large data. On the other hand, jit does a remarkable job!