Just a little physics demo

ball.py

from __future__ import division
"""
This module defines a Ball object for performing lossless kinematics over two dimensional Galilean space.

Author: Pyry Pakkanen
"""

__author__ = "Pyry Pakkanen"
__copyright__ = "Copyright 2013"
__credits__ = ["Pyry Pakkanen"]
__license__ = "MIT"
__version__ = "0.1"
__maintainer__ = "Pyry Pakkanen"
__email__ = "[email protected]"
__status__ = "initial release"

from math import sqrt, hypot

class Ball(object):
    """Two dimensional disc with methods that implement Galilean motion with lossless collisions."""
    def __init__(self, r, m, x, y, vx, vy):
        self.r = r
        self.m = m
        self.x = x
        self.y = y
        self.vx = vx
        self.vy = vy
        self.dt = 0.0

    def step(self, dt):
        """Perform linear motion over time dt"""
        self.x += self.vx * dt
        self.y += self.vy * dt
        self.dt -= dt

    def time_of_impact(self, other):
        """Returns the time to impact other or None if no impact will occur."""
        #r_s(t) = (self.x + self.vx*t)*i + (self.t + self.vy*t)*j
        #r_o(t) = (other.x + other.vx*t)*i + (other.y + other.vy*t)*j
        #R(t) = | r_s(t) - r_o(t) |
        #Solve for smallest positive t in R(t) = self.r + other.r
        dx = self.x - other.x
        dy = self.y - other.y
        dvx = self.vx - other.vx
        dvy = self.vy - other.vy
        r = self.r + other.r
        discriminant = -(dvy*dx - dvx*dy)**2 + (dvx**2 + dvy**2) * r**2
        if discriminant < 0:
            return None

        sqrt_discriminant = sqrt(discriminant)
        a = dvx**2 + dvy**2
        b = -dvx*dx - dvy*dy 
        min_t = (b - sqrt_discriminant) / a
        max_t = (b + sqrt_discriminant) / a
        if min_t > 0:
            return min_t
        if max_t > 0:
            return max_t

        return None

    def collide(self, other):
        """Update velocities of self and other in a lossless collision."""
        # In a center-of-mass coordinates of a two body collision between spheres the
        # speeds are conserved and both objects escape away from the point of impact.
        sm = self.m + other.m
        center_of_mass_vx = (self.m * self.vx + other.m * other.vx) / sm
        center_of_mass_vy = (self.m * self.vy + other.m * other.vy) / sm

        self_speed = hypot(self.vx - center_of_mass_vx, self.vy - center_of_mass_vy)
        other_speed = hypot(other.vx - center_of_mass_vx, other.vy - center_of_mass_vy)

        dx = self.x - other.x
        dy = self.y - other.y
        distance = hypot(dx, dy)
        dx /= distance
        dy /= distance

        self.vx = center_of_mass_vx + self_speed*dx
        self.vy = center_of_mass_vy + self_speed*dy
        other.vx = center_of_mass_vx - other_speed*dx
        other.vy = center_of_mass_vy - other_speed*dy

    def distance(self, other):
        return hypot(self.x - other.x, self.y - other.y)

    @property
    def momentumx(self):
        return self.vx * self.m

    @property
    def momentumy(self):
        return self.vy * self.m

    @property
    def speed(self):
        return hypot(self.vx, self.vy)

    @property
    def energy(self):
        return self.m * (self.vx**2 + self.vy**2)

    def __repr__(self):
        return "Ball{}".format((self.r, self.m, self.x, self.y, self.vx, self.vy))

epsilon = 1e-10

def float_equals(a, b):
    return abs(a - b) < epsilon

def test_impact():
    b0 = Ball(1, 1, 0, 0.5, 0, 0)
    b1 = Ball(1, 1, 4, 0, -1, 0)
    dt = b0.time_of_impact(b1)
    b0.step(dt)
    b1.step(dt)
    assert(float_equals(b0.r + b1.r, b0.distance(b1)))

def test_collision():
    b0 = Ball(2, 1, 0, 0, 0, 1)
    b1 = Ball(3, 2, 3, 4, -1, 0)
    initial_momentumx = b0.momentumx + b1.momentumx
    initial_momentumy = b0.momentumy + b1.momentumy
    initial_energy = b0.energy + b1.energy
    b0.collide(b1)
    assert(float_equals(b0.momentumx + b1.momentumx, initial_momentumx))
    assert(float_equals(b0.momentumy + b1.momentumy, initial_momentumy))
    assert(float_equals(b0.energy + b1.energy, initial_energy))


if __name__ == "__main__":
    # Perform a kinematic simulation animated with pylab

    # Sanity check
    test_impact()
    test_collision()

    from pylab import *
    N = 20  # Number of Balls for the simulation
    xs = random(N)
    ys = random(N)
    angles = random(N)*2*pi
    vs = random(N)
    xvs = [cos(angle)*v for angle, v in zip(angles, vs)]
    yvs = [sin(angle)*v for angle, v in zip(angles, vs)]
    bs = [Ball(0.01, 1, x, y, xv, yv) for (x, y, xv, yv) in zip(xs, ys, xvs, yvs)]

    ion()
    xlim(0, 1)
    ylim(0, 1)
    axes().set_aspect('equal')

    line, = plot(xs, ys, "o")

    for i in range(500):
        for b in bs:
            b.dt = 0.01

        def calculate_collisions():
            ti_pairs = []
            for i in range(N):
                b = bs[i]
                for other in bs[i + 1:]:
                    ti = b.time_of_impact(other)
                    if ti is not None:
                        ti_pairs.append((ti, b, other))
            return ti_pairs

        def collide_all():
            ti_pairs = calculate_collisions()
            if not ti_pairs:
                return
            ti, first, second = min(ti_pairs)
            if epsilon < ti < first.dt:
                for b in bs:
                    b.step(ti)
                    b.x = b.x - floor(b.x)
                    b.y = b.y - floor(b.y)
                first.collide(second)
                collide_all()

        collide_all()
        for b in bs:
            b.step(b.dt)
            b.x = b.x - floor(b.x)
            b.y = b.y - floor(b.y)
        xs = [b.x for b in bs]
        ys = [b.y for b in bs]
        line.set_xdata(xs)
        line.set_ydata(ys)

        draw()