Python - 2D Elastic Collision Simulation

Question

I am currently a Python beginner, who just finished a script simulating collisions between circles (balls) in 2 dimensions.
I would appreciate comments on overall structure, variable names and really anything else.

There are bug present, such as a ball occasionally slipping out of a corner, but the code is functional
Note: The code is auto-formatted with 'black'
Link to Gist

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from random import randint, uniform
class Simulation:
 def __init__(
 self, size, ball_count, max_v_factor, r_range=None, gravity=False, save=False
 ):
 self.size = size
 self.balls = []
 self.max_v_factor = max_v_factor
 self.r_range = r_range
 self.gravity = gravity
 self.fig = plt.figure()
 self.ax = plt.axes()
 self.save = save
 # Generate list of ball objects with random values
 for _ in range(ball_count):
 pos = np.array([randint(10, self.size - 10), randint(10, self.size - 10)])
 v = np.array(
 [
 uniform(-size / self.max_v_factor, size / self.max_v_factor),
 uniform(-size / self.max_v_factor, size / self.max_v_factor),
 ]
 )
 if self.r_range:
 r = uniform(self.r_range[0], self.r_range[1])
 else:
 r = size / 20
 # Calculating mass based on simplified formula for volume of sphere
 m = r ** 3
 self.balls.append(Ball(pos, m, v, r))
 def update(self):
 """Updates the positions and velocities of the balls"""
 for i, ball in enumerate(self.balls):
 ball.move(self.gravity)
 # Get ball's properties
 ball_pos, _, _, ball_r = ball.get_properties()
 for other in self.balls[i + 1 :]:
 # Get other ball's properties
 other_pos, _, _, other_r = other.get_properties()
 # Find the scalar difference between the position vectors
 diff_pos = np.linalg.norm(ball_pos - other_pos)
 # Check if the balls' radii touch. If so, collide
 if diff_pos <= ball_r + other_r:
 ball.collide(other)
 diff_vec = ball_pos - other_pos
 # 'clip' is how much the balls are clipped
 clip = ball_r + other_r - diff_pos
 # Creating normal vector between balls
 diff_norm = diff_vec / diff_pos
 # Creating 'clip_vec' vector that moves ball out of the other
 clip_vec = diff_norm * clip
 # Set new position
 ball.set_pos(ball_pos + clip_vec)
 # Check if the ball's coords is out of bounds
 # X-coord
 if ball_pos[0] - ball_r < 0:
 ball.bounce_x(-ball_pos[0] + ball_r)
 elif ball_pos[0] + ball_r > self.size:
 ball.bounce_x(self.size - ball_pos[0] - ball_r)
 # Y-coord
 elif ball_pos[1] - ball_r < 0:
 ball.bounce_y(-ball_pos[1] + ball_r)
 elif ball_pos[1] + ball_r > self.size:
 ball.bounce_y(self.size - ball_pos[1] - ball_r)
 def animate(self, _):
 """Updates and returns plot. To be used in Matplotlib's FuncAnimation"""
 self.ax.patches = []
 self.update()
 for ball in self.balls:
 pos, _, _, r = ball.get_properties()
 circle = plt.Circle(pos, r, color="blue")
 self.ax.add_patch(circle)
 return (self.ax,)
 def show(self):
 """Draws the table and balls"""
 self.ax.set_aspect(1)
 self.ax.set_xlim(0, self.size)
 self.ax.set_ylim(0, self.size)
 anim = animation.FuncAnimation(
 self.fig, self.animate, frames=1500, interval=5, save_count=1500
 )
 if self.save:
 writer = animation.FFMpegWriter(fps=60)
 anim.save("animation.mp4", writer=writer)
 plt.show()
