Python: loop optimization

Question

I am quite new to Python but I've started using it for some data analysis and now I love it. Before, I used C, which I find just terrible for file I/O.

I am working on a script which computes the radial distribution function between a set of N=10000 (ten thousands) points in a 3D box with periodic boundary conditions (PBC). Basically, I have a file of 10000 lines made like this:

0.037827 0.127853 -0.481895
12.056849 -12.100216 1.607448
10.594823 1.937731 -9.527205
-5.333775 -2.345856 -9.217283
-5.772468 -10.625633 13.097802
-5.290887 12.135528 -0.143371
0.250986 7.155687 2.813220

which represents the coordinates of the N points. What I have to do is to compute the distance between every couple of points (I hence have to consider all the 49995000 combinations of 2 elements) and then do some operation on it.

Of course the most taxing part of the program is the loop over the 49995000 combinations.

My current function looks like this:

    g=[0 for i in range(Nbins)]

for i in range(Nparticles):
    for j in range(i+1,Nparticles):

        #compute distance and apply PBC
        dx=coors[i][0]-coors[j][0]
        if(dx>halfSide):
            dx-=boxSide
        elif(dx<-halfSide):
            dx+=boxSide

        dy=coors[i][1]-coors[j][1]
        if(dy>halfSide):
            dy-=boxSide
        elif(dy<-halfSide):
            dy+=boxSide

        dz=coors[i][2]-coors[j][2]
        if(dz>halfSide):
            dz-=boxSide
        elif(dz<-halfSide):
            dz+=boxSide

        d2=dx**2+dy**2+dz**2

        if(d2<(cutoff*boxSide)**2):
            g[int(sqrt(d2)/dr)]+=2

Notice: coors is a nested array, created using loadtxt() on the data file.

I basically recycled a function that I used in another program, written in C.

I am not using itertool.combinations() because I have noticed that the program runs slightly slower if I use it for some reason (one iteration is around 111 s while it runs in around 106 with this implementation).

This function takes around 106 s to run, which is terrible considering that I have to analyze some 500 configuration files.

Now, my question is: is there general a way to make this kind of loops run faster using Python? I guess that the corresponding C code would be faster, but I would like to stick to Python because it is so much easier to use.

I would like to stress that even though I'm certainly looking for a particular solution to my problem, I would like to know how to iterate more efficiently when using Python in general.

PS Please try to explain the code in your answers as much as possible (if you write any) because I'm a newbie in Python.

Update

First, I want to say that I know that, since there is a cutoff, I could write a more efficient algorithm if I divide the box in smaller boxes and only compute the distances in neighboring boxes, but for now I would only like to speed up this particular algorithm.

I also want to say that using Cython (I followed this tutorial) I managed to speed up everything a little bit, as expected (77 s, where before it took 106).

Answer 1

If memory is not an issue (and it probably isn't given that the actual amount of data won't be different from what you are doing now), you can use numpy to do the math for you and put everything into an NxN array (around 800MB at 8 bytes/float).

Given the operations your code is trying to do, I do not think you need any loops outside numpy:

g = numpy.zeros((Nbins,))
coors = numpy.array(coors)

#compute distance and apply PBC
dx = numpy.subtract.outer(coors[:, 0], coors[:, 0])
dx[dx < -halfSide] += boxSide
dx[dx > halfSide)] -= boxSide
dy = numpy.subtract.outer(coors[:, 1], coors[:, 1])
dy[dy < -halfSide] += boxSide
dy[dy > halfSide] -= boxSide
dz = numpy.subtract.outer(coors[:, 2], coors[:, 2])
dz[dz < -halfSide] += boxSide
dz[dz > halfSide] -= boxSide

d2=dx**2 + dy**2 + dz**2
# Avoid counting diagonal elements: inf would do as well as nan
numpy.fill_diagonal(d2, numpy.nan)

# This is the same length as the number of times the
# if-statement in the loops would have triggered
cond = numpy.sqrt(d2[d2 < (cutoff*boxSide)**2]) / dr
# See http://stackoverflow.com/a/24100418/2988730
np.add.at(g, cond.astype(numpy.int_), 2)

The point is to always stay away from looping in Python when dealing with large amounts of data. Python loops are inefficient. They perform lots of book-keeping operations that slow down the math code.

Libraries such as numpy and dynd provide operations that run the loops you want "under the hood", usually bypassing the bookkeeping with a C implementation. The advantage is that you can do the Python-level book-keeping once for each enormous array and then get on with processing the raw numbers instead of dealing with Python wrapper objects (numbers are full blown-objects in Python, there are no primitive types).

In this particular example, I have recast your code to use a sequence of array-level operations that do the same thing that your original code did. This requires restructuring how you think about the steps a little, but should save a lot on runtime.