Balls hitting

yyan.py

import pygame as pg

R = 10
M = 1
FPS = 30

Margin = 10

ScreenW, ScreenH = 400, 400

BoundX1, BoundX2 = Margin, ScreenW - Margin
BoundY1, BoundY2 = Margin, ScreenH - Margin

DT0 = 0.02

class Ball:
    def __init__(self, x, y, vx, vy, r=R):
        self.x = x + 0.
        self.y = y + 0.
        self.vx = vx + 0.
        self.vy = vy + 0.
        self.r = r + 0.

def is_hitting(b1, b2):
    dx = b2.x - b1.x
    dy = b2.y - b1.y
    dvx = b2.vx - b1.vx
    dvy = b2.vy - b1.vy
    if dx**2 + dy ** 2 < (b1.r + b2.r)**2 and dvx * dx + dvy * dy < 0:
        return True
    else:
        return False

def handle_hit(b1, b2):
    sx = b2.x - b1.x
    sy = b2.y - b1.y
    sl = (sx**2 + sy**2)**0.5
    assert sl != 0
    # normalize vector s
    sx /= sl
    sy /= sl

    # t is the perp of s
    tx = -sy
    ty = sx

    v1 = b1.vx * sx + b1.vy * sy 
    v2 = b2.vx * sx + b2.vy * sy
    # v1 + v2 = v1' + v2'
    # v2 - v1 = - (v2' - v1')
    # => v1' = v2, v2' = v1, swap velocity
    v1p = v2
    v2p = v1
    # b1.v' = b1.v - v1 * s + v1' * s = v1.v + (v1'-v1) * s
    b1.vx = b1.vx + (v1p - v1) * sx
    b1.vy = b1.vy + (v1p - v1) * sy

    b2.vx = b2.vx + (v2p - v2) * sx
    b2.vy = b2.vy + (v2p - v2) * sy
    
def is_exist_hit():
    for i in xrange(len(balls)):
        b1 = balls[i]
        for j in xrange(i+1, len(balls)):
            b2 = balls[j]
            if is_hitting(b1, b2):
                return True
    return False

def init_graphic():
    pg.display.init()
    global screen
    screen = pg.display.set_mode((ScreenW, ScreenH), 0, 32)
    screen.fill((0xff, 0xff, 0xff, 0xff))

def init_balls():
    global balls
    balls = []
    v = 100.
    balls.append(Ball(100, 100, v, 0, R*2))
    balls.append(Ball(170, 105, -v, 0))
    balls.append(Ball(170, 305, -v, v, R*2))
    balls.append(Ball(270, 305, v, -v))
    balls.append(Ball(190, 105, v/2, -v))

def draw_balls():
    for ball in balls:
        pg.draw.circle(screen, (0, 0, 0, 0xff), 
                (int(ball.x), int(ball.y)), int(ball.r), 1)
        pg.draw.line(screen, (0xff, 0, 0, 0xff),
                (int(ball.x), int(ball.y)), (int(ball.x + ball.vx), int(ball.y + ball.vy)))
    pg.draw.rect(screen, (0, 0, 0, 0xff),
            pg.Rect((BoundX1, BoundY1), (BoundX2 - BoundX1, BoundY2 - BoundY1)), 1)

def step(dt):
    for ball in balls:
        ball.x += ball.vx * dt
        ball.y += ball.vy * dt

def bound_hit_handle():
    for ball in balls:
        if ball.x - ball.r <= BoundX1 and ball.vx <= 0:
            ball.vx *= -1
        elif ball.x + ball.r >= BoundX2 and ball.vx >= 0:
            ball.vx *= -1
        if ball.y - ball.r <= BoundY1 and ball.vy <= 0:
            ball.vy *= -1
        elif ball.y + ball.r >= BoundY2 and ball.vy >= 0:
            ball.vy *= -1

def to_next_hit(dt):
    tl, tr = 0., dt
    t = 0. # t is current time
    step(dt)
    t = dt
    if not is_exist_hit():
        return dt
    while tr - tl > 0.001:
        tm = (tr + tl) / 2
        step(tm - t)
        t = tm
        if is_exist_hit():
            tr = tm
        else:
            tl = tm
    return t

def simulate(dt):
    while dt > 0:
        ddt = min(dt, DT0)
        ddtr = to_next_hit(ddt)
        if ddtr < ddt:
            for i in xrange(len(balls)):
                for j in xrange(i+1, len(balls)):
                    if is_hitting(balls[i], balls[j]):
                        handle_hit(balls[i], balls[j])
        bound_hit_handle()
        dt -= ddtr

def main_loop():
    timer = pg.time.Clock()
    quit = 0
    while not quit:
        for e in pg.event.get():
            if e.type == pg.QUIT or e.type == pg.KEYDOWN and e.key in (pg.K_q, pg.K_ESCAPE):
                quit = 1
        simulate(1.0/FPS)
        screen.fill((0xff, 0xff, 0xff, 0xff))
        draw_balls()
        pg.display.flip()
        timer.tick(FPS)

init_graphic()
init_balls()
main_loop()