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 &lt; 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 &lt; projx or abs(size-x[0])-r &lt; projx:
            self._vafter[0] *= -1
        if abs(x[1])-r &lt; projy or abs(size-x[1])-r &lt; 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)
There are a few things that can be improved here.


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

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.
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.
You miss some more whitespace in your step1 function around the != operator. 

if i!=j: 


should be,

if i != j:

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