class Ball:
 def __init__(self, pos, m, v, r):
 """Initialises ball with a position, mass and velocity and radius"""
 self.pos = pos
 self.m = m
 self.v = v
 self.r = r
 def move(self, gravity=False):
 """Moves ball 'v' distance."""
 if gravity:
 self.v[1] -= gravity / 20
 self.pos[0] += self.v[0]
 self.pos[1] += self.v[1]
 def collide(self, other):
 """Computes new velocities for two balls based on old velocities"""
 # Get other ball's properties
 other_pos, other_m, other_v, _ = other.get_properties()
 # Create a normal vector to the collision surface
 norm = np.array([other_pos[0] - self.pos[0], other_pos[1] - self.pos[1]])
 # Convert to unit vector
 unit_norm = norm / (np.sqrt(norm[0] ** 2 + norm[1] ** 2))
 # Create unit vector tagent to the collision surface
 unit_tang = np.array([-unit_norm[1], unit_norm[0]])
 # Project self ball's velocity onto unit vectors
 self_v_norm = np.dot(self.v, unit_norm)
 self_v_tang = np.dot(self.v, unit_tang)
 # Project other ball's velocity onto unit vectors
 other_v_norm = np.dot(other_v, unit_norm)
 other_v_tang = np.dot(other_v, unit_tang)
 # Use 1D collision equations to compute the balls' normal velocity
 self_v_prime = ((self.m - other_m) / (self.m + other_m)) * self_v_norm + (
 (2 * other_m) / (self.m + other_m)
 ) * other_v_norm
 other_v_prime = ((2 * self.m) / (self.m + other_m)) * self_v_norm + (
 (other_m - self.m) / (self.m + other_m)
 ) * other_v_norm
 # Add the two vectors to get final velocity vectors and update.
 self.v = self_v_prime * unit_norm + self_v_tang * unit_tang
 other.set_v(other_v_prime * unit_norm + other_v_tang * unit_tang)
 def bounce_x(self, distance):
 """Bounces off and edge parallel to the x axis.
 Variable, 'distance', is distance to move back out of wall"""
 self.v[0] *= -1
 self.pos[0] += distance
 def bounce_y(self, distance):
 """Bounces off and edge parallel to the x axis.
 Variable, 'distance', is distance to move back out of wall"""
 self.v[1] *= -1
 self.pos[1] += distance
 def get_properties(self):
 """Returns the position, mass and velocity and radius"""
 return self.pos, self.m, self.v, self.r
 def get_pos(self):
 """Returns the position only"""
 return self.pos
 def set_pos(self, pos):
 """Sets the new position"""
 self.pos = pos
 def set_v(self, v):
 """Sets the velocity of the ball"""
 self.v = v
if __name__ == "__main__":
 sim = Simulation(
 size=500,
 ball_count=25,
 max_v_factor=150,
 r_range=(12, 15),
 gravity=6,
 save=False,
 )
 sim.show()

Answer 1

First, welcome to CodeReview! Second, let's tear apart your code ;-)

Python is not Java

To me, this code reads like Java code. It's not "pythonic" in a bunch of small ways.

The most obvious is the getter/setter mentality. Python is a "consenting adult language." (Watch Raymond Hettinger's talk for more.)

Instead of writing get_foo() and set_foo() just create an attribute foo and let users get or set it on their own. If you absolutely need to intercept those accesses later, you can use the @property decorator to insert a method call.

Eliminate comments!

You have a good commenting style. You are not blindly repeating the Python code in English prose. But comments are a code smell! Either you felt something was unclear, or you felt that it was inobvious.

Follow your instincts! If something is unclear or inobvious, give it a name! Instead of a comment like:

# Find the scalar difference between the position vectors
diff_pos = np.linalg.norm(ball_pos - other_pos)

why not write a helper function or helper method to formally bind the description in the comment to that code?

def scalar_difference_between_position_vectors(a, b):
 """Return the scalar difference..."""
 return np.linalg.norm(a - b)

Then update your code and delete the comment:

diff_pos = scalar_difference_between_position_vectors(ball_pos, other_pos)

I'll argue that you can tidy up your code a lot by simply replacing your comments with functions. That includes the giant loop in Simulation.__init__

Names have power

Be careful in your choice of names. Consider this line:

ball.move(self.gravity)

Seriously? What does that do? I was perfectly okay with ball.move( -- I totally felt like that was a good name and that I understood where you were going. I expected to see a velocity vector get passed in, or maybe an absolute destination. But why would you pass in gravity? Now I'm lost and uncertain. I feel like I don't understand this code at all!

Worse, when I looked up the definition, I found this:

def move(self, gravity=False):
 """Moves ball 'v' distance."""
 if gravity:
 self.v[1] -= gravity / 20
 self.pos[0] += self.v[0]
 self.pos[1] += self.v[1]

Ah, of course, gravity isn't something like 9.8 m/s2! It's actually a boolean... that can be divided by 20 ... What's 20?

So there are a couple of problems here. First, you need to take more care with the names you are giving functions. Instead of ball.move how about ball.update or ball.tick or ball.update_position?

I'll skip over the whole "what is gravity?" question until the next point. Instead, let's focus on 20.

20 in this context is a magic number. And magic numbers are subject to one simple but absolute rule:

THERE CAN BE ONLY ONE!

And that one is 0. Any bare number other than zero probably deserves a name. Possibly something like this:

TICKS_PER_SECOND = 20

Use None to check for optional args

Snap back to reality! Oh!
What was gravity? Oh!
That's a parameter: False
Or a bool, you see!
But then we divide by twenty!
There goes @JSog', he
Choked, he's so mad but he
Won't give up that easy!
No, he won't have it!

With all due respect to Mr. Mathers, optional parameters are simple if you treat them simply. Just use None and do your test using is [not] None:

def update_position(delta_t: float, gravity: Optional[float]=None):
 if gravity is not None:
 self.v[VERTICAL] -= gravity * delta_t
 self.pos[HORIZONTAL] += self.v[HORIZONTAL] * delta_t
 self.pos[VERTICAL] += self.v[VERTICAL] * delta_t

Don't try to rely on the equivalence between bool and int to get a zero. Just spell out the check. And even better, spell out the expected type using type annotations.