Quadtree implementation

quadtree.py

@@ -0,0 +1,163 @@
+import pyray as RL
+from pyray import (Rectangle as Rect)
+import math
+import pdb
+import random
+from typing import Optional, Tuple, List
+from dataclasses import dataclass, field
+
+screen_width = 1280
+screen_height = 1024
+
+ball_r = 6
+ball_speed = 3.5
+num_balls = 1000
+qt_capacity = 4
+
+@dataclass
+class Ball:
+    px: float
+    py: float
+    vx: float
+    vy: float
+
+@dataclass
+class QNode:
+    aabb: Rect
+    points: List = field(default_factory=list)
+
+@dataclass
+class Quadtree:
+    node: QNode
+    subdivided = False
+    nw: Optional['Quadtree'] = None
+    ne: Optional['Quadtree'] = None
+    sw: Optional['Quadtree'] = None
+    se: Optional['Quadtree'] = None
+
+@dataclass
+class World:
+    balls = []
+    qt = {}
+    tick = 0
+    paused = False
+    mouse_clicks = []
+
+w = World()
+
+def qt_split(qt: Quadtree):
+    x, y, hw, hh = qt.node.aabb.x, qt.node.aabb.y, qt.node.aabb.width * 0.5, qt.node.aabb.height * 0.5
+    nw = Rect(x     , y     , hw, hh)
+    ne = Rect(x + hw, y     , hw, hh)
+    sw = Rect(x     , y + hh, hw, hh)
+    se = Rect(x + hw, y + hh, hw, hh)
+    qt.nw = Quadtree(QNode(nw))
+    qt.ne = Quadtree(QNode(ne))
+    qt.sw = Quadtree(QNode(sw))
+    qt.se = Quadtree(QNode(se))
+    qt.subdivided = True
+    
+def qt_insert(qt: Quadtree, p):
+    if not RL.check_collision_point_rec(p, qt.node.aabb):
+        return False
+    if qt.subdivided:
+        inserted = False
+        if not inserted: inserted = qt_insert(qt.nw, p)
+        if not inserted: inserted = qt_insert(qt.ne, p)
+        if not inserted: inserted = qt_insert(qt.sw, p)
+        if not inserted: inserted = qt_insert(qt.se, p)
+        return inserted
+    if len(qt.node.points) + 1 >= qt_capacity:
+        qt_split(qt)
+        qt.node.points.append(p)
+        inserted = False
+        for p in qt.node.points:
+            if   qt_insert(qt.nw, p): pass
+            elif qt_insert(qt.ne, p): pass
+            elif qt_insert(qt.sw, p): pass
+            elif qt_insert(qt.se, p): pass
+        qt.node.points.clear()
+        return True
+    else:
+        qt.node.points.append(p)
+        return True
+
+def construct_quadtree(points):
+    root_node = QNode(Rect(0, 0, screen_width, screen_height))
+    qt = Quadtree(root_node)
+    for p in points:
+        qt_insert(qt, p)
+    return qt
+
+def rect_values(r: Rect):
+    return r.x, r.y, r.w, r.h
+
+def init():
+    for n in range(num_balls):
+        # px = random.randrange(ball_r, screen_width  - ball_r)
+        # py = random.randrange(ball_r, screen_height - ball_r)
+        px = random.randrange(ball_r, 50)
+        py = random.randrange(ball_r, 50)
+        angle = random.uniform(0, 360)
+        vx = math.cos(angle) * ball_speed * random.uniform(1, 3)
+        vy = math.sin(angle) * ball_speed * random.uniform(1, 3)
+        w.balls.append(Ball(px, py, vx, vy))
+
+def player_input():
+    if RL.is_key_pressed(RL.KEY_SPACE):
+        w.paused = not w.paused
+    if RL.is_mouse_button_pressed(0):
+        print(RL.get_mouse_position())
+        w.mouse_clicks.append(RL.get_mouse_position())
+    
+def update():
+    # Recontruct quadtree
+    if w.paused:
+        return
+    points = []
+    for b in w.balls:
+        points.append((b.px, b.py))
+    w.qt = construct_quadtree(points)
+    
+    for ball in w.balls:
+        ball.px += ball.vx
+        ball.py += ball.vy
+        if ball.px - ball_r <= 0 or ball.px + ball_r >= screen_width:
+            # Reset position to make sure it's clamped
+            ball.px = RL.clamp(ball.px, ball_r + 0.1, screen_width - ball_r - 0.1)
+            ball.vx *= -1
+        if ball.py - ball_r <= 0 or ball.py + ball_r > screen_height:
+            # Reset position to make sure it's clamped
+            ball.py = RL.clamp(ball.py, ball_r + 0.1, screen_height - ball_r - 0.1)
+            ball.vy *= -1
+
+def draw_qt_dfs(qt: Quadtree):
+    if not qt:
+        return
+    draw_qt_dfs(qt.nw)
+    draw_qt_dfs(qt.ne)
+    draw_qt_dfs(qt.se)
+    draw_qt_dfs(qt.sw)
+    RL.draw_rectangle_lines_ex(qt.node.aabb, 0.5, RL.BLACK)
+
+def draw():
+    RL.begin_drawing()
+
+    RL.clear_background(RL.WHITE)
+    draw_qt_dfs(w.qt)
+    for ball in w.balls:
+        RL.draw_circle_lines_v((ball.px, ball.py), ball_r, RL.BLACK)
+    for mc in w.mouse_clicks:
+        RL.draw_circle_v(mc, 5, RL.RED)
+
+    RL.end_drawing()
+
+RL.init_window(screen_width, screen_height, "Quadtree");
+RL.set_target_fps(60)
+
+init()
+while not RL.window_should_close():
+    player_input()
+    update()
+    draw()
+