Skip to content

Instantly share code, notes, and snippets.

@zeimusu
Last active August 29, 2015 13:56
Show Gist options
  • Save zeimusu/9280735 to your computer and use it in GitHub Desktop.
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
#!/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