The n-body problem visualized with python

n_body_problem.py

import numpy as np
import matplotlib.pyplot as plt

# Constants
G = 6.67430e-11  # gravitational constant

def compute_accelerations(positions, masses):
    num_bodies = len(masses)
    accelerations = np.zeros_like(positions)
    
    for i in range(num_bodies):
        for j in range(num_bodies):
            if i != j:
                r_vec = positions[j] - positions[i]
                r_mag = np.linalg.norm(r_vec)
                accelerations[i] += G * masses[j] * r_vec / r_mag**3
    
    return accelerations

def simulate_orbits(initial_positions, initial_velocities, masses, time_step, num_steps):
    num_bodies = len(masses)
    positions = np.zeros((num_steps, num_bodies, 3))
    velocities = np.zeros_like(positions)
    
    positions[0] = initial_positions
    velocities[0] = initial_velocities
    
    for step in range(1, num_steps):
        accelerations = compute_accelerations(positions[step-1], masses)
        velocities[step] = velocities[step-1] + accelerations * time_step
        positions[step] = positions[step-1] + velocities[step] * time_step
    
    return positions

def plot_orbits(positions):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    
    for body_positions in positions.transpose((1, 0, 2)):
        ax.plot(body_positions[:, 0], body_positions[:, 1], body_positions[:, 2])
    
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    plt.show()

# Initial conditions (configurable)
masses = np.array([1.989e30, 5.972e24, 7.348e22, 6.417e23])  # Example: Sun, Earth, Moon, Mars
initial_positions = np.array([
    [0, 0, 0],                  # Sun
    [1.496e11, 0, 0],           # Earth
    [1.496e11 + 3.844e8, 0, 0], # Moon
    [2.279e11, 0, 0]            # Mars
])
initial_velocities = np.array([
    [0, 0, 0],                  # Sun
    [0, 29780, 0],              # Earth
    [0, 29780 + 1022, 0],       # Moon
    [0, 24007, 0]               # Mars
])

# Simulation parameters
time_step = 100000  # seconds
num_steps = 1000

# Run the simulation
positions = simulate_orbits(initial_positions, initial_velocities, masses, time_step, num_steps)

# Plot the results
plot_orbits(positions)