Optimizing Python loop with Numba for live load analysis

Question

I have written some code that computes flexural moments imposed by different trucks for a bridge with 300 ft length. Truck data are contained in two lists: ax_list and sp_list, which are the axle weights and axle spacings, respectively.

There is nothing much to the code, however, this needs to be repeated for millions of different truck types, and I am trying to optimize my code, which takes real long time when the actual data size set is concerned.

I tried using Numba to see if I can get any speed gains, but it did not change the execution time, whether I add Numba @jit decorators for each function or not. What am I doing wrong here? Any help would be welcome! I also included code to generate representative pseudo data for 1000 records below:

import random
from numba import jit
import numpy as np
from __future__ import division

#Generate Random Data Set

ax_list=[]
sp_list=[]

for i in xrange(1000):
    n = random.randint(3,10)
    ax = []
    sp = [0]
    for i in xrange(n):
        a = round(random.uniform(8,32),1)
        ax.append(a)
    for i in xrange(n-1):
        s = round(random.uniform(4,30), 1)
        sp.append(s)
    ax_list.append(ax)
    sp_list.append(sp)

#Input Parameters
L=300
step_size=4
cstep_size=4
moment_list=[]

@jit
#Simple moment function
def Moment(x):
    if x<L/2.0:
        return 0.5*x
    else:
        return 0.5*(L-x)

#Attempt to vectorize the Moment function, hoping for speed gains
vectMoment = np.vectorize(Moment,otypes=[np.float],cache=False)

@jit
#Truck movement function that uses the vectorized Moment function above
def SimpleSpanMoment(axles, spacings, step_size):
    travel = L + sum(spacings)
    spacings=list(spacings)
    maxmoment = 0
    axle_coords =(0-np.cumsum(spacings))
    while np.min(axle_coords) < L:
        axle_coords = axle_coords + step_size
        moment_inf = np.where((axle_coords >= 0) & (axle_coords <=L), vectMoment(axle_coords), 0)
        moment = sum(moment_inf * axles)
        if maxmoment < moment:
            maxmoment = moment
    return maxmoment

Then to run the loop for 1000 times:

%%timeit
for i in xrange(len(ax_list)):
    moment_list.append(np.around(SimpleSpanMoment(ax_list[i], sp_list[i], step_size),1))

yields:

1 loop, best of 3: 2 s per loop

Answer 1

Just a few quick points on your code:

• any future imports should be at the top of any code. Are you using this intentionally (2to3?) or this is leftover code?
• your loop to generate test data overrides i in two secondary loops - this is bad coding practice.
• you're using xrange, which is pretty much deprecated since Python 2.3. Are you really using Python 2.3 or you just have stuck with xrange?
• no entry point aka if __name__ == "__main__": to demonstrate the start of your code, you have pieces of code scattered around your listing.

• Your code as-is with the random function cannot prove Numba JIT nor caching will have any effect. My results for your functions show:

  CacheInfo(hits=0, misses=665297, maxsize=0, currsize=0)
  CacheInfo(hits=0, misses=1000, maxsize=0, currsize=0)

  Which is to be expected with a random dataset. It's recommended to put a real dataset somewhere on the 'net for people to use.

• SimpleSpanMoment uses a global variable L, you should rewrite the function to inject all variables (or create them inside the function) and not use global variables.

• cstep_size is unused.

Hope this helps you move towards a better implementation, to solve the performance issue. Good Luck!

Answer 2

First off, you should not be using Python 2 anymore, if at all possible. It will be deprecated next year. The major differences (important for this code) are that xrange is now range and print is a function.

Next, since I am currently on mobile I do not have access to numba. But since you said it does not yield any performance gain anyways, I just removed it.

With these small changes, and wrapping your code in a function, your code takes about 20 seconds (on my mobile).

The first improvement is making vectMoment actually properly vectorized. As noted in the documentation, numpy.vectorize is hardly better than writing a for loop by yourself (though it does allow you to use broadcasting on the inputs).

def moment_vect(x, L):
    return np.where(x < L/2, 0.5*x, 0.5*(L-x))

Using this drops the execution time by about a quarter.

Now for the meaty part. Whenever you write a manual while or for while using numpy, you should ask yourself if there is no better way. numpy supplies you with ways to do operations simultaneously on whole arrays. Your while loop in your main function can be realized by building a 2D array of axle_coords. This takes into account that we know where the minimum element is (since numpy.cumsum(spacings) is increasing, it is always the last element). Then we sum along one of the axes and get the maximum of the sums. This uses the numpy functions.

With this we get:

def simple_span_moment(axles, spacings, L, step_size):
    axle_coords = -np.cumsum(spacings)
    steps = np.arange(0, L - axle_coords[-1], step_size)
    axle_coords = axle_coords[:, None] + steps[None, :]
    moments = moment_vect(axle_coords, L) * np.array(axles)[:, None]
    mask = (axle_coords >= 0) & (axle_coords <= L)
    max_moment = np.where(mask,  moments, 0).sum(axis=0).max()
    return max(0, max_moment)

And then, I would put the calling code into its own function:

def graipher(axles, spacings, L, step_size):
    return [np.around(simple_span_moment(ax, sp, L, step_size), 1) for ax, sp in zip(axles, spacings)]

Together, these changes reduce the runtime from about 20 seconds down to 0.5 seconds for me.

The outermost loop can not be as easily vectorized, since your data is a ragged array. It would also probably no longer fit in memory if you have millions of entries instead of a thousand.