-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtest.py
116 lines (94 loc) · 4.99 KB
/
test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
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')