small-projects/physics/main.py
2025-04-08 19:58:12 -05:00

372 lines
13 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
import sys
import random
import math
from pygame.math import Vector2
# -------------------------------
# Component System and Base Classes
# -------------------------------
class Component:
def __init__(self, owner):
self.owner = owner
def update(self, dt):
pass
# RigidBody provides linear and angular motion
class RigidBody(Component):
def __init__(self, owner, mass=1.0):
super().__init__(owner)
self.mass = mass
self.inv_mass = 1.0 / mass if mass != 0 else 0
self.velocity = Vector2(0, 0)
self.force = Vector2(0, 0)
self.angle = 0.0 # in radians
self.angular_velocity = 0.0
self.torque = 0.0
self.inertia = 1.0 # moment of inertia (set later)
self.inv_inertia = 1.0
def update(self, dt):
# --- Linear integration ---
acceleration = self.force * self.inv_mass
self.velocity += acceleration * dt
self.owner.pos += self.velocity * dt
self.force = Vector2(0, 0)
# --- Angular integration ---
angular_acc = self.torque * self.inv_inertia
self.angular_velocity += angular_acc * dt
self.angle += self.angular_velocity * dt
self.torque = 0.0
def apply_force(self, force, point=None):
self.force += force
if point is not None:
# Torque = r x force (scalar in 2D: r.x * force.y - r.y * force.x)
r = point - self.owner.pos
self.torque += r.cross(force)
# -------------------------------
# Collision Components
# -------------------------------
# Base CollisionComponent (extends Component)
class CollisionComponent(Component):
def __init__(self, owner):
super().__init__(owner)
# A draw method is provided for debugging outlines.
def draw(self, surface):
pass
# CircleCollider extends CollisionComponent
class CircleCollider(CollisionComponent):
def __init__(self, owner, radius):
super().__init__(owner)
self.radius = radius
def get_position(self):
return self.owner.pos
def draw(self, surface):
pygame.draw.circle(surface, (255,255,255),
(int(self.owner.pos.x), int(self.owner.pos.y)),
self.radius, 1)
# SquareCollider extends CollisionComponent
class SquareCollider(CollisionComponent):
def __init__(self, owner, width, height):
super().__init__(owner)
self.width = width
self.height = height
def get_corners(self):
hw = self.width / 2
hh = self.height / 2
# Retrieve the current rotation (in radians) from the RigidBody if present.
rb = self.owner.get_component(RigidBody)
angle = rb.angle if rb else 0
# Define corners in local space.
corners = [Vector2(-hw, -hh), Vector2(hw, -hh),
Vector2(hw, hh), Vector2(-hw, hh)]
# Pygames rotate() expects degrees.
rotated = [corner.rotate(math.degrees(angle)) + self.owner.pos for corner in corners]
return rotated
def draw(self, surface):
corners = self.get_corners()
points = [(p.x, p.y) for p in corners]
pygame.draw.polygon(surface, (255,255,255), points, 1)
# PolygonCollider extends CollisionComponent
class PolygonCollider(CollisionComponent):
def __init__(self, owner, points):
super().__init__(owner)
# points is a list of Vector2 objects, defined in local space.
self.points = points
def get_transformed_points(self):
rb = self.owner.get_component(RigidBody)
angle = rb.angle if rb else 0
transformed = [p.rotate(math.degrees(angle)) + self.owner.pos for p in self.points]
return transformed
def draw(self, surface):
points = self.get_transformed_points()
pts = [(p.x, p.y) for p in points]
pygame.draw.polygon(surface, (255,255,255), pts, 1)
# -------------------------------
# GameObject: Entities in the Scene
# -------------------------------
class GameObject:
def __init__(self, pos):
self.pos = Vector2(pos)
self.components = []
def add_component(self, comp):
self.components.append(comp)
def get_component(self, comp_type):
for comp in self.components:
if isinstance(comp, comp_type):
return comp
return None
def update(self, dt):
for comp in self.components:
comp.update(dt)
def draw(self, surface):
# Draw any collision shapes (for visualization)
for comp in self.components:
if isinstance(comp, CollisionComponent):
comp.draw(surface)
# -------------------------------
# Collision Detection Functions
# -------------------------------
def circle_vs_circle(circleA, circleB):
posA = circleA.get_position()
posB = circleB.get_position()
rA = circleA.radius
rB = circleB.radius
diff = posB - posA
dist = diff.length()
if dist < (rA + rB):
penetration = (rA + rB) - dist
normal = diff.normalize() if dist != 0 else Vector2(1, 0)
return normal, penetration
return None, None
def circle_vs_square(circle, square):
center = circle.get_position()
# Obtain the squares rotation angle (in radians)
rb = square.owner.get_component(RigidBody)
angle = rb.angle if rb else 0
# Transform the circle center into the square's local space.
local_center = (center - square.owner.pos).rotate(-math.degrees(angle))
hw = square.width / 2
hh = square.height / 2
closest_x = max(-hw, min(local_center.x, hw))
closest_y = max(-hh, min(local_center.y, hh))
closest = Vector2(closest_x, closest_y)
diff = local_center - closest
dist = diff.length()
if dist < circle.radius:
penetration = circle.radius - dist
local_normal = diff.normalize() if dist != 0 else Vector2(1, 0)
# Rotate the normal back to world space.
normal = local_normal.rotate(math.degrees(angle))
return normal, penetration
return None, None
def square_vs_square(sq1, sq2):
corners1 = sq1.get_corners()
corners2 = sq2.get_corners()
axes = []
def get_axes(corners):
axes = []
for i in range(len(corners)):
edge = corners[(i+1) % len(corners)] - corners[i]
if edge.length() == 0:
continue
normal = Vector2(-edge.y, edge.x).normalize()
axes.append(normal)
return axes
axes.extend(get_axes(corners1))
axes.extend(get_axes(corners2))
min_overlap = float('inf')
collision_normal = None
for axis in axes:
proj1 = [p.dot(axis) for p in corners1]
proj2 = [p.dot(axis) for p in corners2]
min1, max1 = min(proj1), max(proj1)
min2, max2 = min(proj2), max(proj2)
if max1 < min2 or max2 < min1:
return None, None # Separating axis found; no collision.
overlap = min(max1, max2) - max(min1, min2)
if overlap < min_overlap:
min_overlap = overlap
collision_normal = axis
# Ensure the normal points from sq1 to sq2.
if (sq2.owner.pos - sq1.owner.pos).dot(collision_normal) < 0:
collision_normal = -collision_normal
return collision_normal, min_overlap
# Universal collision check selects routine based on collider types.
def check_collision(objA, objB):
colA = objA.get_component(CollisionComponent)
colB = objB.get_component(CollisionComponent)
if colA is None or colB is None:
return None, None
# Circle vs Circle.
if isinstance(colA, CircleCollider) and isinstance(colB, CircleCollider):
return circle_vs_circle(colA, colB)
# Circle vs Square.
if isinstance(colA, CircleCollider) and isinstance(colB, SquareCollider):
return circle_vs_square(colA, colB)
if isinstance(colA, SquareCollider) and isinstance(colB, CircleCollider):
normal, penetration = circle_vs_square(colB, colA)
if normal is not None:
return -normal, penetration
# Square vs Square.
if isinstance(colA, SquareCollider) and isinstance(colB, SquareCollider):
return square_vs_square(colA, colB)
# (PolygonCollider collisions can be implemented similarly using SAT.)
return None, None
# -------------------------------
# Collision Resolution
# -------------------------------
def resolve_collision(objA, objB, normal, penetration):
rbA = objA.get_component(RigidBody)
rbB = objB.get_component(RigidBody)
if rbA is None or rbB is None:
return
total_inv_mass = rbA.inv_mass + rbB.inv_mass
if total_inv_mass == 0:
return
# Correct positions (simple penetration resolution).
correction = normal * (penetration / total_inv_mass * 0.5)
objA.pos -= correction * rbA.inv_mass
objB.pos += correction * rbB.inv_mass
# Relative velocity along collision normal.
relative_vel = rbB.velocity - rbA.velocity
vel_along_normal = relative_vel.dot(normal)
if vel_along_normal > 0:
return # They are separating.
restitution = 0.5 # Bounce factor.
impulse_scalar = -(1 + restitution) * vel_along_normal / total_inv_mass
impulse = normal * impulse_scalar
rbA.velocity -= impulse * rbA.inv_mass
rbB.velocity += impulse * rbB.inv_mass
# --- Angular response ---
# As an approximation, we use the average of the object centers as the contact point.
contact_point = (objA.pos + objB.pos) * 0.5
rA = contact_point - objA.pos
rB = contact_point - objB.pos
# In 2D, the scalar "cross" is computed with Vector2.cross().
rbA.torque -= rA.cross(impulse)
rbB.torque += rB.cross(impulse)
# -------------------------------
# Utility: Set the Moment of Inertia Based on the Collider Type
# -------------------------------
def set_inertia(obj):
rb = obj.get_component(RigidBody)
if rb is None:
return
col = obj.get_component(CollisionComponent)
if isinstance(col, CircleCollider):
rb.inertia = 0.5 * rb.mass * (col.radius ** 2)
elif isinstance(col, SquareCollider):
rb.inertia = rb.mass * ((col.width ** 2 + col.height ** 2) / 12)
else:
rb.inertia = rb.mass # Default value.
rb.inv_inertia = 1.0 / rb.inertia if rb.inertia != 0 else 0
# -------------------------------
# Object Factories
# -------------------------------
def create_circle(pos, radius, mass):
obj = GameObject(pos)
rb = RigidBody(obj, mass)
obj.add_component(rb)
col = CircleCollider(obj, radius)
obj.add_component(col)
set_inertia(obj)
return obj
def create_square(pos, width, height, mass):
obj = GameObject(pos)
rb = RigidBody(obj, mass)
obj.add_component(rb)
col = SquareCollider(obj, width, height)
obj.add_component(col)
set_inertia(obj)
return obj
# -------------------------------
# Main Game Loop and Setup
# -------------------------------
pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Physics Engine with Rotating Collisions")
clock = pygame.time.Clock()
# Create a list of objects.
objects = []
# Create some circles.
for _ in range(3):
pos = (random.randint(100, 700), random.randint(100, 300))
circle = create_circle(pos, random.randint(15, 30), mass=1)
# Give each circle an initial random velocity.
circle.get_component(RigidBody).velocity = Vector2(random.uniform(-100, 100), random.uniform(-100, 100))
objects.append(circle)
# Create some squares ("cubes").
for _ in range(3):
pos = (random.randint(100, 700), random.randint(100, 300))
square = create_square(pos, random.randint(30, 50), random.randint(30, 50), mass=2)
square.get_component(RigidBody).velocity = Vector2(random.uniform(-100, 100), random.uniform(-100, 100))
objects.append(square)
# Gravity (pixels per second^2).
GRAVITY = Vector2(0, 200)
# Main loop.
running = True
while running:
dt = clock.tick(60) / 1000.0 # Delta time in seconds.
for event in pygame.event.get():
if event.type == pygame.QUIT:
running = False
# Apply gravity to each object.
for obj in objects:
rb = obj.get_component(RigidBody)
if rb:
rb.apply_force(GRAVITY * rb.mass)
# Update all objects.
for obj in objects:
obj.update(dt)
# Check collisions between all pairs.
for i in range(len(objects)):
for j in range(i + 1, len(objects)):
normal, penetration = check_collision(objects[i], objects[j])
if normal is not None:
resolve_collision(objects[i], objects[j], normal, penetration)
# Rendering.
screen.fill((30, 30, 30))
for obj in objects:
obj.draw(screen)
pygame.display.flip()
pygame.quit()
sys.exit()