1000 Ball Collision Using Spatial Hash

# 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

SPATIAL_GRID_SIZE = 32

class SpatialHash:
    def __init__(s):
        s.grid = [{}, {}]
        s.currentGrid = 0

    def switchGrid(s):
        s.currentGrid = 1 - s.currentGrid
        s.grid[s.currentGrid] = {}

    def insert(s, entity):
        for cell in s._rect_cells(entity.rect()):
            items = s.grid[s.currentGrid].get(cell)
            if items is None:
                s.grid[s.currentGrid][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 = []
        for cell in s._rect_cells(rect):
            items.extend(s.grid[1 - s.currentGrid].get(cell, ()))
        return items

    def query_point(s, pos):
        x, y = pos // SPATIAL_GRID_SIZE
        return s.grid[1 - s.currentGrid].get((x, y), ())

GRAVITY = 0.1

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

class ball:
    def __init__(s, x, y):
        s.pos = pygame.math.Vector2(x, y)
        s.velocity = pygame.math.Vector2(0, 0)

    def rect(s):
        return pygame.Rect(s.pos - BALL_SIZE, BALL_SIZE_DOUBLE)

    def update(s, sh):
        s.pos += s.velocity
        s.velocity.y += GRAVITY
        if s.pos.y > DH - BALL_RADIUS:
            s.pos.y = DH - BALL_RADIUS
        sh.insert(s)

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 = 1000
balls = []
for index in range(BALL_COUNT):
    newBall = ball(
        random.randint(0, DW),
        random.randint(0, DH)
    )
    balls.append(newBall)
    sh.insert(newBall)

exit = False

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

    PD.fill((0, 0, 0))

    sh.switchGrid()
    for b in balls:
        pygame.draw.circle(PD, BALL_COLOR, b.pos, BALL_RADIUS)
        query = b.rect().inflate(BALL_SIZE)
        collisionPoints = sh.query_point(b.pos)
        for p in collisionPoints:
            if p == b: continue
            if b.pos.distance_to(p.pos) <= BALL_RADIUS * 2:
                angle = math.atan2(p.pos.y - b.pos.y, p.pos.x - b.pos.x)
                p.pos = pygame.math.Vector2(
                    b.pos.x + math.cos(angle) * BALL_RADIUS * 2,
                    b.pos.y + math.sin(angle) * BALL_RADIUS * 2
                )
        b.update(sh)

    pygame.display.update()
    #pygame.time.Clock().tick(60)