4
$\begingroup$

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.

$\endgroup$
  • 2
    $\begingroup$ This seems like it would benefit a lot from parallelizing. Computing different chunks of the product in parallel then combining at the end $\endgroup$ – whpowell96 Feb 7 at 2:12
  • 2
    $\begingroup$ Don't use heap allocated arrays for this. Use a compiled language, like Julia, where you can stack-allocate statically-defined arrays and you'll get pretty massive performance benefits for small matrix operations. $\endgroup$ – Chris Rackauckas Feb 7 at 2:47
  • $\begingroup$ A parallel reducer in a mapreduce style operation will get the job done. $\endgroup$ – Richard Feb 7 at 3:14
  • $\begingroup$ Additionally, one benefit of a parallelized approach is that you can do some regularization, e.g., column scaling in parallel as well and rescale everything after you form the final product. This will help with conditioning of the final matrix as $10^4$ matrix products seems like it could have disastrous conditioning issues if you aren't careful. $\endgroup$ – whpowell96 Feb 7 at 3:59
  • $\begingroup$ How do you generate the list Ms? Probably it would make sense to merge list generation and multiplication: instead of returning a list, multiply matrices one by one as soon as you generate them. (This works if you generate them in a simple order, for instance first to last or last to first). $\endgroup$ – Federico Poloni Feb 7 at 10:02
7
$\begingroup$

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.
| cite | improve this answer | |
$\endgroup$
  • 2
    $\begingroup$ I would just add that I wrote a small Fortran program to test where the 2x2 multplication was written explicitly (no call to a subroutine), and the performance is many orders of magnitude faster even without an optimization flags. It's so fast that I wouldn't even bother with any other approach. $\endgroup$ – KobeGote Feb 7 at 6:00
  • 1
    $\begingroup$ That's a really good point. For the 2x2 matrices numpy is probably using for loops that only add overhead. $\endgroup$ – ElectronsAndStuff Feb 7 at 15:16
  • $\begingroup$ At this point, I'm thinking of writing a python extension in c/c++ to do the repeated matrix multiplication when you pass it a (2,2,n) numpy array. I want to avoid parallelizing this level, since I want to repeat our operation several times and can trivially make that part parallel instead. Unfortunately, I do need the full matrix $\hat{M}$ so I can't try to get away with vector products. $\endgroup$ – ElectronsAndStuff Feb 7 at 19:44
  • $\begingroup$ The use of precompiler lib like numba could also be considered: numba.pydata.org the thing is they require minimal code change (keeping everything in python) and support most useful numpy functions: numba.pydata.org/numba-doc/dev/reference/numpysupported.html $\endgroup$ – G.Clavier Feb 19 at 22:15
3
$\begingroup$

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.

| cite | improve this answer | |
$\endgroup$
  • $\begingroup$ I'd mention that using the StaticArrays.jl will automatically do this transform to non-BLAS hardcoded matmuls, and you get about the same performance benefit. This is a case where languages with a strict type system and an optimizing compiler will shine because you can't rely on computational overhead to be high enough for C-call-style vectorization to be fast. $\endgroup$ – Chris Rackauckas Feb 8 at 21:28

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.