in physics and astronomy a gravitational nbody simulation is a simulation of a dynamical system of particles under the influence of gravity the system consists of a number of bodies each of which exerts a gravitational force on all other bodies these forces are calculated using newton s law of universal gravitation the euler method is used at each timestep to calculate the change in velocity and position brought about by these forces softening is used to prevent numerical divergences when a particle comes too close to another and the force goes to infinity description adapted from https en wikipedia orgwikinbodysimulation see also http www shodor orgrefdeskresourcesalgorithmseulersmethod frame rate of the animation time between time steps in seconds the parameters size color are not relevant for the simulation itself they are only used for plotting euler algorithm for velocity body1 body0 0 0 0 body1 updatevelocity1 0 1 body1 velocity 1 0 0 0 body1 updatevelocity1 0 1 body1 velocity 2 0 0 0 body2 body0 0 5 0 body2 updatevelocity0 10 10 body2 velocity 5 0 100 0 body2 updatevelocity0 10 10 body2 velocity 5 0 200 0 euler algorithm for position body1 body0 0 1 0 body1 updateposition1 body1 position 1 0 0 0 body1 updateposition1 body1 position 2 0 0 0 body2 body10 10 0 2 body2 updateposition1 body2 position 10 0 8 0 body2 updateposition1 body2 position 10 0 6 0 this class is used to hold the bodies the gravitation constant the time factor and the softening factor the time factor is used to control the speed of the simulation the softening factor is used for softening a numerical trick for nbody simulations to prevent numerical divergences when two bodies get too close to each other for each body loop through all other bodies to calculate the total force they exert on it use that force to update the body s velocity bodysystem1 bodysystembody0 0 0 0 body10 0 0 0 lenbodysystem1 2 bodysystem1 updatesystem1 bodysystem1 bodies0 position 0 01 0 0 bodysystem1 bodies0 velocity 0 01 0 0 bodysystem2 bodysystembody10 0 0 0 body10 0 0 0 mass4 1 10 bodysystem2 updatesystem1 bodysystem2 bodies0 position 9 0 0 0 bodysystem2 bodies0 velocity 0 1 0 0 calculation of the distance using pythagoras s theorem extra factor due to the softening technique newton s law of universal gravitation update the body s velocity once all the force components have been added update the positions only after all the velocities have been updated updates the bodysystem and applies the change to the patchlist used for plotting bodysystem1 bodysystembody0 0 0 0 body10 0 0 0 patches1 plt circlebody positionx body positiony body size fcbody colorfor body in bodysystem1 bodies doctest ellipsis updatestepbodysystem1 1 patches1 patches10 center 0 01 0 0 bodysystem2 bodysystembody10 0 0 0 body10 0 0 0 mass4 1 10 patches2 plt circlebody positionx body positiony body size fcbody colorfor body in bodysystem2 bodies doctest ellipsis updatestepbodysystem2 1 patches2 patches20 center 9 0 0 0 update the positions of the bodies update the positions of the patches utility function to plot how the given bodysystem evolves over time no doctest provided since this function does not have a return value each body is drawn as a patch by the pltfunction function called at each step of the animation example 1 figure8 solution to the 3bodyproblem this example can be seen as a test of the implementation given the right initial conditions the bodies should move in a figure8 initial conditions taken from http www artcompsci orgvol1v1webnode56 html bodysystem example1 lenbodysystem 3 example 2 moon s orbit around the earth this example can be seen as a test of the implementation given the right initial conditions the moon should orbit around the earth as it actually does mass velocity and distance taken from https en wikipedia orgwikiearth and https en wikipedia orgwikimoon no doctest provided since this function does not have a return value calculation of the respective velocities so that total impulse is zero i e the two bodies together don t move example 3 random system with many bodies no doctest provided since this function does not have a return value bodies are created pairwise with opposite velocities so that the total impulse remains zero frame rate of the animation time between time steps in seconds the parameters size color are not relevant for the simulation itself they are only used for plotting euler algorithm for velocity body_1 body 0 0 0 0 body_1 update_velocity 1 0 1 body_1 velocity 1 0 0 0 body_1 update_velocity 1 0 1 body_1 velocity 2 0 0 0 body_2 body 0 0 5 0 body_2 update_velocity 0 10 10 body_2 velocity 5 0 100 0 body_2 update_velocity 0 10 10 body_2 velocity 5 0 200 0 euler algorithm for position body_1 body 0 0 1 0 body_1 update_position 1 body_1 position 1 0 0 0 body_1 update_position 1 body_1 position 2 0 0 0 body_2 body 10 10 0 2 body_2 update_position 1 body_2 position 10 0 8 0 body_2 update_position 1 body_2 position 10 0 6 0 this class is used to hold the bodies the gravitation constant the time factor and the softening factor the time factor is used to control the speed of the simulation the softening factor is used for softening a numerical trick for n body simulations to prevent numerical divergences when two bodies get too close to each other for each body loop through all other bodies to calculate the total force they exert on it use that force to update the body s velocity body_system_1 bodysystem body 0 0 0 0 body 10 0 0 0 len body_system_1 2 body_system_1 update_system 1 body_system_1 bodies 0 position 0 01 0 0 body_system_1 bodies 0 velocity 0 01 0 0 body_system_2 bodysystem body 10 0 0 0 body 10 0 0 0 mass 4 1 10 body_system_2 update_system 1 body_system_2 bodies 0 position 9 0 0 0 body_system_2 bodies 0 velocity 0 1 0 0 calculation of the distance using pythagoras s theorem extra factor due to the softening technique newton s law of universal gravitation update the body s velocity once all the force components have been added update the positions only after all the velocities have been updated updates the body system and applies the change to the patch list used for plotting body_system_1 bodysystem body 0 0 0 0 body 10 0 0 0 patches_1 plt circle body position_x body position_y body size fc body color for body in body_system_1 bodies doctest ellipsis update_step body_system_1 1 patches_1 patches_1 0 center 0 01 0 0 body_system_2 bodysystem body 10 0 0 0 body 10 0 0 0 mass 4 1 10 patches_2 plt circle body position_x body position_y body size fc body color for body in body_system_2 bodies doctest ellipsis update_step body_system_2 1 patches_2 patches_2 0 center 9 0 0 0 update the positions of the bodies update the positions of the patches utility function to plot how the given body system evolves over time no doctest provided since this function does not have a return value set section to be plotted fix aspect ratio each body is drawn as a patch by the plt function function called at each step of the animation noqa f841 example 1 figure 8 solution to the 3 body problem this example can be seen as a test of the implementation given the right initial conditions the bodies should move in a figure 8 initial conditions taken from http www artcompsci org vol_1 v1_web node56 html body_system example_1 len body_system 3 example 2 moon s orbit around the earth this example can be seen as a test of the implementation given the right initial conditions the moon should orbit around the earth as it actually does mass velocity and distance taken from https en wikipedia org wiki earth and https en wikipedia org wiki moon no doctest provided since this function does not have a return value calculation of the respective velocities so that total impulse is zero i e the two bodies together don t move example 3 random system with many bodies no doctest provided since this function does not have a return value bodies are created pairwise with opposite velocities so that the total impulse remains zero
from __future__ import annotations

