Simple spatial grid for particle system I am going to simulate a particle system, the particles are infinitesimal points which apply forces on the neighbors and I need a fast way of making proximity checks, they are going to have a maximum check radius but not a minimum and will be almost evenly spaced.

The simulation will be very similar to this: https://youtu.be/SFf3pcE08NM

I thought a spatial grid would be the easiest approach for it and I need it to be as cost efficient as possible to matter how ugly it gets.

Is there any optimization I can make on this code?

Is there a way to compute the optimal cell size from the average distance between particles?

_empty_set = set()


class SpatialHash:
    def __init__(self, cell_size=0.1):
        self.cells = {}
        self.bucket_map = {}
        self.cell_size = cell_size

    def key(self, co):
        return int(co[0] / self.cell_size), int(co[1] / self.cell_size), int(co[2] / self.cell_size)

    def add_item(self, item):
        co = item.co
        k = int(co[0] / self.cell_size), int(co[1] / self.cell_size), int(co[2] / self.cell_size)
        if k in self.cells:
            c = self.cells[k]
        else:
            c = set()
            self.cell_size[k] = c
        c.add(item)
        self.bucket_map[item] = c

    def remove_item(self, item):
        self.bucket_map[item].remove(item)

    def update_item(self, item):
        self.bucket_map[item].remove(item)
        self.add_item(item)

    def check_sphere(self, co, radius, exclude=()):
        r_sqr = radius * radius
        for x in range(int((co[0] - radius) / self.cell_size),
                       int((co[0] + radius) / self.cell_size) + 1):
            for y in range(int((co[1] - radius) / self.cell_size),
                           int((co[1] + radius) / self.cell_size) + 1):
                for z in range(int((co[2] - radius) / self.cell_size),
                               int((co[2] + radius) / self.cell_size) + 1):
                    for item in self.cells.get((x, y, z), _empty_set):
                        if item not in exclude and (item.co - co).length_squared &lt;= r_sqr:
                            yield item
You could rewrite your key() method and use it to much greater effect, reducing repeated code.

def key(self, co, radius=0):
    return tuple((c + radius) // self.cell_size for c in co)


You can then call key() in both add_item() and check_sphere(). In add_item() it will replace k, while in check_sphere() you can use it to define your ranges.

After that I would look for a way to flatten those nested for loops which will likely be an algorithm change. Hopefully someone else has some ideas there.