236 lines
7.4 KiB
Python
236 lines
7.4 KiB
Python
import pyray as RL
|
|
from pyray import (Rectangle as Rect, Vector2 as Vec2)
|
|
import math
|
|
import pdb
|
|
import random
|
|
from typing import Optional, Tuple, List
|
|
from dataclasses import dataclass, field
|
|
|
|
screen_width = 1200
|
|
screen_height = 960
|
|
|
|
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[Tuple[float, float]] = field(default_factory=list)
|
|
|
|
@dataclass
|
|
class Quadtree:
|
|
node: QNode
|
|
subdivided = False
|
|
direction: str = ''
|
|
nw: Optional['Quadtree'] = None
|
|
ne: Optional['Quadtree'] = None
|
|
sw: Optional['Quadtree'] = None
|
|
se: Optional['Quadtree'] = None
|
|
parent: Optional['Quadtree'] = None
|
|
|
|
@dataclass
|
|
class World:
|
|
balls = []
|
|
qt = {}
|
|
tick = 0
|
|
paused = False
|
|
mouse_clicks = []
|
|
nearest_points = []
|
|
visited_quadrants = []
|
|
nearest_pairs = []
|
|
current_visited = 0
|
|
|
|
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), parent=qt, direction='NW')
|
|
qt.ne = Quadtree(QNode(ne), parent=qt, direction='NE')
|
|
qt.sw = Quadtree(QNode(sw), parent=qt, direction='SW')
|
|
qt.se = Quadtree(QNode(se), parent=qt, direction='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 qt_find_nearest_point(qt: Quadtree, point) -> Tuple[float, float]:
|
|
containing_qt = qt
|
|
while containing_qt.subdivided:
|
|
if RL.check_collision_point_rec(point, containing_qt.nw.node.aabb):
|
|
containing_qt = containing_qt.nw
|
|
elif RL.check_collision_point_rec(point, containing_qt.ne.node.aabb):
|
|
containing_qt = containing_qt.ne
|
|
elif RL.check_collision_point_rec(point, containing_qt.sw.node.aabb):
|
|
containing_qt = containing_qt.sw
|
|
elif RL.check_collision_point_rec(point, containing_qt.se.node.aabb):
|
|
containing_qt = containing_qt.se
|
|
|
|
def search_for_nearest(qt: Quadtree, direction = ''):
|
|
nonlocal closest_point, closest_dist
|
|
contains_point = RL.check_collision_point_rec(point, qt.node.aabb)
|
|
if not contains_point:
|
|
px, py = point.x, point.y
|
|
dx, dy = 0,0
|
|
if px < qt.node.aabb.x:
|
|
dx = qt.node.aabb.x - px
|
|
elif px > qt.node.aabb.x + qt.node.aabb.width:
|
|
dx = px - (qt.node.aabb.x + qt.node.aabb.width)
|
|
if py < qt.node.aabb.y:
|
|
dy = qt.node.aabb.y - py
|
|
elif py > qt.node.aabb.y + qt.node.aabb.height:
|
|
dy = py - (qt.node.aabb.y + qt.node.aabb.height)
|
|
dist = RL.vector2_length(Vec2(dx, dy))
|
|
if dist >= closest_dist:
|
|
return
|
|
if qt.subdivided:
|
|
if direction != 'NW': search_for_nearest(qt.nw)
|
|
if direction != 'NE': search_for_nearest(qt.ne)
|
|
if direction != 'SW': search_for_nearest(qt.sw)
|
|
if direction != 'SE': search_for_nearest(qt.se)
|
|
|
|
w.visited_quadrants.append(qt)
|
|
for p in qt.node.points:
|
|
d = RL.vector_2distance(point, Vec2(p[0], p[1]))
|
|
if d < closest_dist:
|
|
closest_point = p
|
|
closest_dist = d
|
|
closest_point = None
|
|
closest_dist = float('inf')
|
|
previous_direction = ''
|
|
while containing_qt is not None:
|
|
search_for_nearest(containing_qt, previous_direction)
|
|
previous_direction = containing_qt.direction
|
|
containing_qt = containing_qt.parent
|
|
|
|
return closest_point
|
|
|
|
|
|
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, 50)
|
|
py = random.randrange(ball_r, 50)
|
|
# px = random.randrange(ball_r, screen_width - ball_r)
|
|
# py = random.randrange(ball_r, screen_height - ball_r)
|
|
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):
|
|
mouse_pos = RL.get_mouse_position()
|
|
nearest = qt_find_nearest_point(w.qt, mouse_pos)
|
|
w.nearest_pairs.append((mouse_pos, nearest))
|
|
w.paused = True
|
|
if RL.is_key_pressed(RL.KEY_ENTER):
|
|
w.current_visited += 1
|
|
|
|
def update():
|
|
# Recontruct quadtree
|
|
if w.paused:
|
|
return
|
|
|
|
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
|
|
|
|
points = []
|
|
for b in w.balls:
|
|
points.append((b.px, b.py))
|
|
w.qt = construct_quadtree(points)
|
|
|
|
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 i in range(w.current_visited):
|
|
RL.draw_rectangle_rec(w.visited_quadrants[i].node.aabb, RL.LIGHTGRAY)
|
|
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)
|
|
for np in w.nearest_points:
|
|
RL.draw_circle_lines_v(Vec2(np[0], np[1]), ball_r, RL.GREEN)
|
|
for mc,(px,py) in w.nearest_pairs:
|
|
pos = Vec2(px, py)
|
|
RL.draw_circle_v(mc, 5, RL.RED)
|
|
RL.draw_circle_lines_v(pos, ball_r, RL.GREEN)
|
|
RL.draw_line_v(mc, pos, RL.BLUE)
|
|
|
|
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()
|
|
|