import random

from matplotlib import animation
from matplotlib import pyplot as plt


INTERVAL = 20


DELTA_TIME = INTERVAL / 1000


class Body:
    def __init__(
        self,
        position_x: float,
        position_y: float,
        velocity_x: float,
        velocity_y: float,
        mass: float = 1.0,
        size: float = 1.0,
        color: str = "blue",
    ) -> None:
        
        self.position_x = position_x
        self.position_y = position_y
        self.velocity_x = velocity_x
        self.velocity_y = velocity_y
        self.mass = mass
        self.size = size
        self.color = color

    @property
    def position(self) -> tuple[float, float]:
        return self.position_x, self.position_y

    @property
    def velocity(self) -> tuple[float, float]:
        return self.velocity_x, self.velocity_y

    def update_velocity(
        self, force_x: float, force_y: float, delta_time: float
    ) -> None:
        
        self.velocity_x += force_x * delta_time
        self.velocity_y += force_y * delta_time

    def update_position(self, delta_time: float) -> None:
        
        self.position_x += self.velocity_x * delta_time
        self.position_y += self.velocity_y * delta_time


class BodySystem:
    

    def __init__(
        self,
        bodies: list[Body],
        gravitation_constant: float = 1.0,
        time_factor: float = 1.0,
        softening_factor: float = 0.0,
    ) -> None:
        self.bodies = bodies
        self.gravitation_constant = gravitation_constant
        self.time_factor = time_factor
        self.softening_factor = softening_factor

    def __len__(self) -> int:
        return len(self.bodies)

    def update_system(self, delta_time: float) -> None:
        
        for body1 in self.bodies:
            force_x = 0.0
            force_y = 0.0
            for body2 in self.bodies:
                if body1 != body2:
                    dif_x = body2.position_x - body1.position_x
                    dif_y = body2.position_y - body1.position_y

                    
                    
                    distance = (dif_x**2 + dif_y**2 + self.softening_factor) ** (
                        1 / 2
                    )

                    
                    force_x += (
                        self.gravitation_constant * body2.mass * dif_x / distance**3
                    )
                    force_y += (
                        self.gravitation_constant * body2.mass * dif_y / distance**3
                    )

            
            body1.update_velocity(force_x, force_y, delta_time * self.time_factor)

        
        for body in self.bodies:
            body.update_position(delta_time * self.time_factor)


