Optimizing numpy matrix operations (currently using a for loop)

Question

I've written some code to compute n matrices based on n items in a list, and then multiply all matrices together at the end.

The code is relatively slow and I'd like to learn more about python optimization. I've used the profiling tools and identified the slow-down in my program is this matrix multiplication loop.

I wonder if anyone has any recommendations for how I might speed this up, perhaps taking advantage of built in C-based functions in Python / NumPy?

def my_matrix(x):

    # Initialise overall matrix as an identity matrix
    # | M_11 M_12 |
    # | M_21 M_22 |
    M = np.matrix([[1, 0],[0, 1]])

    for z in z_all:
        param1 = func1(z)
        param2 = func2(x, z)
        param3 = func3(x, z)

        M_11 = param1 + param2
        M_12 = param1 - param2
        M_21 = param1 * param2
        M_22 = param1 / param2

        # Multiply matrix with overall master matrix
        M = M * np.matrix([[M_11, M_12],[M_21, M_22]])
    return M

From a little background reading, it seems that function calls are computationally expensive, and so, computing arrays for my params and then accessing the arrays could be more efficient than evaluating the function within the loop each time... e.g.

 param1s = funcs(z_all)
 param2s = funcs(x, z_all)
 etc

and then in the for loop:

for i, z in enumerate(z_all):
    param1 = params1[i]
    param2 = params2[i]
 etc.

This is quicker, but only by about 10%, as the saving from fewer function calls is offset by time taken in array access using param1 = params1[i] in the loop.

Does anyone have any recommendations please?

Are any of your params/variables constant throughout your loop? If that's the case, you could define it before the loop and thus save yourself from running the function over and over. — tooty44
– tooty44, Commented Aug 2, 2014 at 20:28

ali_m · Accepted Answer · 2014-08-02 23:12:29Z

1

You could vectorize your computation of M_11, ... M_22 by doing M_11s = params1 + params2 etc.

That way you'll only have to do the matrix multiplications in the loop:

import numpy as np

...

# compute your 'params' over vectors of z-values
param1s = func1(z_all)
param2s = func2(x, z_all)
param3s = func3(x, z_all)  # you don't seem to be using this for anything...

# compute 'M_11, ... M_22'
M_11 = param1s + param2s
M_12 = param1s - param2s
M_21 = param1s * param2s
M_22 = param1s / param2s

# we construct a (2, 2, nz) array from these
M_all = np.array([[M_11, M_12], [M_21, M_22]])

# roll the 'nz' axis to the front so that its shape is (nz, 2, 2)
M_all = np.rollaxis(M_all, -1, 0)

# initialize output with the identity
M_out = np.eye(2)

# loop over each (2, 2) subarray in 'M_all', update the output with the
# corresponding dot product
for mm in M_all:
    M_out = M_out.dot(mm)

edited Aug 2, 2014 at 23:12

answered Aug 2, 2014 at 21:00

ali_m

74.6k28 gold badges230 silver badges314 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

IanRoberts Over a year ago

Thanks, this gives a 3 times speed up! Could you elaborate why this is quicker please? Is it better to work with arrays than matrices in general, since the numpy dot product seems equivalent to matrix multiplication here.

ali_m Over a year ago

The main reason for the speed-up is that the computation of M_11, ... ,M_22 is now vectorized. Operating on whole vectors is always preferable to indexing individual elements. The main argument for using arrays here rather than matrices is that arrays can have arbitrary rank, whereas matrices have a maximum rank of 2. If I was using matrices, M_all would have to be a list of (2, 2) matrices rather than one array. There are lots of other reasons to prefer arrays over matrices, see here.

IanRoberts Over a year ago

Thanks, that's a really helpful reply. I actually want to calculate M_out for a 2D parameter space. At the moment I loop through one parameter (let's call it "time") and for every time in times array, I use the above code to compute M_out. I wonder if this could be further improved by replacing the for loop with some vector operation. I have elaborated and generalised my question as I believe too much discussion here might dilute this question and response: stackoverflow.com/questions/25100046/…

Collectives™ on Stack Overflow

Optimizing numpy matrix operations (currently using a for loop)

1 Answer 1

3 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

3 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related