Last active
November 14, 2023 16:39
-
-
Save benrules2/220d56ea6fe9a85a4d762128b11adfba to your computer and use it in GitHub Desktop.
Python N-Body Orbit Simulation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import math | |
import random | |
import matplotlib.pyplot as plot | |
from mpl_toolkits.mplot3d import Axes3D | |
class point: | |
def __init__(self, x,y,z): | |
self.x = x | |
self.y = y | |
self.z = z | |
class body: | |
def __init__(self, location, mass, velocity, name = ""): | |
self.location = location | |
self.mass = mass | |
self.velocity = velocity | |
self.name = name | |
def calculate_single_body_acceleration(bodies, body_index): | |
G_const = 6.67408e-11 #m3 kg-1 s-2 | |
acceleration = point(0,0,0) | |
target_body = bodies[body_index] | |
for index, external_body in enumerate(bodies): | |
if index != body_index: | |
r = (target_body.location.x - external_body.location.x)**2 + (target_body.location.y - external_body.location.y)**2 + (target_body.location.z - external_body.location.z)**2 | |
r = math.sqrt(r) | |
tmp = G_const * external_body.mass / r**3 | |
acceleration.x += tmp * (external_body.location.x - target_body.location.x) | |
acceleration.y += tmp * (external_body.location.y - target_body.location.y) | |
acceleration.z += tmp * (external_body.location.z - target_body.location.z) | |
return acceleration | |
def compute_velocity(bodies, time_step = 1): | |
for body_index, target_body in enumerate(bodies): | |
acceleration = calculate_single_body_acceleration(bodies, body_index) | |
target_body.velocity.x += acceleration.x * time_step | |
target_body.velocity.y += acceleration.y * time_step | |
target_body.velocity.z += acceleration.z * time_step | |
def update_location(bodies, time_step = 1): | |
for target_body in bodies: | |
target_body.location.x += target_body.velocity.x * time_step | |
target_body.location.y += target_body.velocity.y * time_step | |
target_body.location.z += target_body.velocity.z * time_step | |
def compute_gravity_step(bodies, time_step = 1): | |
compute_velocity(bodies, time_step = time_step) | |
update_location(bodies, time_step = time_step) | |
def plot_output(bodies, outfile = None): | |
fig = plot.figure() | |
colours = ['r','b','g','y','m','c'] | |
ax = fig.add_subplot(1,1,1, projection='3d') | |
max_range = 0 | |
for current_body in bodies: | |
max_dim = max(max(current_body["x"]),max(current_body["y"]),max(current_body["z"])) | |
if max_dim > max_range: | |
max_range = max_dim | |
ax.plot(current_body["x"], current_body["y"], current_body["z"], c = random.choice(colours), label = current_body["name"]) | |
ax.set_xlim([-max_range,max_range]) | |
ax.set_ylim([-max_range,max_range]) | |
ax.set_zlim([-max_range,max_range]) | |
ax.legend() | |
if outfile: | |
plot.savefig(outfile) | |
else: | |
plot.show() | |
def run_simulation(bodies, names = None, time_step = 1, number_of_steps = 10000, report_freq = 100): | |
#create output container for each body | |
body_locations_hist = [] | |
for current_body in bodies: | |
body_locations_hist.append({"x":[], "y":[], "z":[], "name":current_body.name}) | |
for i in range(1,number_of_steps): | |
compute_gravity_step(bodies, time_step = 1000) | |
if i % report_freq == 0: | |
for index, body_location in enumerate(body_locations_hist): | |
body_location["x"].append(bodies[index].location.x) | |
body_location["y"].append(bodies[index].location.y) | |
body_location["z"].append(bodies[index].location.z) | |
return body_locations_hist | |
#planet data (location (m), mass (kg), velocity (m/s) | |
sun = {"location":point(0,0,0), "mass":2e30, "velocity":point(0,0,0)} | |
mercury = {"location":point(0,5.7e10,0), "mass":3.285e23, "velocity":point(47000,0,0)} | |
venus = {"location":point(0,1.1e11,0), "mass":4.8e24, "velocity":point(35000,0,0)} | |
earth = {"location":point(0,1.5e11,0), "mass":6e24, "velocity":point(30000,0,0)} | |
mars = {"location":point(0,2.2e11,0), "mass":2.4e24, "velocity":point(24000,0,0)} | |
jupiter = {"location":point(0,7.7e11,0), "mass":1e28, "velocity":point(13000,0,0)} | |
saturn = {"location":point(0,1.4e12,0), "mass":5.7e26, "velocity":point(9000,0,0)} | |
uranus = {"location":point(0,2.8e12,0), "mass":8.7e25, "velocity":point(6835,0,0)} | |
neptune = {"location":point(0,4.5e12,0), "mass":1e26, "velocity":point(5477,0,0)} | |
pluto = {"location":point(0,3.7e12,0), "mass":1.3e22, "velocity":point(4748,0,0)} | |
if __name__ == "__main__": | |
#build list of planets in the simulation, or create your own | |
bodies = [ | |
body( location = sun["location"], mass = sun["mass"], velocity = sun["velocity"], name = "sun"), | |
body( location = earth["location"], mass = earth["mass"], velocity = earth["velocity"], name = "earth"), | |
body( location = mars["location"], mass = mars["mass"], velocity = mars["velocity"], name = "mars"), | |
body( location = venus["location"], mass = venus["mass"], velocity = venus["velocity"], name = "venus"), | |
] | |
motions = run_simulation(bodies, time_step = 100, number_of_steps = 80000, report_freq = 1000) | |
plot_output(motions, outfile = 'orbits.png') |
Compliments for your work, but, But for the calculation of the distances traveled in a uniformly accelerated motion there is a littele error. The right equation, for uniformly accelerated motion is:
x=1/2 Axt^2; y=1/2 Ayt^2; z=1/2 Az*t^2;
were: Ax, Ay and Az are the components acceleration for x, y and z axis.
While the equation of velocity is right
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
this was really cool !
only i think you have an error in
run_simulation
where you callcompute_gravity_step(bodies, time_step = 1000)
instead ofcompute_gravity_step(bodies, time_step = time_step)