def update_step(
    body_system: BodySystem, delta_time: float, patches: list[plt.Circle]
) -> None:
    
    
    body_system.update_system(delta_time)

    
    for patch, body in zip(patches, body_system.bodies):
        patch.center = (body.position_x, body.position_y)


def plot(
    title: str,
    body_system: BodySystem,
    x_start: float = -1,
    x_end: float = 1,
    y_start: float = -1,
    y_end: float = 1,
) -> None:
    
    fig = plt.figure()
    fig.canvas.manager.set_window_title(title)
    ax = plt.axes(
        xlim=(x_start, x_end), ylim=(y_start, y_end)
    )  
    plt.gca().set_aspect("equal")  

    
    patches = [
        plt.Circle((body.position_x, body.position_y), body.size, fc=body.color)
        for body in body_system.bodies
    ]

    for patch in patches:
        ax.add_patch(patch)

    
    def update(frame: int) -> list[plt.Circle]:
        update_step(body_system, DELTA_TIME, patches)
        return patches

    anim = animation.FuncAnimation(  
        fig, update, interval=INTERVAL, blit=True
    )

    plt.show()


def example_1() -> BodySystem:
    

    position_x = 0.9700436
    position_y = -0.24308753
    velocity_x = 0.466203685
    velocity_y = 0.43236573

    bodies1 = [
        Body(position_x, position_y, velocity_x, velocity_y, size=0.2, color="red"),
        Body(-position_x, -position_y, velocity_x, velocity_y, size=0.2, color="green"),
        Body(0, 0, -2 * velocity_x, -2 * velocity_y, size=0.2, color="blue"),
    ]
    return BodySystem(bodies1, time_factor=3)


def example_2() -> BodySystem:
    

    moon_mass = 7.3476e22
    earth_mass = 5.972e24
    velocity_dif = 1022
    earth_moon_distance = 384399000
    gravitation_constant = 6.674e-11

    
    
    moon_velocity = earth_mass * velocity_dif / (earth_mass + moon_mass)
    earth_velocity = moon_velocity - velocity_dif

    moon = Body(-earth_moon_distance, 0, 0, moon_velocity, moon_mass, 10000000, "grey")
    earth = Body(0, 0, 0, earth_velocity, earth_mass, 50000000, "blue")
    return BodySystem([earth, moon], gravitation_constant, time_factor=1000000)


def example_3() -> BodySystem:
    

    bodies = []
    for _ in range(10):
        velocity_x = random.uniform(-0.5, 0.5)
        velocity_y = random.uniform(-0.5, 0.5)

        
        
        bodies.append(
            Body(
                random.uniform(-0.5, 0.5),
                random.uniform(-0.5, 0.5),
                velocity_x,
                velocity_y,
                size=0.05,
            )
        )
        bodies.append(
            Body(
                random.uniform(-0.5, 0.5),
                random.uniform(-0.5, 0.5),
                -velocity_x,
                -velocity_y,
                size=0.05,
            )
        )
    return BodySystem(bodies, 0.01, 10, 0.1)


if __name__ == "__main__":
    plot("Figure-8 solution to the 3-body-problem", example_1(), -2, 2, -2, 2)
    plot(
        "Moon's orbit around the earth",
        example_2(),
        -430000000,
        430000000,
        -430000000,
        430000000,
    )
    plot("Random system with many bodies", example_3(), -1.5, 1.5, -1.5, 1.5)