smart space partitioning

diff --git a/collision.py b/collision.py
index ff67a669813157ebd6b9d9de7d71d32db3fad332..ff753113c2268a08f779034c126de1ec657070a1 100644 (file)
--- a/collision.py
+++ b/collision.py
@@ -185,3 +185,83 @@ def spacePartitioning(particle_list, width, height):  # broad phase collision de
             checks.extend(tuple(combinations(grid[i][j], 2)))
 
     return list(set(checks))
+
+
+def smarterSpacePartitioning(particle_list):  # broad phase collision detection
+    particles = particle_list.copy()
+    collisions = KDTree(particles)
+    # print(collisions.get_collision(), collisions.n1, collisions.n2)
+    # quit()
+    return list(set(collisions.get_collision()))
+
+
+def median(arr):
+    return arr[len(arr) // 2]
+
+
+class KDTree:
+    default_axis = 0
+    max_depth = 2
+
+    def __init__(self, particles, axis=default_axis, depth=0):
+        self.depth = depth
+        self.particles = particles
+        self.axis = axis
+        if self.axis == 0:
+            particles.sort(key=lambda x: x.position.y)
+        else:
+            particles.sort(key=lambda x: x.position.x)
+
+        self.median = median(particles)
+
+        self.left, self.right = [], []
+
+        if self.axis == 0:
+            for particle in particles:
+                if particle.position.x < self.median.position.x:
+                    self.left.append(particle)
+                else:
+                    self.right.append(particle)
+
+                if particle.left < self.median.position.x:
+                    self.left.append(particle)
+                if particle.right >= self.median.position.x:
+                    self.right.append(particle)
+
+        else:
+            for particle in particles:
+                if particle.position.y < self.median.position.y:
+                    self.left.append(particle)
+                else:
+                    self.right.append(particle)
+
+                if particle.top < self.median.position.y:
+                    self.left.append(particle)
+                if particle.bottom >= self.median.position.y:
+                    self.right.append(particle)
+
+        self.n1, self.n2 = None, None
+        self.build()
+
+    def getAxis(self):
+        return not self.axis
+
+    def getParticlesLeft(self):
+        return self.left
+
+    def getParticlesRight(self):
+        return self.right
+
+    def build(self):
+        if self.depth == self.max_depth:
+            return
+        if len(self.particles) > 1:
+            self.n1 = KDTree(self.getParticlesLeft(), axis=self.getAxis(), depth=self.depth + 1)
+            self.n2 = KDTree(self.getParticlesRight(), axis=self.getAxis(), depth=self.depth + 1)
+
+    def get_collision(self):
+        if len(self.particles) == 1:
+            return []
+        if self.n1 is None or self.n2 is None:
+            return []
+        return list(combinations(self.particles, 2)) + self.n1.get_collision() + self.n2.get_collision()

diff --git a/display.py b/display.py
index 7195de8df156583964709342a29915eb1d1c230c..218f056e066fe2361e1597cbf642dc80a20815fe 100644 (file)
--- a/display.py
+++ b/display.py
@@ -3,7 +3,7 @@ import random
 import pygame
 
 from box import Box
-from collision import sweepAndPrune, handleParticleCollision, detectParticleCollision, spacePartitioning
+from collision import sweepAndPrune, handleParticleCollision, detectParticleCollision, spacePartitioning, smarterSpacePartitioning
 from colours import *
 from particle import Particle
 
@@ -26,8 +26,8 @@ class Display:
         self.collision_objects = {layer: [] for layer in range(self.COLLISION_LAYERS)}
         self.particles = []
 
-        rows = 10
-        cols = 10
+        rows = 3
+        cols = 3
 
         w, h = self.WIDTH / cols, self.HEIGHT / rows
 
@@ -68,7 +68,7 @@ class Display:
         for particle in self.particles:
             particle.update(delta)
 
-        for particle1, particle2 in spacePartitioning(self.particles, self.WIDTH, self.HEIGHT):
+        for particle1, particle2 in smarterSpacePartitioning(self.particles):
             if (particle1, particle2) not in self.collided_last_frame and (
                     particle2, particle1) not in self.collided_last_frame:
                 handleParticleCollision(particle1, particle2)