# Simulation of 2D elastic balls

Following this SO post and this Wikipedia article, I wrote a Python script to simulate physics of 2D elastic balls.

I define the physical behaviour of each ball in the Ball class. Then the function solve_step enables to compute position and velocity of every ball for one time step.

The balls are moving inside a square. Two balls rebound elastically from each other (see compute_coll function). When a ball hits one edge of the square, it bounces off elastically (see compute_refl function).

Note that I'm using 2 functions step1 and step2 for solve_step because that makes the program faster.

Moreover I provide a main to run an example of simulation.

Here I don't show the code used to display the simulation in a window because I think this is another issue.

Any ideas on optimizing and simplifying this code?

import numpy as np


Ball class

class Ball:
"""Define physics of elastic collision."""

def __init__(self, mass, radius, position, velocity):
"""Initialize a Ball object

mass the mass of ball
radius the radius of ball
position the position vector of ball
velocity the velocity vector of ball
"""
self._mass = mass
self._radius = radius
self._position = position
self._velocity = velocity
self._vafter = np.copy(velocity) # temp storage for velocity of next step

def compute_step(self, step):
"""Compute position of next step."""
self._position += step * self._velocity

def new_velocity(self):
"""Store velocity of next step."""
self._velocity = self._vafter

def compute_coll(self, ball, step):
"""Compute velocity after elastic collision with another ball."""
m1 = self._mass
m2 = ball._mass
r1 = self._radius
r2 = ball._radius
v1 = self._velocity
v2 = ball._velocity
x1 = self._position
x2 = ball._position
di = x2-x1
norm = np.linalg.norm(di)
if norm-r1-r2 < step*abs(np.dot(v1-v2,di))/norm:
self._vafter = v1 - 2.*m2/(m1+m2) * np.dot(v1-v2,di)/(np.linalg.norm(di)**2.) * di

def compute_refl(self, step, size):
"""Compute velocity after hitting an edge.

step the step of computation
size the size of a square edge
"""
r = self._radius
v = self._velocity
x = self._position
projx = step*abs(np.dot(v,np.array([1.,0.])))
projy = step*abs(np.dot(v,np.array([0.,1.])))
if abs(x[0])-r < projx or abs(size-x[0])-r < projx:
self._vafter[0] *= -1
if abs(x[1])-r < projy or abs(size-x[1])-r < projy:
self._vafter[1] *= -1.


Solver

def step1(ball_list, step, size):
"""Detect reflection and collision of every ball."""
index_list = range(len(ball_list))
for i in index_list:
ball_list[i].compute_refl(step,size)
for j in index_list:
if i!=j:
ball_list[i].compute_coll(ball_list[j],step)
return ball_list

def step2(ball_list, step):
"""Compute position of every ball."""
index_list = range(len(ball_list))
for i in index_list:
ball_list[i].new_velocity()
ball_list[i].compute_step(step)
return ball_list

def solve_step(ball_list, step, size):
"""Solve a step for every ball."""
ball_list = step1(ball_list, step, size)
ball_list = step2(ball_list, step)
return ball_list


Main

def init_list(N):
"""Generate N Ball objects in a list."""
ball_list = []
r = 10.
for i in range(N):
v = 10.*np.array([(-1.)**i,1.])
pos = 400./float(N+1)*np.array([float(i+1),float(i+1)])
ball_list.append(Ball(r, r, pos, v))
return ball_list

if __name__ == "__main__":
ball_list = init_list(10)
size = 400.
step = 0.05
for i in range(5000):
solve_step(ball_list, step, size)

• I created a specific post about the code used to display the simulation in a window. This is here. Commented Jun 19, 2016 at 0:12
• Is this python 2.7 or 3?
– user95591
Commented Jun 19, 2016 at 1:15
• I'm using python 2.7 but it should work with python 3 if you try. Commented Jun 19, 2016 at 12:39

## 2 Answers

In general for optimising code you should look at a profiler to tell you where the specific hotspots are.

There are also some differences between Python 2.7 and Python 3, but without some confirmation I'll just note that with Python 3 you have a higher number of builtins that generate things on demand rather than building up e.g. whole lists etc. then discarding them immediately. That applies to e.g. the range function.

• In step2 the range(len(ball_list)) is unnecessary. You already have a list, so just iterate over that instead:

def step2(ball_list, step):
"""Compute position of every ball."""
for ball in ball_list:
ball.new_velocity()
ball.compute_step(step)


Since this is purely destructive there's no need to return the list either. It also reads much nicer.

• For step1 again you don't need to return the list. Also, same idea, don't generate a new list with indexes, just do it on the fly:

def step1(ball_list, step, size):
"""Detect reflection and collision of every ball."""
for ball1 in ball_list:
ball1.compute_refl(step, size)
for ball2 in ball_list:
if ball1 is not ball2:
ball1.compute_coll(ball2, step)


Assuming that you don't have any more expensive checks for comparing balls, ball1 is not ball2 should be just a pointer comparison, so reasonably fast. Edit: I update this from the comment; of course is not is better than != for the reason I mentioned already.

For the NumPy things I can just recommend looking at already predefined functions in NumPy/SciPy that are hopefully implemented in C(ython), not recalculating values, i.e. just store every reusable value if possible, possibly looking at other ways of representing your data (matrixes to use fast matrix/matrix operations?) and finally using Cython to run maths directly in C if at all possible.

• Thanks for your answer. It really makes my code simpler :) but it doesn't really improve perf when I check execution time. Maybe I should take a look to Cython, I thought about using fortran and f2py too. Commented Apr 17, 2016 at 14:16
• About ball1 != ball2, I noticed that ball1 is not ball2 is significantly faster. Execution takes 15s instead of 17s. Commented Apr 17, 2016 at 14:21
• My bad, yes, that would be even better. Commented Apr 17, 2016 at 14:35

There are a few things that can be improved here.

1. In Python, using a _ before a class variable denotes a protected member of the class, I don't think that was your intention so you can just make it the following without the preceding _'s in __init__ like this, and throughout the program.

self.mass = mass
self.radius = radius
self.position = position
self.velocity = velocity
self.vafter = np.copy(velocity) # temp storage for velocity of next step

2. Your compute_coil method is mostly just naming variables.

Since we have alot of lines like this,

m1 = self.mass
m2 = ball.mass


we can re-write them on one line like this,

m1, m2 = self.mass, ball.mass


That would change the first part of your compute_coll method to,

m1, m2 = self.mass, ball.mass
r1, r2 = self.radius, ball.radius
v1, v2 = self.velocity, ball.velocity
x1, x2 = self.position, ball.position
di = x2-x1


which I feel reduces a lot of unnecessary clutter.

3. In your compute_coll method you have a little more tidying that can be done.

According to PEP 8,

v1-v2,di


should be,

v1-v2, di


with whitespace after the ,. You also do this in a few of your other methods.

4. You miss some more whitespace in your step1 function around the != operator.

if i!=j:


should be,

if i != j:

5. In your function init_list(N), your argument N is uppercase. According to PEP 8 argument names should be lowercase.