Learning objectives
Recall the problem that we are solving
Gravity Simulator Problem
Input: AUniverse
objectinitialUniverse
, an integernumGens
, and a floating-point numbertime
.
Output: An arraytimePoints
ofnumGens + 1
Universe objects, wheretimePoints[0]
isinitialUniverse
, andtimePoints[i]
is obtained fromtimePoints[i-1]
by approximating the results of the force of gravity over a time interval equal totime
(in seconds).
Setup
[Starter code is the same as that from the previous code along, where we added drawing code.]
Code along summary
[Pseudocode?]
Recall the problem that we are solving
Gravity Simulator Problem
Input: AUniverse
objectinitialUniverse
, an integernumGens
, and a floating-point numbertime
.
Output: An arraytimePoints
ofnumGens + 1
Universe objects, wheretimePoints[0]
isinitialUniverse
, andtimePoints[i]
is obtained fromtimePoints[i-1]
by approximating the results of the force of gravity over a time interval equal totime
(in seconds).
Organizing data
We reproduce our data organization below.
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): if not isinstance(x, (int, float)): raise TypeError(f"x must be a number (int or float), got {type(x).__name__}") if not isinstance(y, (int, float)): raise TypeError(f"y must be a number (int or float), got {type(y).__name__}") self.x: float = float(x) self.y: float = float(y) 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 ): # Name if not isinstance(name, str) or not name.strip(): raise ValueError("Body name must be a non-empty string.") # Physical properties if not isinstance(mass, (int, float)) or mass <= 0: raise ValueError(f"Mass must be a positive number, got {mass}") if not isinstance(radius, (int, float)) or radius < 0: raise ValueError(f"Radius must be a non-negative number, got {radius}") # Vectors for arg, label in ((position, "position"), (velocity, "velocity"), (acceleration, "acceleration")): if not isinstance(arg, OrderedPair): raise TypeError(f"{label} must be an OrderedPair, got {type(arg).__name__}") # Color for c, channel in ((red, "red"), (green, "green"), (blue, "blue")): if not isinstance(c, int): raise TypeError(f"Color component {channel} must be int, got {type(c).__name__}") if not (0 <= c <= 255): raise ValueError(f"Color component {channel} must be in range [0, 255], got {c}") 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 Universe: """ Represents the entire simulation environment. """ gravitational_constant: float = 6.674e-11 # Default; can be overridden def __init__(self, bodies: list[Body], width: float): if not isinstance(bodies, list) or not all(isinstance(b, Body) for b in bodies): raise TypeError("bodies must be a list of Body objects.") if not isinstance(width, (int, float)) or width <= 0: raise ValueError(f"width must be a positive number, got {width}") self.bodies = bodies self.width = float(width)
Starting our highest level function
First, highest level function for solving the Gravity Simulator Problem. It’s an easy function to write, but we have an enormous number of things that we have to keep in mind regarding the parameters to our system, and adding them after the function’s docstring will start to 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
Validation functions
[Transition to validation functions. Explain that they are private. Let’s write two simple ones, for validating num_gens, time, and the gravitational constant associated with Universe. Helper function provided by validate_gravitational_constant and _is_finite_number().]
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 should also validate the universe, but the issue with this is that not only do we have to validate all the fields of the universe, but we need to validate each body in the universe too. Subroutines for the win once again!]
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)
[Now we write _validate_body().]
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)
[This relies on yet another validation function to validate an ordered pair. All of these subroutines will prove useful because they can be reused.]
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")
[Now that we have written these validation functions, we can reuse them however we like, and we can also plug them into simulate_gravity() with minimal intrusion into this 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. """ _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
[update_universe() allows us to practice what we have just learned with validation functions.]
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
[Now we write subroutines, starting with copy_universe and its own helper function copy_body. Note that even though we are validating many parameters, they are all handled with one _validate_universe() function call.]
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
update_acceleration
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)
compute_net_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
compute_force
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)
[This in turn relies on a distance 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)
update_velocity
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)
update_position
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
Here is our updated def 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.")
[Here is a collection of parameters to run these things for.]
[Run simulation. Nothing is really moving?]
python3 main.py jupiterMoons 1000 60 1500
Adding a drawing frequency
[take in one additional parameter that will indicate how frequently to draw the images]
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.")
[updating 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
Improving the visuals by adding trails
[Much like how planes have chemtrails spreading government chemicals across our free skies, we will add trails to our visualization to see where they have been.]
[trails are dictionaries mapping integers to OrderedPair objects. The integers represent the index of the body, and the OrderedPair represent positions of the Body.]
[add three constants alongside JUPITER_MOON_THICKNESS]
# Trail rendering parameters TRAIL_FREQUENCY = 10 NUMBER_OF_TRAIL_FRAMES = 100 TRAIL_THICKNESS_FACTOR = 0.2
[First, we update animate_system]
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
[Next, we update draw_to_canvas().]
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]
[alpha is a way to fade the trail color gradually from white (background) to the body’s color]
[j / num_trails gives you a fraction along the trail: near the head (small j), it’s close to 0; near the tail (large j), it’s close to 1.
That fraction is scaled to 0–255 as alpha.
Each color channel is then a blend:
At the start, more weight on 255 (white).
At the end, more weight on the body’s (r, g, b).
So the trail starts out pale and fades into the body’s true 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
[Visualize simulation. Code runs below]
python3 main.py jupiterMoons 1000 60 1500 10
[See what happens when we double the gravitational constant. What happens if we lower it?]
Three body simulations
[Now for the main event: three body simulations]
[We need to increase the number of generations by a factor of 10x (And decrease the time and increase drawing_Frequency too) on some of these (butterfly, dragonfly, ) … make the point thatt these are unstable. below take about 30 secs each.]
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
[updated runs]
python3 main.py butterfly 1000000 0.00001 1500 3000 python3 main.py dragonfly 1000000 0.00001 1500 3000
Check your work!
[Insert link to cogniterra here]
Time for practice!
[We’ve written code to simulate gravity, but we can build a similar physics engine for different things.]