Manim code to simulate 3-body systems

gistfile1.txt

from manim import *
import numpy as np


class Body:
    def __init__(self, mass, position, color):
        self.mass = mass
        self.position = np.array([position[0], position[1], 0.0])
        self.velocity = np.array([0.0, 0.0, 0.0])
        self.color = color
        self.acceleration = np.array([0.0, 0.0, 0.0])
        self.trajectory = [self.position.copy()]

    def update_velocity(self, velocity):
        self.velocity = np.array([velocity[0], velocity[1], 0.0])

    def update_position(self, dt):
        self.velocity += self.acceleration * dt
        self.position += self.velocity * dt
        self.trajectory.append(self.position.copy())
        self.acceleration = np.array([0.0, 0.0, 0.0])


class ThreeBodySimulation(Scene):
    def construct(self):
        # Constants
        self.G = 100.0  # Gravitational constant
        self.dt = 0.001  # Time step
        self.vX = 10.0  # Velocity multiplier
        self.dv = 0.1  # Small number outside the computer's precision (eg: 2^(-150) in a 32-bit computer)

        rad_length = 2.0

        # Equilateral triangle vertices centered at origin
        self.bodies = [
            Body(100.0, [rad_length * np.cos(0), rad_length * np.sin(0)], BLUE),
            Body(
                100.0,
                [
                    rad_length * np.cos(2 * np.pi / 3),
                    rad_length * np.sin(2 * np.pi / 3),
                ],
                RED,
            ),
            Body(
                100.0,
                [
                    rad_length * np.cos(4 * np.pi / 3),
                    rad_length * np.sin(4 * np.pi / 3),
                ],
                GREEN,
            ),
        ]

        # Initial velocities ~orthogonal to (center-of-mass of system to body) vector
        center_of_mass_position = np.average(
            [body.position for body in self.bodies], axis=0
        )  # Origin, ik. Redundancy if I change initial positions
        for body in self.bodies:
            orth = body.position - center_of_mass_position
            # velocity direction
            velocity = np.array([orth[1], -1 * orth[0]])
            # Add some margin for velocities randomly
            velocity += np.random.choice([-1, 1], size=2) * self.dv
            body.initialize_velocity(velocity * self.vX)

        # Create Manim objects
        self.dots = [Dot(point=body.position, color=body.color) for body in self.bodies]
        self.trails = [
            VMobject(color=body.color, stroke_width=2) for body in self.bodies
        ]

        # Add objects to scene
        for dot, trail in zip(self.dots, self.trails):
            self.add(dot, trail)

        def calculate_forces():
            for i, body1 in enumerate(self.bodies):
                for j, body2 in enumerate(self.bodies):
                    if i != j:
                        r = body2.position - body1.position
                        r_mag = np.linalg.norm(r)
                        r_hat = r / r_mag
                        force = self.G * body1.mass * body2.mass * r_hat / (r_mag**2)
                        body1.acceleration += force / body1.mass

        def update_bodies(mobj, dt):
            # Calculate forces and update positions
            calculate_forces()
            for body in self.bodies:
                body.update_position(dt)

            # Update visual elements
            for dot, body, trail in zip(self.dots, self.bodies, self.trails):
                dot.move_to(body.position)

                # Update trail with recent points
                points = body.trajectory[-100:]
                if len(points) > 1:
                    trail.set_points_smoothly(points)

        # Create an always_redraw updater for continuous animation
        self.always_update = always_redraw(lambda: VGroup(*self.dots, *self.trails))

        # Add an updater that calls our update function
        self.add(self.always_update)

        # Create an animation that updates the simulation
        def simulation_updater(mob, dt):
            update_bodies(mob, self.dt)

        self.always_update.add_updater(simulation_updater)

        # Run the animation
        self.wait(15)  # Run in seconds

        # Clean up
        self.always_update.clear_updaters()
        self.wait(1)


# To render the animation, run:
# manim -pql three_body_simulation.py ThreeBodySimulation