Last active
August 29, 2015 13:56
-
-
Save zeimusu/9280735 to your computer and use it in GitHub Desktop.
A gravity simulator. There are several choices for the integration method. I've found verlet to be most stable
This file contains 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
#!/usr/bin/env python3 | |
import tkinter as tk | |
import numpy as np | |
from numpy import sin, cos, sqrt | |
G = 6.67384e-11 # m^3 kg-1 s-2 | |
DEG = np.pi / 180 | |
AU = 1.49597871e11 # for converting to meters | |
METRE = 1 / 1.49597871e11 # for converting to AU | |
KG = 1 / 1.9891e30 # units of solar masses | |
SEC = 1 / (60 * 60 * 24) # seconds in a day | |
GG = 0.0002959122 # in these units 0.0002959122 | |
class Body: | |
def __init__(self, r, v, mass, col='white'): | |
self.last_r = np.array(r, dtype="float") | |
self.r = np.array(r, dtype="float") | |
self.v = np.array(v, dtype="float") | |
self.mass = mass | |
self.radius = 1.0 | |
self.GM = self.mass * GG | |
self.col = col | |
def orbital_elements(self, principal): | |
"""http://orbitsimulator.com/formulas/OrbitalElements.html""" | |
mu = GG * (principal.mass + self.mass) | |
# calculate relative position,velocity | |
r = self.r - principal.r | |
v = self.v - principal.v | |
try: # catch division by zero | |
R = np.linalg.norm(r) | |
V = np.linalg.norm(v) | |
a = 1 / (2 / R - V ** 2 / mu) # semi major axis | |
h = np.cross(r, v) | |
H = np.linalg.norm(h) | |
P = H ** 2 / mu | |
Q = np.dot(r, v) | |
E = np.sqrt(1 - P / a) # eccentricity | |
e = [1 - R / a, Q / np.sqrt(a * mu)] | |
i = np.arccos(h[2] / H) | |
LAN = 0 | |
if i != 0: | |
LAN = np.arctan2(h[0], -h[1]) # Longitude of acending node | |
ta = [H ** 2 / (R * mu) - 1, H * Q / (R * mu)] | |
TA = np.arctan2(ta[1], ta[0]) | |
Cw = (r[0] * cos(LAN) + r[1] * sin(LAN)) / R | |
if i == 0 or i == np.pi: | |
Sw = (r[1] * cos(LAN) - r[0] * sin(LAN)) / R | |
else: | |
Sw = r[2] / (R * sin(i)) | |
omega = np.arctan2(Sw, Cw) - TA # argument of periapsis | |
if omega < 0: | |
omega += 2 * np.pi | |
u = np.arctan2(e[1], e[0]) | |
M = u - E * sin(u) # mean anomaly | |
return (a, E, omega, LAN, i, M) | |
except ZeroDivisionError: | |
return (0, 0, 0, 0, 0, 0) | |
class System: | |
def __init__(self, mass): | |
self.bodies = { | |
"Sun": Body([0.0, 0.0, 0.0], [0.0, 0.0, 0.0], mass), | |
} | |
self.time = 0.0 | |
def add_body(self, name, r, v, m, col): | |
self.bodies[name] = Body(r, v, m, col) | |
def add_satelite(self, name, principal, a, e, omega, LAN, i, M0, | |
m=1.0, col="white", radius=100, epoch=2000.0): | |
""" | |
http://downloads.rene-schwarz.com/download/M001-Keplerian_Orbit_Elements_to_Cartesian_State_Vectors.pdf | |
""" | |
r, v = self.state_vector(principal, a, e, omega, LAN, i, M0, m) | |
r += self.bodies[principal].r | |
v += self.bodies[principal].v | |
self.add_body(name, r, v, m, col) | |
def add_body_elements(self, name, a, e, omega, LAN, i, M0, | |
m=1.0, col="white", radius=100, epoch=2000.0): | |
self.add_satelite(name, "Sun", a, e, omega, LAN, i, M0, | |
m=m, col=col, radius=radius, epoch=2000.0) | |
def state_vector(self, principal, | |
a, e, omega, LAN, i, M0, m, epoch=2000.0): | |
# Calcuate mean anomaly | |
mu = GG * (m + self.bodies[principal].mass) | |
M = M0 + np.sqrt(mu / a ** 3) * self.time | |
np.fmod(M, 2 * np.pi) | |
# solve Kepler's equation, using Newtons method. | |
epsilon = 1e-6 | |
E = M | |
while True: | |
nextE = E - (E - e * np.sin(E) - M) / (1 - e * cos(E)) | |
E = nextE | |
if abs(E - nextE) < epsilon: | |
break | |
E = nextE | |
# find true anomoly | |
nu = 2 * np.arctan2(np.sqrt(1 + e) * np.sin(E / 2), | |
sqrt(1 - e) * cos(E / 2)) | |
# find radial distance | |
r = a * (1 - e * cos(E)) | |
# find position and velocity in coordinates of orbit | |
o = r * np.array([cos(nu), np.sin(nu), 0]) | |
do = np.sqrt(mu * a) / r * np.array([-np.sin(E), | |
(1 - e ** 2) * cos(E), 0]) | |
# transform to inertial frame | |
co = cos(omega) | |
cO = cos(LAN) | |
so = sin(omega) | |
sO = sin(LAN) | |
si = sin(i) | |
ci = cos(i) | |
r = np.array([ | |
o[0] * (co * cO - so * ci * sO) - o[1] * (so * cO + co * ci * sO), | |
o[0] * (co * sO + so * ci * cO) + o[1] * (co * ci * cO - so * sO), | |
o[0] * so * si + o[1] * co * si]) | |
v = np.array([ | |
do[0] * (co * cO - so * ci * sO) - do[1] * (so * cO + co * ci * sO), | |
do[0] * (co * sO + so * ci * cO) + do[1] * (co * ci * cO - so * sO), | |
do[0] * so * si + do[1] * co * si]) | |
return r, v | |
def momentum(self): | |
momentum = np.zeros(3) | |
mass = 0 | |
for (_, body) in self.bodies.items(): | |
momentum += body.v * body.mass | |
mass += body.mass | |
return momentum, mass | |
def collide(self, body_name1, body_name2): | |
b1 = self.bodies[body_name1] | |
b2 = self.bodies[body_name2] | |
total_momentum = b1.v*b1.mass+b2.v*b2.mass | |
newbody = Body((b1.r + b2.r) * 0.5, | |
total_momentum / (b1.mass + b2.mass), | |
b1.mass+b2.mass, col=b1.col) | |
newname = body_name1[:5]+body_name2[:5] | |
if newname in self.bodies: | |
newname = 'obj' + body_name1 + body_name2 | |
del self.bodies[body_name1] | |
del self.bodies[body_name2] | |
self.bodies[newname] = newbody | |
def aj(self, body_name, r): | |
""" Calculates accelerations due to other bodies at a | |
supposed position 'r'""" | |
acc = np.zeros(3) | |
for each_body_name, each_body in self.bodies.items(): | |
if each_body_name == body_name: | |
continue | |
rvec = each_body.r - r | |
# calculate 1/|rvec|^(3/2) | |
sep3 = (rvec[0] * rvec[0] + | |
rvec[1] * rvec[1] + | |
rvec[2] * rvec[2]) ** (-1.5) | |
# direct calculation is faster than linalg norm | |
# sep = np.linalg.norm(rvec) | |
acc += (each_body.GM * sep3) * rvec | |
return acc | |
def dst_rk(self, body_name, time_step): | |
"""Calculate 4 acceleration vectors: | |
k1 = the acceleration at the current postion. | |
k2 = the acceleration at t+step/2 using a position based on k1 | |
k3 = acceleration at t+step/2 using a postion calulated using k2 | |
k4 = acceleration at t+step using a postion calculated using k3 | |
the estimate of average accelartion is then (k1 + 2k2 + 2k3 + k4)/6 | |
this version doesn't try to be clever about acc | |
http://spiff.rit.edu/richmond/nbody/OrbitRungeKutta4.pdf | |
""" | |
body = self.bodies[body_name] | |
t2 = time_step * 0.5 | |
# calculate the acceleration at time 0 | |
k1v = self.aj(body_name, body.r) | |
k1r = body.v | |
k2v = self.aj(body_name, body.r + k1r * t2) | |
k2r = body.v + k1v * t2 | |
k3v = self.aj(body_name, body.r + k2r * t2) | |
k3r = body.v + k2v * t2 | |
k4v = self.aj(body_name, body.r + k3r * time_step) | |
k4r = body.v + k3v * t2 | |
return ((k1v + 2 * (k2v + k3v) + k4v) / 6, | |
(k2r + 2 * (k2r + k3r) + k4r) / 6) | |
def suvat_rk(self, body_name, time_step): | |
"""Calculate 4 acceleration vectors: | |
k1 = the acceleration at the current postion. | |
k2 = the acceleration at t+step/2 using a position based on k1 | |
k3 = acceleration at t+step/2 using a postion calulated using k2 | |
k4 = acceleration at t+step using a postion calculated using k3 | |
this version attempts to be clever about the positions by using | |
suvat | |
the estimate of average accelartion is then (k1 + 2k2 + 2k3 + k4)/6 | |
""" | |
body = self.bodies[body_name] | |
# k1 = self.aj(body_name, body.r) | |
# often this will be an acceptable estimate | |
# could put some code here to detact small accelarations | |
# for which calcuation of k2-4 is wasteful | |
# or large accelerations for which time step is probably inaccurate | |
t2 = time_step * 0.5 | |
# calculate the acceleration at time 0 | |
k1v = self.aj(body_name, body.r) | |
k1r = body.v | |
k2v = self.aj(body_name, body.r + k1r * t2 + 0.5 * k1v * t2 * t2) | |
k2r = body.v + k1v * t2 | |
k3v = self.aj(body_name, body.r + k2r * t2 + 0.5 * k2v * t2 * t2) | |
k3r = body.v + k2v * t2 | |
k4v = self.aj(body_name, | |
body.r + k3r * time_step + | |
0.5 * k3v * time_step * time_step) | |
k4r = body.v + k3v * t2 | |
return ((k1v + 2 * (k2v + k3v) + k4v) / 6, | |
(k2r + 2 * (k2r + k3r) + k4r) / 6) | |
def calc_acceleration(self, time_step): | |
acc_dict = dict([(body_name, self.rk(body_name, time_step)) | |
for body_name in self.bodies]) | |
return acc_dict | |
def simple_update(self, time_step): | |
""" plug in replacement for update. | |
calculates acceleration without rk | |
simple integration""" | |
acc_dict = dict([(body_name, | |
self.aj(body_name, self.bodies[body_name].r)) | |
for body_name in self.bodies]) | |
for body_name, body in self.bodies.items(): | |
at = acc_dict[body_name] * time_step | |
body.v += at | |
body.r += body.v * time_step # dst | |
self.time += time_step | |
def verlet_init(self, time_step): | |
for body_name, body in self.bodies.items(): | |
acc = self.aj(body_name, body.r) | |
body.r += body.v * time_step + 0.5 * acc * time_step ** 2 | |
def verlet_update(self, time_step): | |
for body_name, body in self.bodies.items(): | |
a = self.aj(body_name, body.r) | |
body.last_r, body.r = (body.r, | |
2 * body.r - body.last_r + | |
a * time_step * time_step) | |
self.time += time_step | |
def time_corrected_verlet(self, old_step, new_step): | |
timeratio = new_step/old_step | |
for body_name, body in self.bodies.items(): | |
a = self.aj(body_name, body.r) | |
body.last_r, body.r = (body.r, | |
body.r + | |
(body.r - body.last_r) * (timeratio) + | |
a * new_step * new_step) | |
self.time += new_step | |
def velocity_verlet(self, time_step): | |
acc_dict = dict([(body_name, | |
self.aj(body_name, self.bodies[body_name].r)) | |
for body_name in self.bodies]) | |
for body_name, body in self.bodies.items(): | |
r_dash = (body.r + body.v * time_step + | |
0.5 * acc_dict[body_name] * time_step * time_step) | |
a = self.aj(body_name, r_dash) | |
body.v += (acc_dict[body_name] + a) * 0.5 * time_step | |
body.r += body.v * time_step | |
def dst_update(self, time_step): | |
acc_dict = dict([(body_name, | |
self.dst_rk(body_name, time_step)) | |
for body_name in self.bodies]) | |
print(acc_dict) | |
for body_name, body in self.bodies.items(): | |
body.v += acc_dict[body_name][0] * time_step | |
body.r += acc_dict[body_name][1] * time_step | |
self.time += time_step | |
def suvat_update(self, time_step): | |
acc_dict = dict([(body_name, | |
self.rk(body_name, time_step)) | |
for body_name in self.bodies]) | |
for body_name, body in self.bodies.items(): | |
at = acc_dict[body_name][0] * time_step | |
vt = acc_dict[body_name][1] * time_step | |
body.v += at | |
body.r += (vt - 0.5 * at * time_step) # suvat | |
self.time += time_step | |
update = verlet_update | |
rk = dst_rk | |
class View: | |
def __init__(self, parent, system): | |
self.parent = parent | |
self.time_step = 20000.0*SEC | |
self.system = system | |
self.width = 600 | |
self.height = 600 | |
self.center_x = self.width / 2 | |
self.center_y = self.height / 2 | |
self.scale = self.height | |
self.view_angle = np.pi / 3 | |
self.cos_view = np.cos(self.view_angle) | |
self.sin_view = sin(self.view_angle) | |
self.focus = 'Sun' | |
self.time = tk.StringVar() | |
self.time.set("Time: " + str(system.time)) | |
self.system_view = tk.Canvas(parent, | |
width=self.width, height=self.height, | |
background='black') | |
self.control = tk.Frame(parent) | |
self.running = tk.BooleanVar() | |
self.running.set(False) | |
self.trace = True | |
tk.Label(self.control, textvariable=self.time).pack(side=tk.TOP) | |
self.run_button = tk.Checkbutton(self.control, command=self.run, | |
text="run", variable=self.running) | |
self.step_button = tk.Button(self.control, | |
command=self.run, text="step") | |
self.zoomin_button = \ | |
tk.Button(self.control, | |
command=lambda factor=1.25: self.zoom(factor), | |
text="Zoom in") | |
self.zoomout_button = \ | |
tk.Button(self.control, | |
command=lambda factor=0.8: self.zoom(factor), | |
text="Zoom out") | |
self.faster_but =\ | |
tk.Button(self.control, | |
command=lambda rate=1.25: self.faster(rate), | |
text="Faster") | |
self.slower_but = \ | |
tk.Button(self.control, | |
command=lambda rate=0.8: self.faster(rate), | |
text="Slower") | |
self.run_button.pack(side=tk.TOP) | |
self.step_button.pack(side=tk.TOP) | |
self.zoomin_button.pack(side=tk.TOP) | |
self.zoomout_button.pack(side=tk.TOP) | |
self.faster_but.pack(side=tk.TOP) | |
self.slower_but.pack(side=tk.TOP) | |
tk.Button(self.control, | |
command=self.obj_info, text="Info").pack() | |
self.system_view.pack(side=tk.LEFT) | |
self.control.pack(side=tk.LEFT) | |
self.sso = self.draw_system() | |
def set_view_angle(self, angle): | |
self.view_angle = angle | |
self.cos_view = np.cos(self.view_angle) | |
self.sin_view = sin(self.view_angle) | |
def faster(self, rate): | |
new_step = self.time_step * rate | |
self.system.time_corrected_verlet(self.time_step, new_step) | |
self.time_step = new_step | |
def zoom(self, factor): | |
self.system_view.delete(tk.ALL) | |
self.scale *= factor | |
self.sso = self.draw_system() | |
def draw_system(self): | |
sso = {} | |
if self.focus: | |
center_r = self.system.bodies[self.focus].r | |
else: | |
center_r = [0, 0] | |
for name, body in self.system.bodies.items(): | |
x = (body.r[0] - center_r[0]) * self.scale + self.center_x | |
y = ((body.r[1] - center_r[1]) * self.cos_view + | |
body.r[2] * self.sin_view) * self.scale + self.center_y | |
sso[name] = self.system_view.create_oval(x - 2, y - 2, | |
x + 2, y + 2, | |
fill=body.col) | |
return sso | |
def redraw(self): | |
if self.focus: | |
center_r = self.system.bodies[self.focus].r | |
else: | |
center_r = [0, 0] | |
for name, ident in self.sso.items(): | |
body = self.system.bodies[name] | |
x = (body.r[0] - center_r[0]) * self.scale + self.center_x | |
y = ((body.r[1] - center_r[1]) * self.cos_view + | |
body.r[2] * self.sin_view) * self.scale + self.center_y | |
self.system_view.create_line(x, y, x + 1, y + 1, fill=body.col) | |
self.system_view.coords(ident, x - 2, y - 2, x + 2, y + 2) | |
def run(self): | |
self.system.update(self.time_step) | |
self.time.set("Time: " + str(int(self.system.time))) | |
self.redraw() | |
if self.running.get(): | |
self.parent.after(1, self.run) | |
def obj_info(self): | |
parent = self.parent | |
system = self.system | |
top = tk.Toplevel(parent) | |
objects = system.bodies.keys() | |
chosenobj = tk.StringVar() | |
chosenobj.set('Sun') | |
planetmenu = tk.OptionMenu(top, chosenobj, *objects) | |
planetmenu.pack() | |
body = system.bodies[chosenobj.get()] | |
elements = body.orbital_elements(system.bodies['Sun']) | |
vectors = (body.r, body.v, body.mass) | |
print(elements) | |
print(vectors) | |
closebutton = tk.Button(top, | |
command=lambda: top.destroy(), | |
text="Close") | |
closebutton.pack() | |
top.update() | |
def make_solar_system(): | |
solar_system = System(1) | |
solar_system.add_body_elements( | |
'Mercury', | |
a=5.7e10*METRE, e=0.2056, omega=29.124*DEG, LAN=48.331*DEG, | |
i=7.0*DEG, M0=174*DEG, | |
m=3.3e23*KG, col='gray') | |
solar_system.add_body_elements( | |
'Venus', | |
a=1.082e11*METRE, e=0.0067, omega=55.186*DEG, LAN=76.687*DEG, | |
i=3.394*DEG, M0=50.115*DEG, | |
m=4.868e24*KG, col='white') | |
solar_system.add_body_elements( | |
'Earth', | |
a=1.49598e11*METRE, e=0.01671, omega=114.207*DEG, LAN=348.739*DEG, | |
i=0, M0=357.517*DEG, | |
m=5.972e24*KG, col='blue') | |
solar_system.add_body_elements( | |
'Mars', | |
a=2.279e11*METRE, e=0.093, omega=286.5*DEG, LAN=49.562*DEG, | |
i=1.85*DEG, M0=19.35*DEG, | |
m=6.418e23*KG, col='red') | |
solar_system.add_body_elements( | |
'Jupiter', | |
a=7.785e11*METRE, e=0.04877, omega=275.1*DEG, LAN=100.5*DEG, | |
i=1.305*DEG, M0=18.818*DEG, | |
m=1.899e27*KG, col='yellow') | |
solar_system.add_body_elements( | |
'Saturn', | |
a=1.433e12*METRE, e=0.055, omega=336*DEG, LAN=113.6*DEG, | |
i=2.485*DEG, M0=320.346*DEG, | |
m=5.68e26*KG, col='brown') | |
solar_system.add_body_elements( | |
'Uranus', | |
a=2.876e12*METRE, e=0.0444, omega=96.541*DEG, LAN=73.989*DEG, | |
i=0.772*DEG, M0=142.9*DEG, | |
m=8.68e25*KG, col='purple') | |
solar_system.add_body_elements( | |
'Neptune', | |
a=4.503e12*METRE, e=0.0112, omega=256.6*DEG, LAN=131.8*DEG, | |
i=1.767*DEG, M0=267.76*DEG, | |
m=1.02e26*KG, col='green') | |
solar_system.add_body_elements( | |
'Halley', | |
a=17.8, e=0.967, omega=111.8657*DEG, LAN=58.8601*DEG, | |
i=162.242*DEG, M0=66.26*DEG, | |
m=1e15*KG, col='white') | |
zero_momentum(solar_system) | |
return (solar_system) | |
def zero_momentum(system): | |
"""change frame to give zero total momentum.""" | |
momen, mass = system.momentum() | |
adj_v = momen / mass | |
for body in system.bodies.values(): | |
body.v -= adj_v | |
def random_sys(): | |
import random | |
system = System(1e30) | |
for i in range(5): | |
system.add_body('body' + str(i), | |
[random.uniform(-1e11*METRE, 1e11*METRE), | |
random.uniform(-1e11*METRE, 1e11*METRE), 0], | |
[0, 0, 0], | |
1e29*KG, 'white') | |
return system | |
def circle_system(): | |
system = System(1e30*KG) | |
system.add_body_elements("Planet", | |
a=1e10*METRE, e=0, omega=0, LAN=0, | |
i=0*DEG, M0=0*DEG, | |
m=1e15*KG, col="red") | |
return system | |
def sem_system(): | |
system = System(1) | |
system.add_body_elements( | |
'Earth', | |
a=1.49598e11*METRE, e=0.01671, omega=114.207*DEG, LAN=348.739*DEG, | |
i=0*DEG, M0=357.517*DEG, | |
m=5.972e24*KG, col='blue') | |
system.add_satelite( | |
'Moon', "Earth", | |
a=3.85e8*METRE, e=0.055, omega=318.1*DEG, LAN=125.1*DEG, | |
i=5.145*DEG, M0=0*DEG, | |
m=7.35e22*KG, col='gray') | |
system.focus = 'Earth' | |
return system | |
def em_system(): | |
system = System(5.97e24*KG) | |
system.add_satelite( | |
'Moon', "Sun", | |
a=3.85e8*METRE, e=0.055, omega=318.1*DEG, LAN=125.1*DEG, | |
i=5.145*DEG, M0=0*DEG, | |
m=7.35e22*KG, col='gray') | |
return system | |
def main(): | |
solar_system = make_solar_system() | |
root = tk.Tk() | |
view = View(root, solar_system) | |
view.time_step = 10000*SEC | |
solar_system.verlet_init(view.time_step) | |
view.set_view_angle(30 * DEG) | |
view.focus = 'Sun' | |
tk.mainloop() | |
if __name__ == "__main__": | |
main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment