simulation

simulation.py

import numpy as np
import rebound
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


def add_galactic_potential(sim, mass_distribution):
    """
    Simplified function to simulate the galactic potential effect on the simulation.
    This function modifies the velocities of particles to simulate a basic galactic potential.
    """
    for p in sim.particles[1:]:  # Skip the first particle assuming it's the central object (e.g., Sun)
        p.vx += np.random.uniform(-1e-3, 1e-3)
        p.vy += np.random.uniform(-1e-3, 1e-3)


def add_relativistic_corrections(sim):
    """
    Simplified function to add relativistic corrections.
    This example adjusts particles' velocities to simulate relativistic effects.
    """
    for p in sim.particles[1:]:  # Again, skip the Sun
        p.vz += np.random.uniform(-1e-5, 1e-5)


def simulate_solar_system(total_years=1000, num_data_points=1000):
    sim = rebound.Simulation()
    sim.integrator = 'ias15'
    sim.add(m=1.0)  # Add the Sun

    # Example addition of Earth for simplicity
    sim.add(m=3e-6, a=1.0)  # Earth, simplified: mass and semi-major axis

    add_galactic_potential(sim, 'simplified_model')
    add_relativistic_corrections(sim)

    times = np.linspace(0, total_years, num_data_points)
    positions = np.zeros((len(sim.particles), num_data_points, 3))

    for i, t in enumerate(times):
        sim.integrate(t)
        for j, p in enumerate(sim.particles):
            positions[j, i] = p.xyz

    return positions


# Constants and initial setup
num_points = 1000
positions = simulate_solar_system(total_years=1000, num_data_points=num_points)

# Plotting
fig = plt.figure(figsize=(12, 9))
ax = fig.add_subplot(111, projection='3d')

# Plot the Sun's and Earth's positions
ax.plot(*positions[0].T, label='Sun')
ax.plot(*positions[1].T, label='Earth')

ax.set_xlabel('X (AU)')
ax.set_ylabel('Y (AU)')
ax.set_zlabel('Z (AU)')
ax.legend()

# Save the figure instead of trying to show it interactively
plt.savefig('./simulation_plot.png')
# Ensure to replace '/path/to/your/directory/' with the actual path where you want to save the image.

print("The plot has been saved as 'simulation_plot.png'.")