# Python - 2D Elastic Collision Simulation

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'

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")

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()


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.

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/s²! 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.