Simple spatial grid for particle system

Question

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 <= 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. — Toby Speight
– Toby Speight, Commented Jul 1, 2019 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. — StackOverflowToxicityVictim
– StackOverflowToxicityVictim, Commented Jul 1, 2019 at 15:42
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 — StackOverflowToxicityVictim
– StackOverflowToxicityVictim, Commented Jul 1, 2019 at 18:04
Welcome to Code Review! Please see What to do when someone answers. I have rolled back Rev 3 → 2 — Sᴀᴍ Onᴇᴌᴀ
– Sᴀᴍ Onᴇᴌᴀ ♦, Commented Jul 1, 2019 at 18:07

AlexV · Accepted Answer · 2019-07-01 19:15:36Z

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[0] / cell_size), int(particle[1] / cell_size), int(particle[2] / 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[0] / cell_size), floor(particle[1] / cell_size), floor(particle[2] / 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[0] // cell_size, particle[1] // cell_size, particle[2] // 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[0] / cell_size), int(particle[1] / cell_size), int(particle[2] / 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.

QuantumChris · Accepted Answer · 2019-07-01 14:33:38Z

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.

Stack Exchange Network

Simple spatial grid for particle system

2 Answers 2

You must log in to answer this question.

Hot Network Questions

Simple spatial grid for particle system

2 Answers 2

You must log in to answer this question.

Related

Hot Network Questions