Simulation of 2D elastic balls

Question

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)

Answer 1

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.

Answer 2

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.