Bouncing balls in Pygame

bounce.py

# 1000 balls find an equilibrium spatial hash collision
# credit to Daniel Pope for showing me the light :)
# fb.com/groups/pygame

import pygame
import random, math
from itertools import product
from pygame.math import Vector2 as v2
import colorsys

SPATIAL_GRID_SIZE = 32

class SpatialHash:
    def __init__(self):
        self.grid = {}
        self.items = set()

    def rebuild(self):
        self.grid = {}
        for item in self.items:
            self.insert(item)

    def insert(self, entity):
        self.items.add(entity)
        for cell in self._rect_cells(entity.rect):
            items = self.grid.get(cell)
            if items is None:
                self.grid[cell] = [entity]
            else:
                items.append(entity)

    def _rect_cells(s, rect):
        x1, y1 = rect.topleft
        x1 //= SPATIAL_GRID_SIZE
        y1 //= SPATIAL_GRID_SIZE
        x2, y2 = rect.bottomright
        x2 = x2 // SPATIAL_GRID_SIZE + 1
        y2 = y2 // SPATIAL_GRID_SIZE + 1
        return product(range(x1, x2), range(y1, y2))

    def query(s, rect):
        items = set()
        for cell in s._rect_cells(rect):
            items.update(s.grid.get(cell, ()))
        return items

GRAVITY = 0.1

BALL_RADIUS = 15
BALL_SIZE = pygame.math.Vector2(BALL_RADIUS)
BALL_SIZE_DOUBLE = BALL_SIZE * 2
BALL_COLOR = (34, 128, 75)


def random_color():
    return tuple(
        round(c * 255) for c in colorsys.hsv_to_rgb(random.random(), 1, 1)
    )


class ball:
    def __init__(self, x, y, radius=BALL_RADIUS):
        self.color = random_color()

        self.velocity = pygame.math.Vector2(0, 0)
        self.radius = radius
        self.mass = radius * radius
        dr = v2(radius, radius)
        self._pos = v2(x, y)
        self.rect = pygame.Rect(
            (*self._pos - dr),
            *(dr * 2)
        )

    @property
    def pos(self):
        return self._pos

    @pos.setter
    def pos(self, pos):
        self._pos = pos
        self.rect.center = pos

    def update(self):
        self.pos += self.velocity
        self.velocity.y += GRAVITY
        if self.pos.y > DH - self.radius:
            self.pos.y = DH - self.radius
            self.velocity.y = -ELASTICITY * abs(self.velocity.y)
            self.velocity.x *= 0.95

    def collides(self, ano):
        minsep = ano.radius + self.radius
        return self.pos.distance_squared_to(ano.pos) < minsep * minsep


DW, DH = 1280, 720
HDW, HDH = DW // 2, DH // 2
DR = pygame.Rect((0, 0), (DW, DH))
pygame.init()
PD = pygame.display.set_mode(DR.size)

sh = SpatialHash()

BALL_COUNT = 100
balls = []
for index in range(BALL_COUNT):
    newBall = ball(
        x=random.randint(0, DW),
        y=random.randint(0, DH),
        radius=random.randint(8, 30)
    )
    balls.append(newBall)
    sh.insert(newBall)

exit = False


ELASTICITY = 0.8


def apply_impact(a, b):
    """Resolve the collision between two balls.

    Calculate their closing momentum and apply a fraction of it back as impulse
    to both objects.
    """
    ab = b.pos - a.pos
    ab.normalize_ip()
    rel_momentum = ab.dot(a.velocity) * a.mass - ab.dot(b.velocity) * b.mass

    rel_momentum *= ELASTICITY

    a.velocity -= ab * rel_momentum / a.mass
    b.velocity += ab * rel_momentum / b.mass


def separate(a, b, frac=0.66) -> bool:
    """Move a and b apart.

    frac is the amount of the overlap to clear; this should be in (0, 1]
    but somewhere in the middle is better for stability.

    Return True if they are now separate.
    """
    ab = a.pos - b.pos
    sep = ab.length()
    overlap = a.radius + b.radius - sep
    if overlap <= 0:
        return True
    ab /= sep
    masses = a.mass + b.mass

    overlap *= frac  # don't try to clear the overlap completely this iteration
    a.pos += ab * (overlap * b.mass / masses)
    b.pos -= ab * (overlap * a.mass / masses)
    return False


def draw():
    PD.fill((0, 0, 0))
    for b in balls:
        pygame.draw.circle(PD, b.color, b.pos, b.radius)
    pygame.display.update()

collisions = set()


def update():
    global collisions
    for b in balls:
        b.update()
    sh.rebuild()

    # Find all collisions occurring this frame
    prev_collisions = collisions
    collisions = set()
    for b in balls:
        possible_collisions = sh.query(b.rect)
        for a in possible_collisions:
            if a is b:
                continue
            if a.collides(b):
                if id(a) > id(b):
                    pair = b, a
                else:
                    pair = a, b

                if pair not in prev_collisions:
                    # We only apply the bounce to the velocity the first time
                    # they collide.
                    apply_impact(*pair)
                collisions.add(pair)

    # Apply several iterations to separate the collisions
    for _ in range(10):
        collisions = {(a, b) for a, b in collisions if separate(a, b)}
        if not collisions:
            break

draw()
while True:
    if any(e.type == pygame.KEYDOWN for e in pygame.event.get()):
        break
    pygame.time.Clock().tick(60)


while True:
    for e in pygame.event.get():
        if e.type == pygame.QUIT:
            exit = True

    if exit or pygame.key.get_pressed()[pygame.K_ESCAPE]: break

    update()
    draw()
    pygame.time.Clock().tick(60)