An exercise in Vpython

bouncing-balls.py

import random
import math
from vpython import sphere, vector, scene

RANGE = 30
X_START = -RANGE - RANGE / 2
X_END = RANGE + RANGE / 2
Y_START = -RANGE
Y_END = RANGE
RADIUS = 0.05 * RANGE
N_BALLS = 20
MIN_VELOCITY = 0.01
MAX_VELOCITY = 0.03
VELOCITY_AFTER_COLLISION = 1


scene.range = RANGE


def get_random_position():
    return vector(
        random.uniform(X_START + RADIUS, X_END - RADIUS),
        random.uniform(Y_START + RADIUS, Y_END - RADIUS),
        0,
    )


def get_random_color():
    return vector(random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1))


def get_random_velocity():
    return random.uniform(MIN_VELOCITY, MAX_VELOCITY)


def get_other_balls(ball):
    return [obj for obj in scene.objects if isinstance(obj, sphere) and obj != ball]


class Ball(sphere):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.vel = vector(get_random_velocity(), get_random_velocity(), 0)

    def is_x_offscreen(self):
        is_above_screen = self.pos.x + self.radius >= X_END
        is_below_screen = self.pos.x - self.radius <= X_START
        return is_above_screen or is_below_screen

    def is_y_offscreen(self):
        is_left_of_screen = self.pos.y + self.radius >= Y_END
        is_right_of_screen = self.pos.y - self.radius <= Y_START
        return is_left_of_screen or is_right_of_screen

    def react_to_collision(self):
        self.color = get_random_color()

    def get_distance_from(self, ball):
        distance_x = self.pos.x - ball.pos.x
        distance_y = self.pos.y - ball.pos.y
        return math.sqrt(distance_x ** 2 + distance_y ** 2)

    def collided_with(self, ball):
        distance = self.get_distance_from(ball)
        return distance < (self.radius + ball.radius)

    def move(self):
        if self.is_x_offscreen():
            self.vel.x *= -1
        if self.is_y_offscreen():
            self.vel.y *= -1

        for other_ball in get_other_balls(self):
            if self.collided_with(other_ball):
                self.react_to_collision()
                other_ball.react_to_collision()

                self.vel.x *= -1
                other_ball.vel.y *= -1

        self.pos += self.vel


def create_random_ball():
    return Ball(
        radius=RADIUS,
        pos=get_random_position(),
        color=get_random_color(),
    )


balls = [create_random_ball() for _ in range(N_BALLS)]

while True:
    for ball in balls:
        ball.move()