Files
rigid-body/rigidbody.py
2026-03-13 14:20:14 -04:00

113 lines
4.8 KiB
Python

from dataclasses import dataclass
from itertools import combinations
import pygame as pg
from tools import debug
from collider.system import intersect
from collider.types import BaseCollider, ColliderContact
from transform import Transform
@dataclass
class RigidBody:
transform: Transform
collider: BaseCollider
velocity: pg.Vector2
angular_velocity: float = 0
mass: float = 1.0
restitution: float = 0.5
coef_friction: float = 0.6
def apply_impulse(self, j: pg.Vector2, point: pg.Vector2) -> None:
moment_arm = point - self.transform.global_position
self.velocity += j * self.inv_mass
self.angular_velocity += moment_arm.cross(j) * self.inv_moment_of_inertia
@property
def inv_mass(self) -> float:
return 0.0 if self.mass == 0.0 else 1/self.mass
@property
def moment_of_inertia(self) -> float:
return 0.0 if self.mass == 0.0 else self.collider.moment_of_inertia(self.mass)
@property
def inv_moment_of_inertia(self) -> float:
return 0.0 if self.mass == 0.0 else 1/self.moment_of_inertia
class PhysicsSystem:
def __init__(self, gravity: int = 250):
self.bodies: list[RigidBody] = []
self.gravity = pg.Vector2(0,gravity)
def add_body(self, body: RigidBody) -> None:
self.bodies.append(body)
def update(self, dtr: float):
substeps = 2
for _ in range(substeps):
dt = dtr / substeps
g = self.gravity * dt
for body in self.bodies:
if body.mass != 0.0:
body.velocity += g
body.transform.position += dt * body.velocity
body.transform.rotation += dt * body.angular_velocity
debug.draw_collider(body.collider, body.transform)
for a, b in combinations(self.bodies, 2):
if collision := intersect(a.collider, b.collider, a.transform, b.transform):
self.resolve_collision(a, b, collision, dt)
def resolve_collision(self, a: RigidBody, b: RigidBody, collision: ColliderContact, dt: float) -> None: #TODO: refactor this and clean up
if a.mass == 0.0 and b.mass == 0.0: return
SLACK=0.3
SLOP = 0.001
BAUMGARTE_CONSTANT = 0.3
correction = collision.penetration * SLACK * collision.normal
baumgarte_bias = BAUMGARTE_CONSTANT * (collision.penetration - SLOP / dt)
if a.mass != 0.0:
a.transform.position += correction
if b.mass != 0.0:
b.transform.position -= correction
debug.draw_contact(collision)
restitution = a.restitution * b.restitution
friction = a.coef_friction * b.coef_friction
tangent = collision.normal.rotate(90)
a_w = a.angular_velocity
b_w = b.angular_velocity
v_rel_linear = pg.Vector2(a.velocity - b.velocity)
for point in collision.points:
r_a = point - a.transform.position
r_b = point - b.transform.position
w_cross_r_a = pg.Vector2(-r_a.y * a_w, r_a.x * a_w) #omega is not encoded as a vector, so we pretend it is
w_cross_r_b = pg.Vector2(-r_b.y * b_w, r_b.x * b_w)
normal_mass = a.inv_mass + b.inv_mass + (a.inv_moment_of_inertia * r_a.cross(collision.normal)**2) +(b.inv_moment_of_inertia * r_b.cross(collision.normal) ** 2)
v_rel = v_rel_linear + w_cross_r_a - w_cross_r_b
v_rel_tangent = v_rel.dot(tangent)
if v_rel_linear.dot(collision.normal) < 0.0:
collision_impulse = -(1+restitution)*(v_rel.dot(collision.normal))\
/ (a.inv_mass + b.inv_mass + (r_a.cross(collision.normal)**2 * a.inv_moment_of_inertia)\
+ (r_b.cross(collision.normal)**2 * b.inv_moment_of_inertia)) / len(collision.points)
a.apply_impulse(collision_impulse * collision.normal, point)
b.apply_impulse(-1 * collision_impulse * collision.normal, point)
friction_impulse = (-v_rel_tangent / (a.inv_mass + b.inv_mass +\
(r_a.cross(tangent))**2 * a.inv_moment_of_inertia +\
(r_b.cross(tangent)**2 * b.inv_moment_of_inertia))) / len(collision.points)
friction_impulse = pg.math.clamp(friction_impulse, -abs(collision_impulse)*friction, abs(collision_impulse)*friction)
a.apply_impulse( friction_impulse * tangent, point)
b.apply_impulse( -1 * friction_impulse * tangent, point)
baumgarte_impulse = -normal_mass * (v_rel_linear.length() + baumgarte_bias)
baumgarte_impulse = max(baumgarte_impulse, 0)
a.apply_impulse(baumgarte_impulse * collision.normal, point)
b.apply_impulse(-1 * baumgarte_impulse * collision.normal, point)