Quickly counting particles in grid

Question

I've written some python code to calculate a certain quantity from a cosmological simulation. It does this by checking whether a particle in contained within a box of size 8,000^3, starting at the origin and advancing the box when all particles contained within it are found. As I am counting ~2 million particles altogether, and the total size of the simulation volume is 150,000^3, this is taking a long time.

I'll post my code below, does anybody have any suggestions on how to improve it?

Thanks in advance.

from __future__ import division
import numpy as np




def check_range(pos, i, j, k):
    a = 0
    if i <= pos[2] < i+8000:
            if j <= pos[3] < j+8000:
                    if k <= pos[4] < k+8000:
                            a = 1
    return a


def sigma8(data):

    N = []

    to_do = data

    print 'Counting number of particles per cell...'

    for k in range(0,150001,8000):
            for j in range(0,150001,8000):
                    for i in range(0,150001,8000):
                            temp = []
                            n = []
                            for count in range(len(to_do)):
                                    n.append(check_range(to_do[count],i,j,k))
                                    to_do[count][1] = n[count]
                                    if to_do[count][1] == 0:
                                           temp.append(to_do[count])
                                    #Only particles that have not been found are
                                    # searched for again

                            to_do = temp
                            N.append(sum(n))
                    print 'Next row'
            print 'Next slice, %i still to find' % len(to_do)

    print 'Calculating sigma8...'

    if not sum(N) == len(data):
            return 'Error!\nN measured = {0}, total N = {1}'.format(sum(N), len(data))

    else:
            return 'sigma8 = %.4f, variance = %.4f, mean = %.4f' % (np.sqrt(sum((N-np.mean(N))**2)/len(N))/np.mean(N), np.var(N),np.mean(N))

Answer 1

I'll try to post some code, but my general idea is the following: create a Particle class that knows about the box that it lives in, which is calculated in the __init__. Each box should have a unique name, which might be the coordinate of the bottom left corner (or whatever you use to locate your boxes).

Get a new instance of the Particle class for each particle, then use a Counter (from the collections module).

Particle class looks something like:

# static consts - outside so that every instance of Particle doesn't take them along
# for the ride...
MAX_X = 150,000
X_STEP = 8000
# etc.

class Particle(object):

    def __init__(self, data):
        self.x = data[xvalue]
        self.y = data[yvalue]
        self.z = data[zvalue]
        self.compute_box_label()

    def compute_box_label(self):
        import math

        x_label = math.floor(self.x / X_STEP)
        y_label = math.floor(self.y / Y_STEP)
        z_label = math.floor(self.z / Z_STEP) 
        self.box_label = str(x_label) + '-' + str(y_label) + '-' + str(z_label)

Anyway, I imagine your sigma8 function might look like:

def sigma8(data):
    import collections as col

    particles = [Particle(x) for x in data]
    boxes = col.Counter([x.box_label for x in particles])
    counts = boxes.most_common()

    #some other stuff

counts will be a list of tuples which map a box label to the number of particles in that box. (Here we're treating particles as indistinguishable.)

Using list comprehensions is much faster than using loops---I think the reason is that you're basically relying more on the underlying C, but I'm not the person to ask. Counter is (supposedly) highly-optimized as well.

Note: None of this code has been tested, so you shouldn't try the cut-and-paste-and-hope-it-works method here.