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 / self.cell_size), int(co / self.cell_size), int(co / self.cell_size)

co = item.co
k = int(co / self.cell_size), int(co / self.cell_size), int(co / self.cell_size)
if k in self.cells:
c = self.cells[k]
else:
c = set()
self.cell_size[k] = c
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)

for x in range(int((co - radius) / self.cell_size),
int((co + radius) / self.cell_size) + 1):
for y in range(int((co - radius) / self.cell_size),
int((co + radius) / self.cell_size) + 1):
for z in range(int((co - radius) / self.cell_size),
int((co + 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 <= r_sqr:
yield item
• Did you consider using an octree for this? If so, perhaps explaining why you rejected that may help inform the reviews. Jul 1 '19 at 15:35
• Well,The particles will be massive but well distributed, If I do it right I can fit snugly an average amount of particles on the grid, with an octree all nodes would have about athe same depth so cutting the overhead of travessing it would be equivalent to a spatial hash. Jul 1 '19 at 15:42
• Do you have a function to generate test data? Jul 1 '19 at 17:58
• No, but I'm testing some "massive" simulations on blender with it right now and it seems to be behaving properly, edited it now, because I found a bug but its extremely slow yet i.postimg.cc/1tPxSh0d/image.png Jul 1 '19 at 18:04
• Welcome to Code Review! Please see What to do when someone answers. I have rolled back Rev 3 → 2 Jul 1 '19 at 18:07

HoboProber's answer is good with regards to clean code, but more than likely slower than what you already have.

At the moment we basically have to following variants of the key function:

from math import floor

def key_op(particle, cell_size):
return int(particle / cell_size), int(particle / cell_size), int(particle / cell_size)

def key_hobo(particle, cell_size):
return tuple(coordinate // cell_size for coordinate in particle)

def key_op_edit(particle, cell_size):
"""This was taken from your edit that was rolled back due to answer invalidation"""
return floor(particle / cell_size), floor(particle / cell_size), floor(particle / cell_size)

HoboProber already sneakily introduced floor division (think \\) to you, so the explicit version of this is also up for discussion:

def key_floor_div(particle, cell_size):
return particle // cell_size, particle // cell_size, particle // cell_size

The timings for them are as follows:

key_op          1.28 µs ± 16.2 ns (mean ± std. dev. of 7 runs, 1000000 loops each)
key_op_edit     1.62 µs ± 13.8 ns (mean ± std. dev. of 7 runs, 1000000 loops each)
key_hobo        2.09 µs ± 34.3 ns (mean ± std. dev. of 7 runs, 100000 loops each)
key_floor_div   1.11 µs ± 14.3 ns (mean ± std. dev. of 7 runs, 1000000 loops each)

The timing was done in an IPython environment running %timeit func(example, cell_size) with example = (23849.234, 1399283.8923, 2137842.24357) and cell_size = 10.

Based on these results there seems to be a narrow win for the floor div version over the original implementation.

But can we do better? Enter numba, a just-in-time compiler for Python code. At this early testing stage basically all you have to do is from numba import jit and then

@jit(nopython=True)
def key_op_numba(particle, cell_size):
return int(particle / cell_size), int(particle / cell_size), int(particle / cell_size)

Act accordingly for the other versions. Now lets look at the timings:

key_op          623 ns ± 8.42 ns (mean ± std. dev. of 7 runs, 1000000 loops each)
key_op_edit     618 ns ± 10.5 ns (mean ± std. dev. of 7 runs, 1000000 loops each)
key_hobo*       N/A
key_floor_div   628 ns ± 7.06 ns (mean ± std. dev. of 7 runs, 1000000 loops each)

As you can see numba can easily double the performance of this functions, and that is with the default settings. You might be able to squeeze out a little bit more of you are really going for it.

Bug: There is likely a bug in this piece of code:

if k in self.cells:
c = self.cells[k]
else:
c = set()
self.cell_size[k] = c
#            ^-- should likely be self.cells here

Note: a defaultdict can help to simplify this part of the code, though I cannot really say something on how its performance compares to that of a normal dict and your member-check/if-construct.

*numba does not seem to like to instantiate a tuple from a generator expression, transforming it into a list comprehension yields 1.25 µs ± 25.1 ns. You then have to convert it to a tuple to make it hashable, e.g. tuple(key_hobo_numba(example, cell_size), which leaves us with the final timing of 1.45 µs ± 14.8 ns.

You could rewrite your key() method and use it to much greater effect, reducing repeated code.