Оптимизированный пространственный хеш Python

Question

У меня есть система частиц, ограниченная на поверхности треугольной сетки и в основном хорошо распределенная до определенной средней плотности и организованная в виде квадрата.

это для каркаса пересчета, который я делаю.

Итак, чтобы разрешить итерации между частицами, я сделал этот контейнер пространственного хеширования, но он довольно медленный для моей задачи, где есть тысячи частиц и миллионы проверок соседства.

Есть ли способ дальнейшей оптимизации этой структуры данных?или для такой системы частиц существует более быстрая структура данных?

class SpatialHash:
    def __init__(self, cell_size=0.1):
        self.buckets = {}
        self.items = {}
        self.size = cell_size

    def get_key(self, location):
        return (
            round(location[0] / self.size),
            round(location[1] / self.size),
            round(location[2] / self.size)
        )

    # item.coord refers to a vector containing the 3D location of the item
    def insert(self, item, key=None):
        if not key:
            key = self.get_key(item.coord)
        if key in self.buckets:
            self.buckets[key].add(item)
        else:
            self.buckets[key] = {item, }
        self.items[item] = self.buckets[key]

    def remove(self, item):
        if item in self.items:
            self.items[item].remove(item)
            del self.items[item]

    def update(self, item):
        self.remove(item)
        self.insert(item)

    def test_sphere(self, coord, radius, exclude=()):
        radius_sqr = radius ** 2
        radius = radius / self.size
        location = coord / self.size
        min_x = math.floor(location[0] - radius)
        max_x = math.ceil(location[0] + radius)
        min_y = math.floor(location[1] - radius)
        max_y = math.ceil(location[1] + radius)
        min_z = math.floor(location[2] - radius)
        max_z = math.ceil(location[2] + radius)
        for x in range(min_x, max_x + 1):
            for y in range(min_y, max_y + 1):
                for z in range(min_z, max_z + 1):
                    key = (x, y, z)
                    if key in self.buckets:
                        for item in self.buckets[key]:
                            if (item.coord - coord).length_squared <= radius_sqr:
                                if item in exclude:
                                    continue
                                yield item