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
.
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$