Fastest Way to Mutiply $10^4$ 2x2 Matrices

Question

In a code that I work with (written in python, but also tagging as matlab because numpy is so close and I could use it if need be), we use a transfer matrix method to compute the properties of a physical system. That is, for a particle at initial position $\vec{x}_i$, we compute the final position as $$\vec{x}_f = \hat{M}(z)\cdot\vec{x}_i \approx \hat{M}(z_n)\cdots\hat{M}(z_1)\cdot\hat{M}(z_0)\cdot\vec{x}_i,$$ where for convergence reasons $n\sim 10^4$. After profiling our code, something like 90% of of the CPU's time is spent in the final matrix multiplications. I am currently using the naive implementation of this which looks like:

# The list of matrices
Ms = [M1, M2, M3, ..., Mn]

# Start with the identity matrix
result = np.identity(2)

# Multiply the matrices
for M in Ms:
    result = M @ result

My question is: is there a clever way to speed up the matrix multiplication step? Alternatively, I would also be interested in less than clever ways to shave off time with numpy voodoo.

Unfortunately, the matrices don't commute and so I can't take a logarithm, sum, and then take a matrix exponential which I assume would be faster.

---

Edit: The matrices are generated as follows:

# Calculate the constant matrices and edge matrices
Ms = get_M_const(E, B, gammas[:-1], delta_z)
rising_Ms = np.concatenate((np.array([[[1.0,], [0.0,]], [[0.0,], [1.0,]]]), get_M_edge(E[1:], gammas[1:-1], 'rising')), axis=2)
falling_Ms = get_M_edge(E, gammas[1:], 'falling')

# Interleave the arrays
c = np.empty((2,2, Ms.shape[-1]+rising_Ms.shape[-1]+falling_Ms.shape[-1],), dtype=Ms.dtype)
c[:,:,0::3] = rising_Ms
c[:,:,1::3] = Ms
c[:,:,2::3] = falling_Ms

There are technically three different types of the matrices $\hat{M}_{\text{total}} = \hat{M}_{\text{falling}}\cdot\hat{M}_{\text{const}}\cdot\hat{M}_{\text{rising}}$ that I compute using numpy functions to take advantage of vectorized routines. The variables E, B, and gammas are numpy arrays of shape (n) and delta_z is just a number. Those functions return (2,2,n) arrays which I then interleave to get the full (2,2,3n) array of matrices which get multiplied.

I guess I simplified my code too much in the first code block by listing the matrices as being in a native python list. The rest of it is how I perform matrix multiplication, however. I run the for loop over the elements of the transpose of c.

Answer 1

In general, I agree with Chris's comment that using a compiled language with the allocation of the matrices on the stack can help significantly.

Several possibilities if we are limited to Python and numpy:

• consider np.array vs np.matrix, it might happen that np.matrix is faster than np.array matrix-matrix product (it is unclear what you are using now, and how $2\times2$ size will influence the result)
• consider parallelizing computation of the final matrix as per comment from whpowell96
• maybe, you do not need to compute the overall matrix $\hat{M}(z)$. Instead of computing $(10^4-1)$ matrix-matrix products and $1$ matrix-vector products, the alternative is $10^4$ matrix-vector products that might be better if no other computations are needed.
• consider Cython and/or distributions of Python targeted at performance.

Answer 2

I wanted to follow up with this question because I'm absolutely astounded by the performance improvement I was able to achieve using a C extension to python.

I wrote a simple function in C that takes my (2,2,n) numpy array and performs the repeated matrix multiplication on it. I followed KobeGote's recommendation to hard-code the 2x2 matrix multiplication as well to avoid the overhead of for loops in numpy's version. With these changes I tested the performance with the following python scripts.

import spam
import numpy as np
import timeit
import functools as f

# Make a test array
arr = np.random.rand(2,2,10000)

# Test the speed of the new method
testfun_new = lambda: spam.numpy_test(arr)
print("Execution Time New: {:.0f} us".format(timeit.timeit(testfun_new, number=100)/100*1e6))

# Test the speed of the old method
testfun_old = lambda: f.reduce(np.dot, arr.T).T
print("Execution Time Old: {:.0f} ms".format(timeit.timeit(testfun_old, number=100)/100*1e3))

# Make sure they are the same
print()
print('Relative Difference of Elements:')
print((testfun_new() - testfun_old())/testfun_old())

spam is just the dumb name I gave my C extension for the time being. The output was:

Execution Time New: 35 us
Execution Time Old: 22 ms

Relative Difference of Elements:
[[ 8.04395924e-15 -3.77388718e-15]
 [ 7.98691127e-15 -3.94433965e-15]]

I was expecting some performance gain, but I never imagined that it would be nearly three orders of magnitude faster! I even had to double check that they were returning the same value because I thought I was just forgetting something, but they are the same to basically machine precision. With this change, the run-time of my full matrix calculation is reduced by 95% which is well below the threshold of what I care about anymore and now it's limited by other numerical operations.

I'm surprised numpy/scipy don't already have a function like this considering the performance difference. Maybe I'll suggest it to them later, or at least release my own code after it's polished up.

Edit: For anyone encountering this problem in the future, please check out the python library I have written which contains my implementation. Find it on PyPi at https://pypi.org/project/matprod/ or just install it with pip install matprod.