快速统计网格中的粒子数量

Question

我写了一些Python代码，用来从一个宇宙模拟中计算某个量。这个过程是通过检查一个粒子是否在一个大小为8,000立方的盒子里来完成的。这个盒子从原点开始，当盒子里的所有粒子都被找到后，就会移动到下一个位置。因为我总共要计算大约200万个粒子，而整个模拟的体积是150,000立方，所以这个过程花费的时间很长。

我会把我的代码贴在下面，有没有人能给我一些改进的建议呢？

提前谢谢大家。

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

我会尽量贴一些代码，但我大致的想法是这样的：创建一个粒子类，这个类要知道它所处的盒子，这个盒子是在__init__里计算出来的。每个盒子应该有一个独特的名字，这个名字可以是左下角的坐标（或者你用来定位盒子的其他方式）。

为每个粒子获取一个新的粒子类实例，然后使用一个计数器（来自collections模块）。

粒子类大概长这样：

# 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)

总之，我想你的sigma8函数可能看起来像这样：

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将是一个元组列表，这些元组将盒子的标签和该盒子里的粒子数量对应起来。（在这里，我们把粒子当作是不可区分的。）

使用列表推导式比使用循环要快得多——我觉得原因是你基本上是更依赖底层的C语言，但我不是这方面的专家。计数器（Counter）也是（据说）经过高度优化的。

注意：这些代码都没有经过测试，所以你不应该尝试直接复制粘贴然后希望它能工作。