Learning objectives
In this code along, we will build a physics engine to simulate the force of gravity and how it affects the motions of a collection of heavenly bodies, solving the problem below whose solution we discussed in a language-neutral fashion in the core text. We will then run this simulation on the moons of Jupiter as well as a collection of three-body simulations that exhibit beautiful patterns.
Gravity Simulator Problem
Input: A Universe object initialUniverse, an integer numGens, and a floating-point number time.
Output: An array timePoints of numGens + 1 Universe objects, where timePoints[0] is initialUniverse, and timePoints[i] is obtained from timePoints[i-1] by approximating the results of the force of gravity over a time interval equal to time (in seconds).
Setup
This code along will build upon the starter code that we expanded in the previous code along. In particular, we will be working with the gravity folder, which has the following structure:
data: a directory containing.txtfiles that represent different initial systems that we will simulate.output: an initially empty directory that will contain images and .mp4 videos that we draw.datatypes.py: a file that contains type declarations.gravity.py: a file that will contain functions that we will write in this code along for implementing our physics engine.custom_io.py: a file that contains code for reading in a system of heavenly bodies from file.drawing.py: a file that contains code for drawing a system to file and creating animations over multiple time steps.main.py: a file that containsdef main(), where we will call functions to read in systems, run our gravity simulation, and draw the results.
We will be editing gravity.py in this code along to implement our physics engine; at the end of the code along, we will write some code in main.py, so make sure that main() does not have any code from the previous code along. Here is what our own main() contains at the start of this code along.
def main():
print("Buiding a gravity simulator.")
if __name__ == "__main__":
main()
Finally, after we implement the physics engine, we will generate videos visualizing the changes of our systems over multiple generations.
Code along summary
Reminder of object-oriented design
We discussed our data organization in the previous code along, but we provide it below as a refresher. We have omitted parameter checks and __repr()__ for clarity.
class Universe:
"""
Represents the entire simulation environment.
"""
gravitational_constant: float = 6.674e-11 # Default; can be overridden
def __init__(self, bodies: list[Body], width: float):
self.bodies = bodies
self.width = float(width)
class Body:
"""
Represents a celestial body within the simulation universe.
"""
def __init__(
self,
name: str,
mass: float,
radius: float,
position: OrderedPair,
velocity: OrderedPair,
acceleration: OrderedPair,
red: int,
green: int,
blue: int
):
self.name = name
self.mass = float(mass)
self.radius = float(radius)
self.position = position
self.velocity = velocity
self.acceleration = acceleration
self.red = red
self.green = green
self.blue = blue
class OrderedPair:
"""
Represents a 2D vector or coordinate pair (x, y).
Attributes:
x: The x-coordinate or horizontal component (float).
y: The y-coordinate or vertical component (float).
"""
def __init__(self, x: float = 0.0, y: float = 0.0):
self.x: float = float(x)
self.y: float = float(y)
Starting our highest level function
First, we will implement our highest level function, which will solve the Gravity Simulator Problem; its pseudocode representation from the core text is shown below.
SimulateGravity(initialUniverse, numGens, time)
timePoints ← array of numGens+1 Universe objects
timePoints[0] ← initialUniverse
for every integer i between 1 and numGens
timePoints[i] ← UpdateUniverse(timePoints[i-1], time)
return timePoints
This function is an easy one to write, since we are being lazy as always and passing the work to an update_universe() subroutine. However, we need to perform many parameter checks to ensure that the parameters given to simulate_gravity() are valid. For example, should we consult every Body object in initial_universe.bodies and check each of its attributes? Doing so is thorough and good practice but will really clutter our function.
def simulate_gravity(initial_universe: Universe, num_gens: int, time: float) -> list[Universe]:
"""
Simulate an N-body system for a fixed number of generations.
Args:
initial_universe: The starting state of the universe.
num_gens: Number of simulation steps to advance (>= 0).
time: Time step (Δt) between generations (> 0).
Returns:
A list of Universe snapshots of length num_gens + 1.
"""
# how should we run parameter checks?
time_points = [initial_universe]
# range from 1 to num_gens, and call update_universe to fill time_points[i]
for i in range(1, num_gens + 1):
updated = update_universe(time_points[i - 1], time)
time_points.append(updated)
return time_points
A short detour to write private validation functions
Notice that simulate_gravity() already has a lot going on. Rather than filling it with long chains of isinstance() checks, we will write a set of private helper functions, or internal implementation details not meant to be called from outside the module. Each validator’s name begins with an underscore to indicate that it is private. We start with three simple such functions: _validate_time_step(), _validate_num_gens(), and _validate_gravitational_constant(), each of which relies on a shared helper _is_finite_number() that checks whether an input value is a real, finite number (and not, for example, math.inf or math.nan).
def _validate_time_step(time: float) -> None:
if not _is_finite_number(time) or time <= 0:
raise ValueError("time_step must be a positive finite number")
def _validate_num_gens(num_gens: int) -> None:
if not isinstance(num_gens, int) or num_gens < 0:
raise ValueError("num_gens must be an integer >= 0")
def _validate_gravitational_constant(G: float) -> None:
if not _is_finite_number(G) or G <= 0:
raise ValueError("Universe.gravitational_constant must be a positive finite number")
def _is_finite_number(x: float) -> bool:
return isinstance(x, (int, float)) and math.isfinite(x)
We also need to validate the Universe object passed to simulate_gravity(). However, this task is not as simple as checking a single value: a Universe contains a width field and a list of Body objects, and each Body in turn has its own collection of attributes that need to be validated. Rather than pack all of this logic into one unwieldy function, we will as always delegate the work to subroutines.
def _validate_universe(u: Universe) -> None:
if not isinstance(u, Universe):
raise TypeError("Expected a Universe instance")
if not _is_finite_number(u.width) or u.width <= 0:
raise ValueError("Universe.width must be a positive finite number")
if not isinstance(u.bodies, list):
raise TypeError("Universe.bodies must be a list of Body instances")
for b in u.bodies:
_validate_body(b)
The _validate_body() function checks each attribute of a Body, and in turn calls _validate_pair() to check each of the three OrderedPair attributes.
def _validate_body(b: Body) -> None:
if not isinstance(b, Body):
raise TypeError("Expected a Body instance")
if not _is_finite_number(b.mass) or b.mass <= 0:
raise ValueError("Body.mass must be a positive finite number")
if not _is_finite_number(b.radius) or b.radius < 0:
raise ValueError("Body.radius must be a nonnegative finite number")
_validate_pair(b.position)
_validate_pair(b.velocity)
_validate_pair(b.acceleration)
Finally, _validate_pair() is the most foundational of our validators: it checks that a value is an OrderedPair containing two finite numbers. Because it is called by _validate_body(), which is called by _validate_universe(), reusing it anywhere else in our code costs us nothing as the subroutine is already written and tested.
def _validate_pair(p: OrderedPair) -> None:
if not isinstance(p, OrderedPair):
raise TypeError("Expected an OrderedPair")
if not _is_finite_number(p.x) or not _is_finite_number(p.y):
raise ValueError("OrderedPair must contain finite numeric components")
With our private validators in place, we can now rewrite simulate_gravity() so that all parameter checking is delegated cleanly to four function calls. Notice how little the body of the function has changed, since we have only added four lines at the top of the function, and yet we have gained thorough validation of every field of every Body object belonging to the input Universe.
def simulate_gravity(initial_universe: Universe, num_gens: int, time: float) -> list[Universe]:
"""
Simulate an N-body system for a fixed number of generations.
Args:
initial_universe: The starting state of the universe.
num_gens: Number of simulation steps to advance (>= 0).
time: Time step (Δt) between generations (> 0).
Returns:
A list of Universe snapshots of length num_gens + 1.
"""
_validate_universe(initial_universe)
_validate_num_gens(num_gens)
_validate_time_step(time)
_validate_gravitational_constant(Universe.gravitational_constant)
time_points = [initial_universe]
# range from 1 to num_gens, and call update_universe to fill time_points[i]
for i in range(1, num_gens + 1):
next_universe = update_universe(time_points[i - 1], time)
time_points.append(next_universe)
return time_points
Updating a universe
Our update_universe() function advances the state of every Body in the universe by one time step. It validates its inputs using the private helpers we just wrote, creates a deep copy of the current universe (we will explain why in a moment), and then updates each body’s acceleration, velocity, and position in that copy, delegating each computation to a subroutine.
def update_universe(current_universe: Universe, time: float) -> Universe:
"""
Advance the universe by a single time step.
Uses a velocity-Verlet style update: compute new accelerations from the
current state, then advance velocity and position accordingly.
Args:
current_universe: Universe state at the current time.
time: Time step (Δt) to advance.
Returns:
A new Universe instance representing the next state.
"""
_validate_universe(current_universe)
_validate_time_step(time)
new_universe = copy_universe(current_universe)
# Update every body in the cloned universe based on forces from current_universe
for b in new_universe.bodies:
old_acc, old_vel = b.acceleration, b.velocity
b.acceleration = update_acceleration(current_universe, b)
b.velocity = update_velocity(b, old_acc, time)
b.position = update_position(b, old_acc, old_vel, time)
return new_universe
Copying a universe
You might wonder why we cannot simply write new_universe = current_universe to get a copy. The answer is that Python objects are passed by reference: that assignment does not create a new Universe — it creates a second name that points to the same object in memory. Any change we make to a Body inside new_universe would silently modify current_universe as well, corrupting the simulation history we are building up in time_points. To avoid this, we need a deep copy: a completely independent Universe object whose Body objects and OrderedPair objects are all freshly allocated.
We now write copy_universe() and its helper copy_body(). Note how copy_universe() validates the universe once and then delegates each body’s copy to copy_body(). Even though there are many attributes to copy, the validation cost is minimal: a single call to _validate_universe() already checks every field of every body in the list.
def copy_universe(current_universe: Universe) -> Universe:
"""
Deep-copy a Universe (bodies and width).
The gravitational constant `G` is a class attribute and does not need to be copied.
Args:
current_universe (Universe): The universe to copy. Must contain
a list of bodies and a width value.
Returns:
Universe: A new Universe instance with deep-copied bodies and
the same width as the original.
"""
_validate_universe(current_universe)
new_bodies = []
for b in current_universe.bodies:
new_bodies.append(copy_body(b))
return Universe(new_bodies, current_universe.width)
def copy_body(b: Body) -> Body:
"""
Deep-copy a Body, including position, velocity, and acceleration.
Args:
b (Body): The body to copy. Must contain name, mass, radius,
position, velocity, acceleration, and color attributes.
Returns:
Body: A new Body instance with identical properties and
deep-copied OrderedPair objects for position, velocity,
and acceleration.
"""
_validate_body(b)
new_body = Body(
b.name, b.mass, b.radius,
OrderedPair(b.position.x, b.position.y),
OrderedPair(b.velocity.x, b.velocity.y),
OrderedPair(b.acceleration.x, b.acceleration.y),
b.red, b.green, b.blue,
)
return new_body
Using Newton’s laws to compute the net force (acceleration) on a body
update_acceleration() computes the new acceleration of a body by first finding the net gravitational force acting on it, then applying Newton’s second law: a = F / m. It delegates the force calculation to compute_net_force(), which we will write next.
def update_acceleration(current_universe: Universe, b: Body) -> OrderedPair:
"""
Compute the body's acceleration from the net gravitational force.
Uses:
a = F / m
Args:
current_universe (Universe): The universe containing all bodies that
contribute gravitational force.
b (Body): The body whose acceleration is being computed.
Returns:
OrderedPair: A 2D vector (ax, ay) representing the updated acceleration.
"""
_validate_universe(current_universe)
_validate_body(b)
net_force = compute_net_force(current_universe, b) # OrderedPair Fx, Fy
# apply Newton's second law over components
ax = net_force.x / b.mass
ay = net_force.y / b.mass
return OrderedPair(ax, ay)
The net gravitational force on a body is the vector sum of the individual forces exerted on it by every other body in the universe. compute_net_force() loops over all bodies in the universe, skips the body itself (a body does not exert a gravitational force on itself), and accumulates the x- and y-components of each pairwise force, delegating each pairwise calculation to compute_force().
def compute_net_force(current_universe: Universe, b: Body) -> OrderedPair:
"""
Compute the net gravitational force on a body from all other bodies.
Args:
current_universe (Universe): The universe containing all bodies.
Must have a list of bodies and a valid gravitational constant.
b (Body): The body on which the net gravitational force is computed.
Returns:
OrderedPair: A 2D vector (x, y) representing the net gravitational
force acting on the given body.
"""
_validate_universe(current_universe)
_validate_body(b)
net_force = OrderedPair(0.0, 0.0)
G = Universe.gravitational_constant
for cur_body in current_universe.bodies:
if not (cur_body is b):
current_force = compute_force(b, cur_body, G)
net_force.x += current_force.x
net_force.y += current_force.y
return net_force
The pairwise gravitational force between two bodies is given by Newton’s law of universal gravitation: F = G · m1 · m2 / d2, directed from b1 toward b2. compute_force() computes this magnitude and then decomposes it into x and y components by multiplying by the unit vector pointing from b1 to b2. Note that we handle the edge case of two bodies occupying exactly the same position by returning a zero force, which avoids a division-by-zero error.
def compute_force(b1: Body, b2: Body, G: float) -> OrderedPair:
"""
Compute the gravitational force exerted on one body by another.
Args:
b1 (Body): The body on which the force is acting.
b2 (Body): The body exerting the gravitational force.
G (float): Gravitational constant.
Returns:
OrderedPair: A 2D vector (x, y) representing the force exerted
on `b1` by `b2`.
"""
_validate_body(b1)
_validate_body(b2)
_validate_gravitational_constant(G)
d = distance(b1.position, b2.position)
if d == 0.0:
return OrderedPair(0.0, 0.0) # treat as no force
F_magnitude = G * b1.mass * b2.mass / (d * d)
dx = b2.position.x - b1.position.x
dy = b2.position.y - b1.position.y
Fx = F_magnitude * dx / d
Fy = F_magnitude * dy / d
return OrderedPair(Fx, Fy)
The computation of force requires knowing the distance between two bodies. We compute this using the standard Euclidean distance formula, which is its own small but reusable subroutine.
def distance(p1: OrderedPair, p2: OrderedPair) -> float:
"""
Compute the Euclidean distance between two position vectors.
Args:
p1 (OrderedPair): The first position vector.
p2 (OrderedPair): The second position vector.
Returns:
float: The distance between p1 and p2.
"""
_validate_pair(p1)
_validate_pair(p2)
dx = p2.x - p1.x
dy = p2.y - p1.y
return math.sqrt(dx * dx + dy * dy)
Updating the velocity and position of a body
We update the velocity of a body using the velocity-Verlet method, which averages the old and new accelerations over the time step. This gives a more accurate result than simply using either the old or the new acceleration alone, and is a standard technique in physics simulations. Note that by the time update_velocity() is called from update_universe(), b.acceleration has already been updated to the new value, which is why we also pass the old acceleration as a separate argument.
def update_velocity(b: Body, old_acceleration: OrderedPair, time: float) -> OrderedPair:
"""
Update velocity using average acceleration over the step.
Formula:
v_{t+Δt} = v_t + 0.5 * (a_t + a_{t+Δt}) * Δt
Args:
b (Body): The body whose velocity is being updated. Must have
a `velocity` attribute (OrderedPair) and a current
`acceleration` attribute (OrderedPair).
old_acceleration (OrderedPair): The acceleration of the body at
the previous time step.
time (float): The time step Δt over which to update the velocity.
Must be a positive value.
Returns:
OrderedPair: A new OrderedPair containing the updated velocity
components (vx, vy).
"""
_validate_body(b)
_validate_pair(old_acceleration)
_validate_time_step(time)
vx = b.velocity.x + 0.5 * (b.acceleration.x + old_acceleration.x) * time
vy = b.velocity.y + 0.5 * (b.acceleration.y + old_acceleration.y) * time
return OrderedPair(vx, vy)
We update position using the standard kinematic formula, which accounts for both the initial velocity and the acceleration over the time interval. We use the old velocity and old acceleration here — both captured before any updates are applied in update_universe() — which is why update_universe() saves them explicitly before calling update_acceleration().
def update_position(b: Body, old_acc: OrderedPair, old_vel: OrderedPair, time: float) -> OrderedPair:
"""
Update position using constant-acceleration kinematics.
Formula:
p_{t+Δt} = p_t + v_t * Δt + 0.5 * a_t * Δt²
Args:
b (Body): The body whose position is being updated. Must have
a `position` attribute (OrderedPair).
old_acc (OrderedPair): The acceleration of the body at the previous
time step.
old_vel (OrderedPair): The velocity of the body at the previous
time step.
time (float): The time step Δt over which to update the position.
Must be a positive value.
Returns:
OrderedPair: A new OrderedPair containing the updated position
components (px, py).
"""
_validate_body(b)
_validate_pair(old_acc)
_validate_pair(old_vel)
_validate_time_step(time)
px = b.position.x + old_vel.x * time + 0.5 * old_acc.x * time * time
py = b.position.y + old_vel.y * time + 0.5 * old_acc.y * time * time
return OrderedPair(px, py)
Running the simulation
Now that we have implemented our physics engine, we can write main() to tie everything together. Our program will read five command-line arguments: the scenario name (matching a .txt file in the data/ directory), the number of simulation generations, the time step in seconds, the canvas width in pixels, and — as we will discuss momentarily — a drawing frequency. Here is the complete main().
def main():
"""
Run the full pipeline:
1) read universe from file
2) simulate gravity for N generations
3) render selected frames to pygame surfaces
4) encode frames to an MP4 video
"""
print("Let's simulate gravity!")
# Expect 5 user arguments (plus program name)
if len(sys.argv) != 6:
raise ValueError(
"Error: incorrect number of command line arguments. Six desired."
)
scenario = sys.argv[1]
# establish input file and output prefix
input_file = f"data/{scenario}.txt"
video_path = f"output/{scenario}.mp4"
# Parse CLI arguments
num_gens = int(sys.argv[2])
time_step = float(sys.argv[3])
canvas_width = int(sys.argv[4])
print("Command line arguments read!")
# Read initial universe and gravitational constant
initial_universe = read_universe(input_file)
# --- Simulate ---
print("Simulating gravity now.")
time_points = simulate_gravity(initial_universe, num_gens, time_step)
print("Gravity simulation complete.")
# --- Draw frames ---
print("Rendering frames.")
surfaces = animate_system(time_points, canvas_width)
print("Frames drawn.")
# --- Encode video ---
print("Encoding MP4 video.")
save_video_from_surfaces(surfaces, video_path, fps=10, codec="libx264", quality=8)
print("Success! MP4 video produced.")
print("Animation finished! Exiting normally.")
Let’s start by running the simulation on the Jupiter moons system. The command below simulates 1,000 generations with a time step of 60 seconds (one minute per step) and a canvas 1,500 pixels wide.
python3 main.py jupiterMoons 1000 60 1500
After running this command, you may notice that the moons barely seem to have moved. This makes sense: 1,000 steps of 60 seconds each represents only about 17 hours of simulated time, while the innermost moon Io takes roughly 42 hours to complete a single orbit. To see meaningful orbital motion, we would need either many more generations or a larger time step — but rendering every one of those frames into a video would be slow and produce a very large file. We will solve this next with a drawing frequency parameter.
Adding a drawing frequency
Rather than rendering every simulation step as a video frame, we add a drawing_frequency parameter: we only produce a surface for every drawing_frequency-th time step. This allows us to run many more total generations while keeping the output video a manageable length. We update main() to read this new argument (highlighted) and pass it through to animate_system().
def main():
"""
Run the full pipeline:
1) read universe from file
2) simulate gravity for N generations
3) render selected frames to pygame surfaces
4) encode frames to an MP4 video
"""
print("Let's simulate gravity!")
# Expect 5 user arguments (plus program name)
if len(sys.argv) != 6:
raise ValueError(
"Error: incorrect number of command line arguments. Six desired."
)
scenario = sys.argv[1]
# establish input file and output prefix
input_file = f"data/{scenario}.txt"
video_path = f"output/{scenario}.mp4"
# Parse CLI arguments
num_gens = int(sys.argv[2])
time_step = float(sys.argv[3])
canvas_width = int(sys.argv[4])
drawing_frequency = int(sys.argv[5])
print("Command line arguments read!")
# Read initial universe and gravitational constant
initial_universe = read_universe(input_file)
# --- Simulate ---
print("Simulating gravity now.")
time_points = simulate_gravity(initial_universe, num_gens, time_step)
print("Gravity simulation complete.")
# --- Draw frames ---
print("Rendering frames.")
surfaces = animate_system(time_points, canvas_width, drawing_frequency)
print("Frames drawn.")
# --- Encode video ---
print("Encoding MP4 video.")
save_video_from_surfaces(surfaces, video_path, fps=10, codec="libx264", quality=8)
print("Success! MP4 video produced.")
print("Animation finished! Exiting normally.")
We also update animate_system() in drawing.py to accept and use the drawing_frequency parameter, so that it only renders a frame when the current index is a multiple of drawing_frequency.
def animate_system(
time_points: list[Universe],
canvas_width: int,
drawing_frequency: int
) -> list[pygame.Surface]:
"""
Render snapshots of the Universe into pygame surfaces.
Delegates the drawing to draw_to_canvas(u, canvas_width).
Only draws frames when i % drawing_frequency == 0.
"""
if not isinstance(time_points, list):
raise TypeError("time_points must be a list of Universe objects")
if len(time_points) == 0:
raise ValueError("time_points must not be empty")
for u in time_points:
if not isinstance(u, Universe):
raise TypeError("all elements of time_points must be Universe objects")
if not isinstance(canvas_width, int) or canvas_width <= 0:
raise ValueError("canvas_width must be a positive integer")
if not isinstance(drawing_frequency, int) or drawing_frequency <= 0:
raise ValueError("drawing_frequency must be a positive integer")
surfaces = []
for i, u in enumerate(time_points):
if i % drawing_frequency == 0:
surface = draw_to_canvas(u, canvas_width)
surfaces.append(surface)
return surfaces
(Optional) Improving the visuals by adding trails
Much like how planes leave contrails across the sky, we will add trails to our visualization to show where each body has been. Trails let us see the path each body has traced through space over time, which reveals the beautiful structure of orbital motion far more vividly than a static snapshot.
We will represent trails as a dictionary mapping each body’s index in u.bodies to a list of recent OrderedPair positions. As we step through the simulation, we append each body’s current position to its trail list, and we cap the list at a fixed maximum length so that old positions gradually fall off the back of the trail, giving the impression of fading.
We add four constants to the top of drawing.py, after the package imports, alongside the existing JUPITER_MOON_MULTIPLIER. TRAIL_FREQUENCY controls how often a position is recorded in the trail history relative to the drawing frequency; NUMBER_OF_TRAIL_FRAMES limits how many past positions are kept; and TRAIL_THICKNESS_FACTOR scales the trail line width relative to the body’s on-canvas radius.
# Trail rendering parameters TRAIL_FREQUENCY = 10 NUMBER_OF_TRAIL_FRAMES = 100 TRAIL_THICKNESS_FACTOR = 0.2
We first update animate_system() to build and maintain the trails dictionary as we step through the time points, and to pass it to draw_to_canvas() on each draw frame.
def animate_system(
time_points: list[Universe],
canvas_width: int,
drawing_frequency: int
) -> list[pygame.Surface]:
"""
Render selected frames of the system into pygame surfaces.
Frames are sampled every `drawing_frequency` simulation steps; trail history
is updated more frequently using `TRAIL_FREQUENCY` so trails look smooth.
Args:
time_points: Snapshots of the Universe over time (0..N).
canvas_width: Width/height (px) of the square canvas.
drawing_frequency: Draw a frame when i % drawing_frequency == 0.
Returns:
A list of pygame.Surface objects (one per drawn frame).
"""
if not isinstance(time_points, list):
raise TypeError("time_points must be a list of Universe objects")
if len(time_points) == 0:
raise ValueError("time_points must not be empty")
for u in time_points:
if not isinstance(u, Universe):
raise TypeError("all elements of time_points must be Universe objects")
if not isinstance(canvas_width, int) or canvas_width <= 0:
raise ValueError("canvas_width must be a positive integer")
if not isinstance(drawing_frequency, int) or drawing_frequency <= 0:
raise ValueError("drawing_frequency must be a positive integer")
surfaces = []
trails = {}
for i, u in enumerate(time_points):
# Update trails at the configured frequency
if (i * TRAIL_FREQUENCY) % drawing_frequency == 0:
for body_index in range(len(u.bodies)):
if body_index not in trails:
trails[body_index] = []
trails[body_index].append(u.bodies[body_index].position)
# Keep only the most recent max_len entries
max_len = NUMBER_OF_TRAIL_FRAMES * TRAIL_FREQUENCY
if len(trails[body_index]) > max_len:
start_index = len(trails[body_index]) - max_len
else:
start_index = 0
trails[body_index] = trails[body_index][start_index:]
# Draw a frame
if i % drawing_frequency == 0:
surface = draw_to_canvas(u, canvas_width, trails)
surfaces.append(surface)
return surfaces
We next update draw_to_canvas() to accept the trails dictionary as a new parameter, and to call a new draw_trails() subroutine before drawing the bodies themselves (highlighted), so that body circles appear on top of their trails.
def draw_to_canvas(
u: Universe,
canvas_width: int,
trails: dict[int, list[OrderedPair]],
) -> pygame.Surface:
"""
Draw a single Universe snapshot onto a new pygame Surface.
"""
# --- lightweight parameter checks ---
if not isinstance(u, Universe):
raise TypeError("u must be a Universe")
if not isinstance(canvas_width, int) or canvas_width <= 0:
raise ValueError("canvas_width must be a positive integer")
if not isinstance(trails, dict):
raise TypeError("trails must be a dict[int, list[OrderedPair]]")
# --- create drawing surface ---
surface = pygame.Surface((canvas_width, canvas_width))
surface.fill((255, 255, 255)) # white background
# Trails first (so bodies appear on top)
draw_trails(surface, trails, u.width, canvas_width, u.bodies)
# Draw bodies
for b in u.bodies:
color = (b.red, b.green, b.blue)
center_x = int((b.position.x / u.width) * canvas_width)
center_y = int((b.position.y / u.width) * canvas_width)
radius = (b.radius / u.width) * canvas_width
if b.name in ("Io", "Ganymede", "Callisto", "Europa"):
radius *= JUPITER_MOON_MULTIPLIER
pygame.draw.circle(surface, color, (center_x, center_y), radius)
return surface
Finally, we write draw_trails(). For each body, we draw a sequence of short line segments connecting consecutive recorded trail positions. To give the impression of a fading tail, we blend each segment’s color from white (the background) toward the body’s true color. The parameter alpha increases from 0 at the oldest trail point to 255 at the newest: near the head of the trail (small j), alpha is close to 0 and the color is close to white; near the body itself (large j), alpha is close to 255 and the color approaches the body’s true RGB values. Each channel is computed as a weighted blend between 255 (white) and the body’s color, with alpha / 255 as the weight on the body’s color.
def draw_trails(
surface: pygame.Surface,
trails: dict[int, list[OrderedPair]],
u_width: float,
canvas_width: int,
bodies: list[Body],
) -> None:
"""
Draw fading trails for each body onto the given surface.
"""
for body_index, b in enumerate(bodies):
trail = trails.get(body_index, [])
num_trails = len(trail)
# Line width scales with body radius
line_width = int((b.radius / u_width) * canvas_width * TRAIL_THICKNESS_FACTOR)
# Thicker lines for the big Jupiter moons
if b.name in {"Ganymede", "Io", "Callisto", "Europa"}:
line_width = int(line_width * JUPITER_MOON_MULTIPLIER)
if line_width <= 0:
line_width = 1
# Draw segments with a simple color fade toward the body's color
for j in range(num_trails - 1):
alpha = 255.0 * j / num_trails
red = int((1 - alpha / 255.0) * 255.0 + (alpha / 255.0) * b.red)
green = int((1 - alpha / 255.0) * 255.0 + (alpha / 255.0) * b.green)
blue = int((1 - alpha / 255.0) * 255.0 + (alpha / 255.0) * b.blue)
start_x = int((trail[j].x / u_width) * canvas_width)
start_y = int((trail[j].y / u_width) * canvas_width)
end_x = int((trail[j + 1].x / u_width) * canvas_width)
end_y = int((trail[j + 1].y / u_width) * canvas_width)
pygame.draw.line(
surface, (red, green, blue),
(start_x, start_y), (end_x, end_y), line_width
)
Rerunning our simulation
Let’s rerun the simulation on the Jupiter moons system, this time passing a drawing frequency of 10 so that we capture one frame for every 10 simulation steps. With 1,000 generations at 60 seconds each, this yields 100 frames showing roughly 17 hours of orbital motion.
python3 main.py jupiterMoons 1000 60 1500 10
Once you have verified that your code is working, try modifying Universe.gravitational_constant at the top of gravity.py — double it, halve it, or experiment with other values. Notice how the orbital periods and trajectories change: a stronger gravitational constant pulls the moons into tighter, faster orbits, while a weaker one causes them to drift more slowly.
Three-body simulations
Now for the main event. Our data/ directory includes several classic three-body initial conditions — beautifully choreographed orbits discovered numerically. Unlike the Jupiter system, these simulations use a gravitational constant of 1 and much smaller time steps, since the bodies are much lighter and closer together relative to their masses. Three-body orbits are also famously unstable: small numerical errors accumulate over time and the bodies eventually diverge from their ideal trajectories, so we need many more generations at a finer time step to capture the full pattern before chaos sets in. The commands below each take about 30 seconds to run.
python3 main.py butterfly 100000 0.0001 1500 300 python3 main.py dragonfly 100000 0.0001 1500 300 python3 main.py figureEight 100000 0.0001 1500 300 python3 main.py goggles 100000 0.0001 1500 300 python3 main.py moth 100000 0.0001 1500 300 python3 main.py yarn 100000 0.0001 1500 300
For even higher resolution and smoother trails, we can simulate ten times as many generations with a proportionally smaller time step. These runs will take longer but produce the most visually striking results.
python3 main.py butterfly 1000000 0.00001 1500 3000 python3 main.py dragonfly 1000000 0.00001 1500 3000
After running these simulations, you will find the following videos in your output/ directory.
The dragonfly three-body orbit:
The butterfly three-body orbit:
Time for practice!
After you have the chance to check your work below, you are ready to apply what you have learned to the chapter’s practice problems. Or, you can carry on to the next chapter, where we will learn about how to construct evolutionary trees from DNA data.
