initial commit

main.py

@@ -0,0 +1,178 @@
+import pygame as pg
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from itertools import combinations
+from collections import namedtuple
+from random import random
+
+GLOBAL_GRAVITY=250
+DIMENSIONS = namedtuple("Dimensions", ["WIDTH", "HEIGHT"])(500,500)
+
+@dataclass
+class ColliderContact:
+    __slots__ = ["point", "normal", "penetration"]
+    point: pg.Vector2
+    normal: pg.Vector2
+    penetration: float
+
+@dataclass
+class RigidBody:
+    position: pg.Vector2
+    collider: "BaseCollider"
+    velocity: pg.Vector2
+    mass: float = 1.0
+    restitution: float = 0.5
+
+    @property
+    def inv_mass(self):
+        return 0.0 if self.mass == 0 else 1.0 / self.mass
+
+class BaseCollider(ABC):
+     
+    def collide(self, other: "BaseCollider", this_position: pg.Vector2, other_position: pg.Vector2) -> ColliderContact | None:
+        if isinstance(other, CircleCollider):
+            return self.collide_circle(other, this_position, other_position)
+        
+
+    @abstractmethod
+    def collide_circle(self, other: "CircleCollider", this_position: pg.Vector2, other_position: pg.Vector2) -> ColliderContact | None:
+        pass
+
+class CircleCollider(BaseCollider):
+
+    def __init__(self, radius: float):
+        self.radius = radius
+
+    def collide_circle(self, other: "CircleCollider", this_position: pg.Vector2, other_position: pg.Vector2) -> ColliderContact | None:
+        delta = this_position - other_position
+        dist = delta.length()
+        radii = self.radius + other.radius
+        if dist >= radii or dist == 0:
+            return None
+        normal = delta.normalize()
+        return ColliderContact(
+            point=other_position + normal * other.radius,
+            normal=normal,
+            penetration=radii - dist,
+        )  
+
+class PhysicsSystem:
+    
+    def __init__(self):
+        self.bodies: list[RigidBody] = []
+        self.gravity = pg.Vector2(0,GLOBAL_GRAVITY)
+    
+    def add_body(self, body: RigidBody) -> None:
+        self.bodies.append(body)
+
+    def update(self, dt: float):
+        g = self.gravity * dt
+        for body in self.bodies:
+            body.velocity += g
+            body.position += dt * body.velocity
+            self.resolve_bounds(body)
+
+        for a, b in combinations(self.bodies, 2):
+            if collision := a.collider.collide(b.collider, a.position, b.position):
+                self.resolve_collision(a, b, collision)
+        
+    def resolve_collision(self, a: RigidBody, b: RigidBody, collision: ColliderContact) -> None:
+        SLACK=0.4
+        correction = collision.penetration / (a.inv_mass + b.inv_mass) * SLACK * collision.normal
+        a.position += correction
+        b.position -= correction
+
+        v_rel = a.velocity - b.velocity
+        restitution = a.restitution * b.restitution
+        impulse = (-(1 + restitution) * (collision.normal.dot(v_rel))) / (a.inv_mass + b.inv_mass)
+        a.velocity += collision.normal*impulse*a.inv_mass
+        b.velocity -= collision.normal*impulse*a.inv_mass
+
+    def resolve_bounds(self, body: RigidBody) -> None:
+        r = body.collider.radius
+
+        if body.position.x - r < 0:
+            body.position.x = r
+            body.velocity.x = abs(body.velocity.x) * body.restitution
+        elif body.position.x + r > DIMENSIONS.WIDTH:
+            body.position.x = DIMENSIONS.WIDTH - r
+            body.velocity.x = -abs(body.velocity.x) * body.restitution
+
+        if body.position.y - r < 0:
+            body.position.y = r
+            body.velocity.y = abs(body.velocity.y) * body.restitution
+        elif body.position.y + r > DIMENSIONS.HEIGHT:
+            body.position.y = DIMENSIONS.HEIGHT - r
+            body.velocity.y = -abs(body.velocity.y) * body.restitution
+
+class Circle:
+
+    def __init__(self, radius: float, body: RigidBody, sprite: pg.Surface):
+        self.radius = radius
+        self.body = body
+        self.sprite = sprite
+
+    def draw(self, screen: pg.Surface):
+        screen.blit(
+            self.sprite,
+            (
+                self.body.position.x - self.radius,
+                self.body.position.y - self.radius
+            )
+        )
+
+class World:
+
+    def __init__(self, screen: pg.Surface):
+        self.physics = PhysicsSystem()
+        self.screen = screen
+        self.entities = []
+
+    def update(self, dt) -> None:
+        self.physics.update(dt)
+        
+    def draw(self) -> None:
+        for entity in self.entities:
+            entity.draw(self.screen)
+
+    def spawn_circle(self, radius: float, position: pg.Vector2, velocity: pg.Vector2 | None = None, restitution:float = 0.5, mass: float = 1.0) -> None:
+        if velocity is None:
+            velocity = pg.Vector2(0,0)
+        collider = CircleCollider(radius)
+        body = RigidBody(position=position, collider=collider, velocity=velocity, restitution=restitution,mass=mass)
+        self.physics.add_body(body)
+        sprite = pg.Surface((2*radius, 2*radius),pg.SRCALPHA)
+        sprite.fill((255,255,255,0))
+        pg.draw.circle(
+            sprite,
+            color=(0,255,0,255),
+            center=(radius, radius),
+            radius=radius
+        )
+        self.entities.append(Circle(radius,body,sprite))
+
+def main():
+    running=True
+    pg.init()
+    clock=pg.time.Clock()
+    screen = pg.display.set_mode(size=DIMENSIONS)
+    world=World(screen)
+
+    while running:
+        dt = clock.tick(144) / 1000
+        screen.fill((0,0,0,0))
+        world.update(dt)
+        world.draw()
+        pg.display.flip()
+        for event in pg.event.get():
+            if event.type == pg.MOUSEBUTTONDOWN:
+                world.spawn_circle(
+                    10,
+                    pg.Vector2(pg.mouse.get_pos()),
+                    restitution=random()
+                )
+            if event.type == pg.QUIT:
+                running = False
+        
+if __name__ == "__main__":
+    main()