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.


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

def step2(ball_list, step):
    """Compute position of every ball."""
    index_list = range(len(ball_list)) 
    for i in index_list:
    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


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)
  • 1
    \$\begingroup\$ I created a specific post about the code used to display the simulation in a window. This is here. \$\endgroup\$
    – cromod
    Jun 19, 2016 at 0:12
  • \$\begingroup\$ Is this python 2.7 or 3? \$\endgroup\$
    – user95591
    Jun 19, 2016 at 1:15
  • \$\begingroup\$ I'm using python 2.7 but it should work with python 3 if you try. \$\endgroup\$
    – cromod
    Jun 19, 2016 at 12:39

2 Answers 2


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:

    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.

  • \$\begingroup\$ 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. \$\endgroup\$
    – cromod
    Apr 17, 2016 at 14:16
  • 1
    \$\begingroup\$ About ball1 != ball2, I noticed that ball1 is not ball2 is significantly faster. Execution takes 15s instead of 17s. \$\endgroup\$
    – cromod
    Apr 17, 2016 at 14:21
  • \$\begingroup\$ My bad, yes, that would be even better. \$\endgroup\$
    – ferada
    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,


    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.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.