Balls simulation

balls.py

from pymunk import Body, Circle, Vec2d, Segment, Poly, Space
import pymunk
import pygame
import random
from select import select

def vint(v):
    return int(v[0]), int(v[1])

def randarea(x0, y0, w, h):
    return x0 + random.random() * w, y0 + random.random() * h

def randcolor():
    return ([random.randint(0, 155) for i in xrange(3)])

class Balls:
    def __init__(self, rs, w, h):
        pygame.display.init()
        self.screen = pygame.display.set_mode((800, 600))
        self.rs = rs
        self.w = w
        self.h = h

    def render(self):
        screen = self.screen
        screen.fill(0xffffff)
        # render walls
        for wall, seg in self.walls:
            pygame.draw.line(screen, 0x000000, vint(seg.a+wall.position), 
                    vint(seg.b+wall.position), 1)
            # print 'seg', seg.a + wall.position, seg.b + wall.position
        # render balls
        for ball, circle in self.balls:
            pygame.draw.circle(screen, ball.color, vint(ball.position),
                    int(circle.radius))
        pygame.display.flip()

    def solve(self):
        k = 5
        w, h = self.w, self.h
        w1 = w * k
        h1 = h * k
        w2 = w * (k-1)
        h2 = h * (k-1)
        #init a space
        space = Space(5000)
        # walldata = (pos, a, b)
        wallV = 40
        wallM = 1000000
        x0, y0 = 400, 300
        wallsData = [((0, -h1/2), (0, wallV), (-w1/2, 0), (w1/2, 0)),
                ((w1/2, 0), (-wallV, 0), (0, -h1/2), (0, h1/2)),
                ((0, h1/2), (0, -wallV), (-w1/2, 0), (w1/2, 0)),
                ((-w1/2, 0), (wallV, 0), (0, -h1/2), (0, h1/2))]
        walls = []
        # add walls
        for pos, v, a, b in wallsData:
            body = Body(wallM, pymunk.moment_for_segment(wallM, a, b))
            body.position = pos + Vec2d(x0, y0)
            # body.velocity = v
            # body.apply_force(Vec2d(v) * body.mass)
            shape = Segment(body, a, b, 10)
            shape.elasticity = 0.9
            shape.group = 1
            space.add(body, shape)
            walls.append((body, shape))

        balls = []
        # add balls
        V = 100
        for r in self.rs:
            mass = 10 * r ** 2
            body = Body(mass, pymunk.moment_for_circle(mass, 0, r))
            shape = Circle(body, r)
            shape.elasticity = 0.9
            body.position = Vec2d(randarea(-w2/2, -h2/2, w2, h2)) + (x0, y0)
            body.velocity = Vec2d(randarea(-V, -V, V, V))
            body.color = randcolor()
            space.add(body, shape)
            balls.append((body, shape))

        self.walls = walls
        self.balls = balls

        self.render()
        # raw_input()

        SPS = 80
        FPS = 30
        iterDt = 1.0 / SPS 
        fpsDt = 1.0/FPS

        timer = pygame.time.Clock()
        frameTime = 0
        timer.tick()
        # start solve
        dwl, dhl = 10000000, 10000000
        cnt = 0
        while True:
            cnt += 1
            space.step(iterDt)
            dh = abs(walls[0][0].position.y - walls[2][0].position.y)
            dw = abs(walls[1][0].position.x - walls[3][0].position.x)
            if dh > dhl or dw > dwl: break
            dwl, dhl = dw, dh
            for wall,shape in walls:
                wall.reset_forces()
                wall.apply_force(((x0, y0) - wall.position) * 100000)
            frameTime += iterDt
            if frameTime >= fpsDt:
                for e in pygame.event.get():
                    if e.type == pygame.KEYDOWN and e.unicode == ' ':
                        pause = 1
                        while pause:
                            for e in pygame.event.get():
                                if e.type == pygame.KEYDOWN and e.unicode == ' ':
                                    pause = 0
                            pygame.time.delay(10)
                self.render()
                frameTime -= fpsDt
            timer.tick(SPS)
        self.render()
        print 'finish in %d loop' % (cnt,)

if __name__ == '__main__':
    rs = [10, 10, 10, 10, 10,
            5, 5, 5, 5,
            3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3]  
    rs = sum(([r]*(60/r) for r in [5, 10, 12, 15]), [])
    solver = Balls(rs, 200, 200)
    solver.solve()
    raw_input()

Use space bar to pause.