text: import numpy as np
  
  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 = np.array(position)
          self.velocity = np.array(velocity)
          self.vafter = np.copy(velocity) # temporary 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 computeEnergy(ball_list):
          """Compute kinetic energy."""
          return self.mass / 2. * np.linalg.norm(self.velocity)**2
      
      def compute_coll(self, ball, step):
          """Compute velocity after collision with another ball."""
          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
          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 computation step
          size the medium size
          """
          r, v, x = self.radius, self.velocity, 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 solve_step(ball_list, step, size):
      """Solve a step for every ball."""
      
      # Detect edge-hitting 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)
                  
      # Compute position of every ball  
      for ball in ball_list:
          ball.new_velocity()
          ball.compute_step(step)