small-projects/bounding-box-thing/main.py
2025-04-05 18:26:40 -05:00

494 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import pygame, sys, random, math
pygame.init()
screen = pygame.display.set_mode((800,600))
pygame.display.set_caption("Realistic Rotational Physics Derived from Collisions")
clock = pygame.time.Clock()
font = pygame.font.SysFont(None, 20)
# Physics constants
GRAVITY = 500.0 # pixels/s^2
RESTITUTION = 0.8 # bounciness coefficient
DAMPING = 0.99 # linear damping per frame
ANGULAR_DAMPING = 0.99 # angular damping per frame
# Friction coefficient for collisions
FRICTION_COEFF = 0.5
# Modes
MODE_EDIT = "edit"
MODE_SIMULATE = "simulate"
mode = MODE_EDIT
# Global lists and undo/redo stacks
shapes = []
undo_stack = []
redo_stack = []
# State for drawing a new shape
drawing_new_shape = False
current_drawing_points = [] # absolute coordinates
# State for dragging a shape
dragging_shape = None
drag_offset = (0,0)
drag_start_offset = None
# Collision mesh grid square size (for visual oriented boxes)
square_size = 10
# Debug flag and counters
debug_mode = False
collision_checks = 0 # number of SAT axis tests
impulse_resolutions_count = 0 # count of impulse resolutions performed
def rotate_point(x, y, angle):
"""Rotate a point (x,y) by angle (in radians)."""
cos_a = math.cos(angle)
sin_a = math.sin(angle)
return (x*cos_a - y*sin_a, x*sin_a + y*cos_a)
def compute_moment_of_inertia(local_poly, mass):
"""Approximate moment of inertia (using maximum distance from the local center)."""
cx = sum(p[0] for p in local_poly)/len(local_poly)
cy = sum(p[1] for p in local_poly)/len(local_poly)
max_dist_sq = max((p[0]-cx)**2 + (p[1]-cy)**2 for p in local_poly)
return mass * max_dist_sq / 2
def point_in_polygon(x, y, poly):
"""Raycasting algorithm to test if (x,y) is inside poly."""
num = len(poly)
j = num - 1
inside = False
for i in range(num):
xi, yi = poly[i]
xj, yj = poly[j]
if ((yi > y) != (yj > y)) and (x < (xj-xi)*(y-yi)/(yj-yi+1e-9)+xi):
inside = not inside
j = i
return inside
def greedy_mesh(grid, rows, cols, square_size):
"""Combine adjacent True cells into larger rectangles (in local coordinates)."""
used = [[False for _ in range(cols)] for _ in range(rows)]
rects = []
for row in range(rows):
for col in range(cols):
if grid[row][col] and not used[row][col]:
width = 0
while col+width < cols and grid[row][col+width] and not used[row][col+width]:
width += 1
height = 0
valid = True
while valid and row+height < rows:
for i in range(width):
if not grid[row+height][col+i] or used[row+height][col+i]:
valid = False
break
if valid:
height += 1
for r in range(row, row+height):
for c in range(col, col+width):
used[r][c] = True
rects.append(pygame.Rect(col*square_size, row*square_size, width*square_size, height*square_size))
return rects
def compute_local_collision_boxes(local_poly, square_size):
"""Compute grid cells inside local_poly and merge them using greedy meshing."""
xs = [p[0] for p in local_poly]
ys = [p[1] for p in local_poly]
max_x = int(max(xs))
max_y = int(max(ys))
cols = (max_x // square_size) + 1
rows = (max_y // square_size) + 1
grid = [[False for _ in range(cols)] for _ in range(rows)]
for row in range(rows):
for col in range(cols):
cx = col*square_size + square_size/2
cy = row*square_size + square_size/2
if point_in_polygon(cx, cy, local_poly):
grid[row][col] = True
rects = greedy_mesh(grid, rows, cols, square_size)
return rects
def sat_collision_polygon(polyA, polyB):
"""
Uses the Separating Axis Theorem (SAT) to check collision between two convex polygons.
Returns (colliding, mtv_normal, penetration_depth).
Increments the global collision_checks count for each axis test.
"""
global collision_checks
min_overlap = float('inf')
mtv_axis = None
for polygon in (polyA, polyB):
for i in range(len(polygon)):
p1 = polygon[i]
p2 = polygon[(i+1)%len(polygon)]
edge = (p2[0]-p1[0], p2[1]-p1[1])
axis = (-edge[1], edge[0])
length = math.hypot(axis[0], axis[1])
if length == 0:
continue
axis = (axis[0]/length, axis[1]/length)
collision_checks += 1
minA, maxA = float('inf'), float('-inf')
for p in polyA:
proj = p[0]*axis[0] + p[1]*axis[1]
minA = min(minA, proj)
maxA = max(maxA, proj)
minB, maxB = float('inf'), float('-inf')
for p in polyB:
proj = p[0]*axis[0] + p[1]*axis[1]
minB = min(minB, proj)
maxB = max(maxB, proj)
overlap = min(maxA, maxB) - max(minA, minB)
if overlap < 0:
return (False, None, 0)
if overlap < min_overlap:
min_overlap = overlap
mtv_axis = axis
centerA = (sum(p[0] for p in polyA)/len(polyA), sum(p[1] for p in polyA)/len(polyA))
centerB = (sum(p[0] for p in polyB)/len(polyB), sum(p[1] for p in polyB)/len(polyB))
d = (centerB[0]-centerA[0], centerB[1]-centerA[1])
if d[0]*mtv_axis[0] + d[1]*mtv_axis[1] < 0:
mtv_axis = (-mtv_axis[0], -mtv_axis[1])
return (True, mtv_axis, min_overlap)
class Shape:
def __init__(self, local_poly, offset):
"""
local_poly: list of points (x,y) in local coordinates (with minimum at (0,0)).
offset: (x,y) position in the window.
"""
self.local_poly = local_poly
self.offset = list(offset)
self.collision_rects = compute_local_collision_boxes(self.local_poly, square_size)
self.velocity = [0.0, 0.0]
self.angular_velocity = 0.0 # starts at zero; rotation will be derived from collisions
self.angle = 0.0
self.color = (random.randint(50,255), random.randint(50,255), random.randint(50,255))
area = abs(0.5 * sum(local_poly[i][0]*local_poly[(i+1)%len(local_poly)][1] - local_poly[(i+1)%len(local_poly)][0]*local_poly[i][1]
for i in range(len(local_poly))))
self.mass = area if area > 0 else 1.0
self.inertia = compute_moment_of_inertia(local_poly, self.mass)
self.local_center = (sum(p[0] for p in local_poly)/len(local_poly),
sum(p[1] for p in local_poly)/len(local_poly))
self.colliding_boxes = set()
def get_absolute_polygon(self):
"""Return the rotated and translated polygon."""
abs_poly = []
for p in self.local_poly:
rel = (p[0]-self.local_center[0], p[1]-self.local_center[1])
rot = rotate_point(rel[0], rel[1], self.angle)
abs_poly.append((rot[0] + self.offset[0] + self.local_center[0],
rot[1] + self.offset[1] + self.local_center[1]))
return abs_poly
def get_absolute_collision_polys(self):
"""Return oriented collision boxes (for visualization)."""
polys = []
for rect in self.collision_rects:
corners = [(rect.x, rect.y),
(rect.x+rect.width, rect.y),
(rect.x+rect.width, rect.y+rect.height),
(rect.x, rect.y+rect.height)]
transformed = []
for pt in corners:
rel = (pt[0]-self.local_center[0], pt[1]-self.local_center[1])
rot = rotate_point(rel[0], rel[1], self.angle)
transformed.append((rot[0] + self.offset[0] + self.local_center[0],
rot[1] + self.offset[1] + self.local_center[1]))
polys.append(transformed)
return polys
def get_bounding_rect(self):
poly = self.get_absolute_polygon()
xs = [p[0] for p in poly]
ys = [p[1] for p in poly]
return pygame.Rect(min(xs), min(ys), max(xs)-min(xs), max(ys)-min(ys))
def get_center(self):
"""Return the global center (local center translated)."""
return (self.offset[0] + self.local_center[0],
self.offset[1] + self.local_center[1])
def update(self, dt):
# Apply gravity.
self.velocity[1] += GRAVITY * dt
# Update position.
self.offset[0] += self.velocity[0] * dt
self.offset[1] += self.velocity[1] * dt
# Update rotation.
self.angle += self.angular_velocity * dt
# Bounce off window edges.
br = self.get_bounding_rect()
if br.left < 0:
self.offset[0] -= br.left
self.velocity[0] = -RESTITUTION * self.velocity[0]
if br.right > screen.get_width():
self.offset[0] -= (br.right - screen.get_width())
self.velocity[0] = -RESTITUTION * self.velocity[0]
if br.top < 0:
self.offset[1] -= br.top
self.velocity[1] = -RESTITUTION * self.velocity[1]
if br.bottom > screen.get_height():
self.offset[1] -= (br.bottom - screen.get_height())
self.velocity[1] = -RESTITUTION * self.velocity[1]
self.velocity[0] *= DAMPING
self.velocity[1] *= DAMPING
self.angular_velocity *= ANGULAR_DAMPING
def draw(self, surface):
pygame.draw.polygon(surface, self.color, self.get_absolute_polygon(), 2)
polys = self.get_absolute_collision_polys()
for idx, poly in enumerate(polys):
col = (255,0,0) if idx in self.colliding_boxes else (0,0,255)
pygame.draw.polygon(surface, col, poly, 1)
def resolve_collision_realistic(shapeA, shapeB, normal, penetration):
"""
Applies positional correction and impulsebased collision resolution (with friction)
so that rotation is derived from collision impulses.
"""
global impulse_resolutions_count
percent = 0.2
slop = 0.01
inv_massA = 1/shapeA.mass
inv_massB = 1/shapeB.mass
correction = (max(penetration - slop, 0) / (inv_massA + inv_massB)) * percent
shapeA.offset[0] -= correction * inv_massA * normal[0]
shapeA.offset[1] -= correction * inv_massA * normal[1]
shapeB.offset[0] += correction * inv_massB * normal[0]
shapeB.offset[1] += correction * inv_massB * normal[1]
centerA = shapeA.get_center()
centerB = shapeB.get_center()
contact_point = ((centerA[0]+centerB[0])/2, (centerA[1]+centerB[1])/2)
rA = (contact_point[0]-centerA[0], contact_point[1]-centerA[1])
rB = (contact_point[0]-centerB[0], contact_point[1]-centerB[1])
velA_contact = (shapeA.velocity[0] + -shapeA.angular_velocity * rA[1],
shapeA.velocity[1] + shapeA.angular_velocity * rA[0])
velB_contact = (shapeB.velocity[0] + -shapeB.angular_velocity * rB[1],
shapeB.velocity[1] + shapeB.angular_velocity * rB[0])
rv = (velB_contact[0]-velA_contact[0], velB_contact[1]-velA_contact[1])
vel_along_normal = rv[0]*normal[0] + rv[1]*normal[1]
if vel_along_normal > 0:
return
e = RESTITUTION
rA_cross_n = rA[0]*normal[1] - rA[1]*normal[0]
rB_cross_n = rB[0]*normal[1] - rB[1]*normal[0]
inv_inertiaA = 1/shapeA.inertia if shapeA.inertia != 0 else 0
inv_inertiaB = 1/shapeB.inertia if shapeB.inertia != 0 else 0
denom = inv_massA + inv_massB + (rA_cross_n**2)*inv_inertiaA + (rB_cross_n**2)*inv_inertiaB
if denom == 0:
return
j = -(1+e)*vel_along_normal/denom
impulse = (j*normal[0], j*normal[1])
shapeA.velocity[0] -= impulse[0]*inv_massA
shapeA.velocity[1] -= impulse[1]*inv_massA
shapeB.velocity[0] += impulse[0]*inv_massB
shapeB.velocity[1] += impulse[1]*inv_massB
shapeA.angular_velocity -= rA_cross_n * j * inv_inertiaA
shapeB.angular_velocity += rB_cross_n * j * inv_inertiaB
tangent = (rv[0]-vel_along_normal*normal[0],
rv[1]-vel_along_normal*normal[1])
t_length = math.hypot(tangent[0], tangent[1])
if t_length != 0:
tangent = (tangent[0]/t_length, tangent[1]/t_length)
jt = - (rv[0]*tangent[0] + rv[1]*tangent[1]) / denom
jt = max(-j*FRICTION_COEFF, min(jt, j*FRICTION_COEFF))
friction_impulse = (jt*tangent[0], jt*tangent[1])
shapeA.velocity[0] -= friction_impulse[0]*inv_massA
shapeA.velocity[1] -= friction_impulse[1]*inv_massA
shapeB.velocity[0] += friction_impulse[0]*inv_massB
shapeB.velocity[1] += friction_impulse[1]*inv_massB
shapeA.angular_velocity -= (rA[0]*tangent[1] - rA[1]*tangent[0]) * jt * inv_inertiaA
shapeB.angular_velocity += (rB[0]*tangent[1] - rB[1]*tangent[0]) * jt * inv_inertiaB
global impulse_resolutions_count
impulse_resolutions_count += 1
class Button:
def __init__(self, rect, text):
self.rect = pygame.Rect(rect)
self.text = text
def draw(self, surface):
pygame.draw.rect(surface, (100,100,100), self.rect)
txt = font.render(self.text, True, (255,255,255))
txt_rect = txt.get_rect(center=self.rect.center)
surface.blit(txt, txt_rect)
def is_clicked(self, pos):
return self.rect.collidepoint(pos)
new_shape_btn = Button((10,10,100,30), "New Shape")
play_btn = Button((120,10,100,30), "Play")
undo_btn = Button((230,10,100,30), "Undo")
redo_btn = Button((340,10,100,30), "Redo")
edit_btn = Button((450,10,100,30), "Edit")
running = True
while running:
dt = clock.tick(60)/1000.0
for event in pygame.event.get():
if event.type == pygame.QUIT:
running = False
if event.type == pygame.KEYDOWN:
if event.key == pygame.K_d:
debug_mode = not debug_mode
if mode == MODE_EDIT:
if event.type == pygame.MOUSEBUTTONDOWN:
pos = pygame.mouse.get_pos()
if new_shape_btn.is_clicked(pos):
drawing_new_shape = True
current_drawing_points = []
dragging_shape = None
elif play_btn.is_clicked(pos):
mode = MODE_SIMULATE
# Launch shapes with random linear velocities but zero angular velocity.
for shape in shapes:
shape.velocity = [random.choice([-200, -150, 150, 200]),
random.choice([-200, -150, 150, 200])]
shape.angular_velocity = 0.0
elif undo_btn.is_clicked(pos):
if undo_stack:
action = undo_stack.pop()
if action[0] == "add":
shape = action[1]
if shape in shapes:
shapes.remove(shape)
redo_stack.append(("add", shape))
elif action[0] == "move":
shape, old_offset, new_offset = action[1], action[2], action[3]
shape.offset = list(old_offset)
redo_stack.append(("move", shape, new_offset, old_offset))
elif redo_btn.is_clicked(pos):
if redo_stack:
action = redo_stack.pop()
if action[0] == "add":
shape = action[1]
shapes.append(shape)
undo_stack.append(("add", shape))
elif action[0] == "move":
shape, old_offset, new_offset = action[1], action[2], action[3]
shape.offset = list(new_offset)
undo_stack.append(("move", shape, old_offset, new_offset))
elif edit_btn.is_clicked(pos):
pass
else:
if drawing_new_shape:
current_drawing_points.append(pos)
else:
for shape in reversed(shapes):
if point_in_polygon(pos[0], pos[1], shape.get_absolute_polygon()):
dragging_shape = shape
drag_offset = (pos[0]-shape.offset[0], pos[1]-shape.offset[1])
drag_start_offset = shape.offset.copy()
break
elif event.type == pygame.MOUSEBUTTONUP:
if dragging_shape:
if drag_start_offset != dragging_shape.offset:
undo_stack.append(("move", dragging_shape, drag_start_offset, dragging_shape.offset.copy()))
redo_stack.clear()
dragging_shape = None
elif event.type == pygame.KEYDOWN:
if event.key == pygame.K_RETURN and drawing_new_shape and len(current_drawing_points) >= 3:
xs = [p[0] for p in current_drawing_points]
ys = [p[1] for p in current_drawing_points]
min_x, min_y = min(xs), min(ys)
local_poly = [(p[0]-min_x, p[1]-min_y) for p in current_drawing_points]
new_shape = Shape(local_poly, (min_x, min_y))
shapes.append(new_shape)
undo_stack.append(("add", new_shape))
redo_stack.clear()
drawing_new_shape = False
elif event.key == pygame.K_z:
if undo_stack:
action = undo_stack.pop()
if action[0] == "add":
shape = action[1]
if shape in shapes:
shapes.remove(shape)
redo_stack.append(("add", shape))
elif action[0] == "move":
shape, old_offset, new_offset = action[1], action[2], action[3]
shape.offset = list(old_offset)
redo_stack.append(("move", shape, new_offset, old_offset))
elif event.key == pygame.K_y:
if redo_stack:
action = redo_stack.pop()
if action[0] == "add":
shape = action[1]
shapes.append(shape)
undo_stack.append(("add", shape))
elif action[0] == "move":
shape, old_offset, new_offset = action[1], action[2], action[3]
shape.offset = list(new_offset)
undo_stack.append(("move", shape, old_offset, new_offset))
elif mode == MODE_SIMULATE:
if event.type == pygame.MOUSEBUTTONDOWN:
pos = pygame.mouse.get_pos()
if edit_btn.is_clicked(pos):
mode = MODE_EDIT
for shape in shapes:
shape.velocity = [0.0, 0.0]
shape.angular_velocity = 0.0
if event.type == pygame.KEYDOWN:
if event.key == pygame.K_ESCAPE:
running = False
if mode == MODE_EDIT and dragging_shape:
mpos = pygame.mouse.get_pos()
dragging_shape.offset[0] = mpos[0]-drag_offset[0]
dragging_shape.offset[1] = mpos[1]-drag_offset[1]
if mode == MODE_SIMULATE:
for shape in shapes:
shape.update(dt)
shape.colliding_boxes = set()
collision_checks = 0
impulse_resolutions_count = 0
for i in range(len(shapes)):
for j in range(i+1, len(shapes)):
shapeA = shapes[i]
shapeB = shapes[j]
polyA = shapeA.get_absolute_polygon()
polyB = shapeB.get_absolute_polygon()
colliding, normal, penetration = sat_collision_polygon(polyA, polyB)
if colliding:
resolve_collision_realistic(shapeA, shapeB, normal, penetration)
screen.fill((30,30,30))
new_shape_btn.draw(screen)
play_btn.draw(screen)
undo_btn.draw(screen)
redo_btn.draw(screen)
edit_btn.draw(screen)
if drawing_new_shape:
if len(current_drawing_points) > 1:
pygame.draw.lines(screen, (255,255,255), False, current_drawing_points, 2)
for p in current_drawing_points:
pygame.draw.circle(screen, (255,0,0), p, 4)
for shape in shapes:
shape.draw(screen)
if mode == MODE_SIMULATE:
sim_text = font.render("Simulation Mode - Click Edit to return", True, (255,255,255))
screen.blit(sim_text, (10,50))
else:
edit_text = font.render("Edit Mode - Draw/Drag shapes. Enter to finish shape.", True, (255,255,255))
screen.blit(edit_text, (10,50))
if debug_mode:
debug_lines = []
debug_lines.append(f"Mode: {mode}")
debug_lines.append(f"Shapes: {len(shapes)}")
total_boxes = sum(len(s.get_absolute_collision_polys()) for s in shapes)
debug_lines.append(f"Oriented Collision Boxes: {total_boxes}")
debug_lines.append(f"SAT Axis Tests: {collision_checks}")
debug_lines.append(f"Impulse Resolutions: {impulse_resolutions_count}")
debug_lines.append(f"FPS: {int(clock.get_fps())}")
for i, shape in enumerate(shapes):
debug_lines.append(f"Shape {i}: angle={math.degrees(shape.angle):.1f}°, ang_vel={shape.angular_velocity:.2f}")
for idx, line in enumerate(debug_lines):
txt = font.render(line, True, (255,255,0))
screen.blit(txt, (10, 90+idx*18))
pygame.display.flip()
pygame.quit()
sys.exit()