# Some preliminary stuff:
from manimlib.imports import * # Import Manim Python library
def add(v1, v2):
return [v1[0] + v2[0], v1[1] + v2[1]] # Matrix vector addition with two-dimensional vectors
def mul(v, s):
return [v[0] * s, v[1] * s] # Scale a 2d vector by a constant
def dot(v1, v2):
return v1[0] * v2[0] + v1[1] * v2[1] # Dot product for 2d vectors
# Noticeably, I used simple lists in Python to represent vectors because numpy arrays are slow.
# Most of the effort in making the code was in optimizing. The one hour and ten minute render time is the
# fastest I could make it.
# Taking the square root for Pythagorean distance slows down the code, so I opted to return squared values.
def squared_distance(a, b):
return (a[0] - b[0]) * (a[0] - b[0]) + (a[1] - b[1]) * (a[1] - b[1])
# Class AABB is simply an object that stores parameters of a rectangle (center, width and height).
class AABB:
def __init__(self, center, width, height):
self.center = center
self.width = width
self.height = height
def contains(self, mob):
return self.center[0] - self.width <= mob.pos[0] < self.center[0] + self.width and self.center[1] \
- self.height <= mob.pos[1] < self.center[1] + self.height
def intersects(self, rect):
return not (self.center[0] - self.width >= rect.center[0] + rect.width or self.center[0] + self.width <=
rect.center[0] - rect.width or self.center[1] + self.height <= rect.center[1] - rect.height or
self.center[1] - self.height >= rect.center[1] + rect.height)
def encompass(self, rect):
return self.center[0] - self.width <= rect.center[0] - rect.width and self.center[1] - self.height <= \
rect.center[1] - rect.height and rect.center[0] + rect.width <= self.center[0] + self.width and \
rect.center[1] + rect.height <= self.center[1] + self.height
# Quad trees are the heart of optimizing ball-to-ball collision detection.
class QuadTree:
def __init__(self, boundary, cap):
self.boundary = boundary
self.cap = cap
self.array = []
self.divided = False
def subdivide(self):
x = self.boundary.center[0]
y = self.boundary.center[1]
w = self.boundary.width / 2
h = self.boundary.height / 2
self.northwest = QuadTree(AABB((x - w, y + h, 0), w, h), self.cap)
self.northeast = QuadTree(AABB((x + w, y + h, 0), w, h), self.cap)
self.southeast = QuadTree(AABB((x + w, y - h, 0), w, h), self.cap)
self.southwest = QuadTree(AABB((x - w, y - h, 0), w, h), self.cap)
self.divided = True
def insert(self, mob):
if not self.boundary.contains(mob):
return
if len(self.array) < self.cap:
self.array.append(mob)
else:
if not self.divided:
self.subdivide()
self.northwest.insert(mob)
self.northeast.insert(mob)
self.southeast.insert(mob)
self.southwest.insert(mob)
def query(self, rect):
found = []
if not self.boundary.intersects(rect):
return found
if rect.encompass(self.boundary):
found.extend(self.array)
else:
for p in self.array:
if rect.contains(p):
found.append(p)
if self.divided:
found.extend(self.northwest.query(rect))
found.extend(self.northeast.query(rect))
found.extend(self.southeast.query(rect))
found.extend(self.southwest.query(rect))
return found
# Particle class carries all information about a ball's position, velocity, and other useful variables.
class Particle(Dot):
def __init__(self, position_init, vector_init, bounce_init=0.5, **kwargs):
self.pos = position_init
self.vector = vector_init
self.bounce_coeff = bounce_init
self.category = 0
self.funnelled = False
self.airborne = True
Dot.__init__(self, point=np.array(position_init + [0]), **kwargs)
# This is where all animation begins.
class scene_1(Scene):
CONFIG = { # Naming constants for convenience
"runtime": 20,
"particle_amount": 1000,
"gravity": -9.8,
"steps_per_frame": 10,
"obstacle_rows": 17,
"obstacle_init_columns": 8,
"obstacle_radius": DEFAULT_SMALL_DOT_RADIUS * 1.3,
"obs_color": "#888888",
"p_radius": DEFAULT_SMALL_DOT_RADIUS
}
def construct(self):
self.particles = VGroup()
self.obstacles = VGroup()
self.initialize()
self.add(self.particles, self.obstacles)
self.wait(self.runtime)
def initialize(self): # This up until line 197 sets up all objects on the screen, and initial conditions.
random.seed()
self.obstacles.add(Arc(start_angle=0, angle=PI / 4, stroke_width=3, radius=1, color=self.obs_color),
Arc(start_angle=PI, angle=-PI / 4, stroke_width=3, radius=1, color=self.obs_color))
self.obstacles[0].move_arc_center_to((-2 * self.p_radius - 1, 3.35, 0))
self.obstacles[1].move_arc_center_to((2 * self.p_radius + 1, 3.35, 0))
dist = 2 * self.obstacle_radius + 3 * self.p_radius
height = dist * 1.732 / 2
for i in range(self.obstacle_rows):
for j in range(self.obstacle_init_columns + i):
self.obstacles.add(Particle([dist * (j - (self.obstacle_init_columns + i - 1) / 2), 3.15 - i * height],
0, color=self.obs_color, radius=self.obstacle_radius))
last_row_columns = self.obstacle_init_columns + self.obstacle_rows - 1
num_obstacles = int(self.obstacle_rows * (self.obstacle_init_columns + last_row_columns) / 2)
x_barriers = []
for i in range(2 + num_obstacles - last_row_columns, 2 + num_obstacles):
line_1 = Line((self.obstacles[i].pos[0] - self.obstacle_radius * 0.75,
self.obstacles[i].pos[1] - dist, 0), (self.obstacles[i].pos[0] -
self.obstacle_radius * 0.75, -5, 0),
color=self.obs_color, stroke_width=1, stroke_opacity=0.5)
line_2 = Line((self.obstacles[i].pos[0] + self.obstacle_radius * 0.75,
self.obstacles[i].pos[1] - dist, 0), (self.obstacles[i].pos[0] +
self.obstacle_radius * 0.75, -5, 0),
color=self.obs_color, stroke_width=1, stroke_opacity=0.5)
arc_1 = ArcBetweenPoints(np.array([self.obstacles[i].pos[0], self.obstacles[i].pos[1] -
dist / 3, 0]), np.array([self.obstacles[i].pos[0] -
self.obstacle_radius * 0.75,
self.obstacles[i].pos[1] - dist, 0]),
PI * 5 / 36, color=self.obs_color, stroke_width=1, stroke_opacity=0.5)
arc_2 = ArcBetweenPoints(np.array([self.obstacles[i].pos[0] + self.obstacle_radius * 0.75,
self.obstacles[i].pos[1] - dist, 0]),
np.array([self.obstacles[i].pos[0], self.obstacles[i].pos[1] -
dist / 3, 0]), PI * 5 / 36, color=self.obs_color, stroke_width=1,
stroke_opacity=0.5)
black = Rectangle(width=1.5 * self.p_radius, height=self.obstacles[i].pos[1] - dist + 5,
fill_opacity=1, fill_color=BLACK, stroke_opacity=0)
black.move_to((self.obstacles[i].pos[0], self.obstacles[i].pos[1] - black.get_height() /
2 - 0.25, 0))
x_barriers.append(self.obstacles[i].pos[0])
self.add(black, line_1, line_2, arc_1, arc_2)
decorations = VGroup()
ref_point = np.array(self.obstacles[2 + num_obstacles - last_row_columns].pos + [0])
decorations.add(FunctionGraph(lambda x: 1.732 * (x - ref_point[0] + dist) + ref_point[1] - 0.05,
x_min=ref_point[0] - dist + 0.1, stroke_width=1, color=self.obs_color,
stroke_opacity=0.5))
decorations.add(ArcBetweenPoints(ref_point + np.array([-dist + 0.1, 0.1 * 1.732 - 0.05, 0]), ref_point +
np.array([-dist, -0.25, 0]), PI / 6, stroke_width=1, color=self.obs_color,
stroke_opacity=0.5))
decorations.add(Line(ref_point + np.array([-dist, -0.25, 0]), ref_point + np.array([-dist, -0.25, 0]) + DOWN
* 5, stroke_width=1, color=self.obs_color, stroke_opacity=0.5))
decorations.add(Line(ref_point + np.array([-1.5 * dist, 5, 0]), ref_point + np.array([-1.5 * dist, -2.05, 0]),
color=self.obs_color))
decorations.add(ArcBetweenPoints(ref_point + np.array([-1.5 * dist - 2, -4.05, 0]), ref_point +
np.array([-1.5 * dist, -2.05, 0]), PI / 2, color=self.obs_color))
decorations.add(Line(ref_point + np.array([-dist, 0.2, 0]), ref_point + LEFT * dist + UP * 5,
stroke_width=1, color=self.obs_color, stroke_opacity=0.5))
decorations.add(ArcBetweenPoints(ref_point + np.array([-dist, 0.2, 0]), ref_point +
np.array([-dist + 7 * (1 - 1.732 / 2), 4, 0]), PI / 6, stroke_width=1,
color=self.obs_color, stroke_opacity=0.5).shift(UP * 1.732 + RIGHT))
decorations.add(Line(ref_point + np.array([-dist, 0.2, 0]), ref_point +
np.array([-dist + 1, 0.2 + 1.732, 0]), stroke_width=1, color=self.obs_color,
stroke_opacity=0.5))
self.add(decorations)
self.add(decorations.copy().flip(UP).move_to(decorations.get_center() -
RIGHT * 2 * decorations.get_center()[0]))
self.counter = 0
# This is where everything about gravity, collision, and bouncing off obstacles is simulated.
def updater(mob, dt):
if self.counter < self.particle_amount: # Randomly spawn a new particle for every subsequent frame
particle_colors = [MAROON_A, MAROON_B, MAROON_C, MAROON_D, MAROON_E]
mob.add(Particle([-(0.28 + self.p_radius) * (1 - 2 * random.random()), 4], [0, self.gravity / 60],
radius=self.p_radius, color=particle_colors[random.randint(0, 4)]))
self.counter += 1
for i in range(self.counter):
mob[i].move_to(np.array(mob[i].pos + [0])) # Update every ball's position on the screen
if mob[i].pos[1] >= 3.15 and not mob[i].funnelled: # Check if a ball has entered the Galton Board
arc_center = self.obstacles[math.ceil(mob[i].pos[0])].get_arc_center().tolist()
delta = squared_distance(arc_center, mob[i].pos)
if delta <= (self.p_radius + 1) * (self.p_radius + 1):
delta = math.sqrt(delta)
unit_collision = mul(add(arc_center, mul(mob[i].pos, -1)), 1 / delta)
mob[i].vector = add(mob[i].vector, mul(unit_collision, -(1 + mob[i].bounce_coeff) *
dot(mob[i].vector, unit_collision)))
mob[i].pos = add(mob[i].pos, mul(unit_collision, delta - 1 - self.p_radius))
else:
mob[i].funnelled = True
if mob[i].pos[1] <= ref_point[1]: # Check if a ball is done falling through obstacles
if mob[i].category == 0:
mob[i].category = (mob[i].pos[0] - ref_point[0]) / dist
x_thres = -(mob[i].pos[1] - ref_point[1] + 0.05) / 1.732 - ref_point[0] + dist - 2 * \
self.p_radius / 1.732
if not abs(mob[i].pos[0]) < x_thres: # Check if a ball hits the walls on the side
index = 2 * (mob[i].pos[0] > 0) - 1
unit_collision = [index * 1.732 / 2, 1 / 2]
mob[i].vector = add(mob[i].vector, mul(unit_collision, -2 * dot(mob[i].vector,
unit_collision)))
mob[i].pos[0] -= index * (abs(mob[i].pos[0]) - x_thres)
x_ref = self.obstacles[2].pos[0] # Up until line 249, check if a ball hits an obstacle
y_ref = self.obstacles[2].pos[1]
if y_ref + height / 2 > mob[i].pos[1] >= y_ref - height * (self.obstacle_rows - 1 / 2):
row = math.ceil((y_ref + height / 2 - mob[i].pos[1]) / height)
if x_ref - row * dist / 2 < mob[i].pos[0] <= x_ref + dist * ((self.obstacle_init_columns - 1)
+ row / 2):
column = math.ceil((mob[i].pos[0] - x_ref + row * dist / 2) / dist)
max_columns = self.obstacle_init_columns + row - 1
index = 1 + int(row * (self.obstacle_init_columns + max_columns) / 2) - max_columns \
+ column
ob_center = self.obstacles[index].pos
delta = squared_distance(ob_center, mob[i].pos)
threshold = self.p_radius + self.obstacle_radius
if delta <= threshold * threshold:
delta = math.sqrt(delta)
unit_collision = mul(add(ob_center, mul(mob[i].pos, -1)), 1 / delta)
mob[i].vector = add(mob[i].vector, mul(unit_collision, -(1 + mob[i].bounce_coeff) *
dot(mob[i].vector, unit_collision)))
mob[i].pos = add(mob[i].pos, mul(unit_collision, delta - self.p_radius -
self.obstacle_radius))
# Here we run steps_per_frame = 10 iterations per frame for ball-to-ball collision. This eliminates
# some problems due to large time steps, but in turn slows down the code. In particular, this solves
# the problem of stacking the balls on top of each other once they have landed. Lower values for
# steps_per_frame make ball stacking look badly realistic, so the tradeoff was worth it.
for j in range(self.steps_per_frame):
ptree = QuadTree(AABB(ORIGIN, FRAME_HEIGHT / 2, FRAME_HEIGHT / 2), 2) # Applying quad trees
for i in range(self.counter):
if mob[i].airborne: # Gravity is simulated here.
mob[i].vector[1] += self.gravity * dt / self.steps_per_frame
mob[i].pos = add(mob[i].pos, mul(mob[i].vector, dt / self.steps_per_frame))
mob[i].airborne = True
if mob[i].pos[1] <= ref_point[1] + 0.2: # Check if a ball hits the walls separating the slots
if not abs(mob[i].pos[0]) <= dist - ref_point[0] - self.p_radius:
index = 2 * (mob[i].pos[0] > 0) - 1
mob[i].vector[0] *= -mob[i].bounce_coeff
mob[i].pos[0] = index * (dist - ref_point[0] - self.p_radius)
if not mob[i].category == 0:
if mob[i].category >= 0:
x = x_barriers[min(last_row_columns - 1, math.floor(mob[i].category))] + 0.75 * \
self.obstacle_radius
if mob[i].pos[0] - self.p_radius <= x:
mob[i].vector[0] *= -mob[i].bounce_coeff
mob[i].pos[0] = x + self.p_radius
if mob[i].category <= last_row_columns - 1:
x = x_barriers[max(0, math.ceil(mob[i].category))] - 0.75 * self.obstacle_radius
if mob[i].pos[0] + self.p_radius >= x:
mob[i].vector[0] *= -mob[i].bounce_coeff
mob[i].pos[0] = x - self.p_radius
if mob[i].pos[1] - self.p_radius < -FRAME_HEIGHT / 2: # Check if a ball hits the floor
mob[i].vector[1] *= -mob[i].bounce_coeff / 3
mob[i].pos[1] = self.p_radius - FRAME_HEIGHT / 2
mob[i].airborne = False
ptree.insert(mob[i])
# Finally, handling ball-to-ball collision. This also conveniently handles ball stacking.
for i in range(self.counter):
neighbors = ptree.query(AABB(mob[i].pos, 2 * self.p_radius + 0.001, 2 * self.p_radius + 0.001))
for m in neighbors:
if mob[i].pos == m.pos:
continue
delta = squared_distance(mob[i].pos, m.pos)
if delta < 4 * self.p_radius * self.p_radius:
delta = math.sqrt(delta)
unit_collision = mul(add(m.pos, mul(mob[i].pos, -1)), 1 / delta)
error = self.p_radius - delta / 2 + 0.001
v_2 = dot(unit_collision, mob[i].vector)
v_1 = dot(unit_collision, m.vector)
mob[i].pos = add(mob[i].pos, mul(unit_collision, -error))
m.pos = add(m.pos, mul(unit_collision, error))
mob[i].vector = add(mob[i].vector, mul(unit_collision, v_1 - v_2))
m.vector = add(m.vector, mul(unit_collision, v_2 - v_1))
if not mob[i].category == 0:
mob[i].airborne = False
m.airborne = False
self.particles.add_updater(updater)