Python numpy loop acceleration?

Question

Consider the following code using numpy arrays, which are very slow:

# Intersection of an octree and a trajectory def intersection(octree, trajectory): # Initialize numpy arrays ox = octree.get("x") oy = octree.get("y") oz = octree.get("z") oe = octree.get("extent")/2 tx = trajectory.get("x") ty = trajectory.get("y") tz = trajectory.get("z") result = np.zeros(np.size(ox)) # Loop over elements for i in range(0, np.size(tx)): for j in range(0, np.size(ox)): if (tx[i] > ox[j]-oe[j] and tx[i] < ox[j]+oe[j] and ty[i] > oy[j]-oe[j] and ty[i] < oy[j]+oe[j] and tz[i] > oz[j]-oe[j] and tz[i] < oz[j]+oe[j]): result[j] += 1 # Finalize return result 

How to rewrite a function to speed up the calculation? ( np.size(tx) == 10000 and np.size(ox) == 100000 )

Answer 1

You allocate 10,000 lists of size 100,000. The first thing to do is stop using range for the nested j loop and use the xrange generator version instead. This will save time and space by highlighting all of these lists.

The following will be the use of vectorized operations:

 for i in xrange(0, np.size(tx)): index = (ox-oe < tx[i]) & (ox+oe > tx[i]) & (oy-oe < ty[i]) & (oy+oe > ty[i]) & (oz-oe < tz[i]) & (oz+oe > tz[i]) result[index] += 1 

Answer 2

You will probably get good results by running this code under PyPy: http://pypy.org/ (instructions for our NumPy integration at https://bitbucket.org/pypy/numpy )

Answer 3

I think this will give the same result for a double loop and be faster:

 for j in xrange(np.size(ox)): result[j] += sum( abs(tx-ox[j])<oe[j] & abs(ty-oy[j])<oe[j] & abs(tz-oz[j])<oe[j] ) 

To get this: 1) change the order of the cycles (i.e. replace them) that are valid, since nothing changes in the loops; 2) pull result[j] beyond the bounds of cycle i ; 3) convert all t>ox-oe and t<ox+oe to abs(t-ox)<oe (although this may not be very fast acceleration, it is easier to read).

Since you do not have executable code, and I did not want to create a test for this, I am not 100% sure that is correct.