Python: optimising loops

Question

I wish to optimise some python code consisting of two nested loops. I am not so familar with numpy, but I understand it should enable me to improve the efficiency for such a task. Below is a test code I wrote that reflects what happens in the actual code. Currently using the numpy range and iterator is slower than the usual python one. What am I doing wrong? What is the best solution to this problem?

Thanks for your help!

import numpy
import time

# setup a problem analagous to that in the real code
npoints_per_plane = 1000
nplanes = 64
naxis = 1000
npoints3d = naxis + npoints_per_plane * nplanes
npoints = naxis + npoints_per_plane
specres = 1000

# this is where the data is being mapped to
sol = dict()
sol["ems"] = numpy.zeros(npoints3d)
sol["abs"] = numpy.zeros(npoints3d)

# this would normally be non-random input data
data = dict()
data["ems"] = numpy.zeros((npoints,specres))
data["abs"] = numpy.zeros((npoints,specres))
for ip in range(npoints):
    data["ems"][ip,:] = numpy.random.random(specres)[:]
    data["abs"][ip,:] = numpy.random.random(specres)[:]
ems_mod = numpy.random.random(1)[0]
abs_mod = numpy.random.random(1)[0]
ispec = numpy.random.randint(specres)

# this the code I want to optimize

t0 = time.time()

# usual python range and iterator
for ip in range(npoints_per_plane):
    jp = naxis + ip
    for ipl in range(nplanes):
        ip3d = jp + npoints_per_plane * ipl
        sol["ems"][ip3d] = data["ems"][jp,ispec] * ems_mod
        sol["abs"][ip3d] = data["abs"][jp,ispec] * abs_mod

t1 = time.time()

# numpy ranges and iterator
ip_vals = numpy.arange(npoints_per_plane)
ipl_vals = numpy.arange(nplanes)
for ip in numpy.nditer(ip_vals):
    jp = naxis + ip
    for ipl in numpy.nditer(ipl_vals):
        ip3d = jp + npoints_per_plane * ipl
        sol["ems"][ip3d] = data["ems"][jp,ispec] * ems_mod
        sol["abs"][ip3d] = data["abs"][jp,ispec] * abs_mod


t2 = time.time()

print "plain python: %0.3f seconds" % ( t1 - t0 )
print "numpy: %0.3f seconds" % ( t2 - t1 )

edit: put "jp = naxis + ip" in the first for loop only

additional note:

I worked out how to get numpy to quickly do the inner loop, but not the outer loop:

# numpy vectorization
for ip in xrange(npoints_per_plane):
    jp = naxis + ip
    sol["ems"][jp:jp+npoints_per_plane*nplanes:npoints_per_plane] = data["ems"][jp,ispec] * ems_mod
    sol["abs"][jp:jp+npoints_per_plane*nplanes:npoints_per_plane] = data["abs"][jp,ispec] * abs_mod

Joe's solution below shows how to do both together, thanks!

Answer 1

The best way of writing loops in numpy is not writing loops and instead using vectorized operations. For example:

c = 0
for i in range(len(a)):
    c += a[i] + b[i]

becomes

c = np.sum(a + b, axis=0)

For a and b with a shape of (100000, 100) this takes 0.344 seconds in the first variant, and 0.062 seconds in the second.

In the case presented in your question the following does what you want:

sol['ems'][naxis:] = numpy.ravel(
    numpy.repeat(
        data['ems'][naxis:,ispec,numpy.newaxis] * ems_mod,
        nplanes,
        axis=1
    ),
    order='F'
)

This could be further optimized with some tricks, but that would reduce clarity and is probably premature optimization because:

plain python: 0.064 seconds

numpy: 0.002 seconds

The solution works as follows:

Your original version contains jp = naxis + ip which merely skips the first naxis elements [naxis:] selects all but the first naxis elements. Your inner loop repeats the value of data[jp,ispec] for nplanes times and writes it to multiple locations ip3d = jp + npoints_per_plane * ipl which is equivalent to a flattened 2D array offset by naxis. Therefore a second dimension is added via numpy.newaxis to the (previously 1D) data['ems'][naxis:, ispec], the values are repeated nplanes times along this new dimension via numpy.repeat. The resulting 2D array is then flattened again via numpy.ravel (in Fortran order, i.e., with the lowest axis having the smallest stride) and written to the appropriate subarray of sol['ems']. If the target array was actually 2D, the repeat could be skipped by using automatic array broadcasting.

If you run into a situation where you cannot avoid using loops, you could use Cython (which supports efficient buffer views on numpy arrays).