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/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.