Skip to content

Commit 5bce725

Browse files
author
PhySimdev
committed
docs: finished Simulation.py
1 parent 44a48e8 commit 5bce725

File tree

2 files changed

+142
-60
lines changed

2 files changed

+142
-60
lines changed

Simulation.py

Lines changed: 139 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,20 @@
88

99

1010
class Simulation:
11+
"""Handles data and connection between simulation bodies."""
1112

1213
def __init__(self, theta=1, rc=0, absolute_pos=True, focus_index=0):
14+
"""Setup of the simulation.
1315
16+
Method that sets up the simulation with parameters and type of
17+
focus.
18+
19+
Args:
20+
theta: Theta value for the Barnes Hut simulation.
21+
rc: Restitution coefficient for collisions.
22+
absolute_pos: Bool value to determine type of movement.
23+
focus_index: Index of the list focus_options form 0 to 2.
24+
"""
1425
self.restitution_coefficient = rc
1526
self.focus_options = ["none", "body", "cm"]
1627
self.absolute_pos = absolute_pos
@@ -36,6 +47,9 @@ def __init__(self, theta=1, rc=0, absolute_pos=True, focus_index=0):
3647
def add_body(self, body):
3748
self.bodies.append(body)
3849

50+
def get_bodies(self):
51+
return self.bodies
52+
3953
def remove_body(self, body):
4054
self.bodies.remove(body)
4155
if body == self.focused_body and len(self.bodies) > 0:
@@ -51,7 +65,54 @@ def set_focus(self, body):
5165
else:
5266
self.focused_body = None
5367

68+
def update_center_of_mass(self, timestep):
69+
new_pos = np.array(self.tree.root_node.center_of_mass)
70+
old_pos = self.cm_pos
71+
self.cm_velo = (new_pos - old_pos) / timestep
72+
self.cm_pos = new_pos
73+
74+
def get_focus_pos(self):
75+
if self.focus_type == "body":
76+
pos = self.focused_body.position
77+
elif self.focus_type == "cm":
78+
pos = self.cm_pos
79+
else:
80+
pos = np.array([0, 0, 0])
81+
return pos
82+
83+
def switch_focus(self, direction):
84+
if self.focus_type == "body":
85+
focused_index = self.bodies.index(self.focused_body)
86+
87+
if direction == "previous":
88+
focused_index += 1
89+
if focused_index > len(self.bodies) - 1:
90+
focused_index = 0
91+
92+
elif direction == "next":
93+
focused_index -= 1
94+
if focused_index < 0:
95+
focused_index = len(self.bodies) - 1
96+
97+
self.set_focus(self.bodies[focused_index])
98+
99+
def clear_trail(self):
100+
for body in self.bodies:
101+
body.trail = []
102+
54103
def calculate(self, timestep, draw_box):
104+
"""Method that calculates a simulation physics step.
105+
106+
Method that calls functions for physics calculations.
107+
Also includes caclulations for the total energy. If draw_box is
108+
true the boxes of the octree are also extracted.
109+
110+
Args:
111+
timestep: Amount of seconds that are counted with in the
112+
physics calculations.
113+
draw_box: Boolean value that determines if cube data should be
114+
extracted. Implemented to run faster.
115+
"""
55116
if self.first:
56117
self.first = False
57118
self.update_interactions()
@@ -86,25 +147,16 @@ def calculate(self, timestep, draw_box):
86147
self.time += timestep
87148
self.iteration += 1
88149

89-
def update_center_of_mass(self, timestep):
90-
new_pos = np.array(self.tree.root_node.center_of_mass)
91-
old_pos = self.cm_pos
92-
self.cm_velo = (new_pos - old_pos) / timestep
93-
self.cm_pos = new_pos
150+
def get_data(self):
151+
"""Method that gets simulation data.
94152
95-
def get_bodies(self):
96-
return self.bodies
153+
Method that exports simulation data in the form of a list.
154+
The list includes almost all values of the simulation excluding the
155+
bodies themselves.
97156
98-
def get_focus_pos(self):
99-
if self.focus_type == "body":
100-
pos = self.focused_body.position
101-
elif self.focus_type == "cm":
102-
pos = self.cm_pos
103-
else:
104-
pos = np.array([0, 0, 0])
105-
return pos
106-
107-
def get_data(self):
157+
Returns:
158+
Body data, Octree data and system data for further usage.
159+
"""
108160
default_pos = self.get_focus_pos()
109161

110162
body_data = []
@@ -133,26 +185,6 @@ def get_data(self):
133185

134186
return body_data, self.tree_nodes, system_data, self.cm_pos - default_pos
135187

136-
def switch_focus(self, direction):
137-
if self.focus_type == "body":
138-
focused_index = self.bodies.index(self.focused_body)
139-
140-
if direction == "previous":
141-
focused_index += 1
142-
if focused_index > len(self.bodies) - 1:
143-
focused_index = 0
144-
145-
elif direction == "next":
146-
focused_index -= 1
147-
if focused_index < 0:
148-
focused_index = len(self.bodies) - 1
149-
150-
self.set_focus(self.bodies[focused_index])
151-
152-
def clear_trail(self):
153-
for body in self.bodies:
154-
body.trail = []
155-
156188
def update_interactions(self):
157189
center = self.get_focus_pos()
158190

@@ -174,9 +206,15 @@ def update_interactions(self):
174206
self.compute_collisions(self.tree.collision_dic)
175207

176208
def compute_collisions(self, collisions):
177-
# sort collisions from lowest to highest mass
178-
# we know that body has bigger mass than other so other is deleted
179-
# we have to check if other was deleted by a previous collision with another planet
209+
"""Method that computes body collisions.
210+
211+
Method that computes body collisions based on the resitution
212+
coefficient. Sort collisions from lowest to highest mass. So all
213+
collisions can be determined. Calls body to body inelastic collision.
214+
215+
Args:
216+
collisions: All collisons as extractes from the Barnes-Hut program.
217+
"""
180218
self.destroyed = []
181219
bodies = list(collisions.keys())
182220
bodies.sort(key=lambda element: element.mass)
@@ -188,10 +226,25 @@ def compute_collisions(self, collisions):
188226

189227

190228
class SimulationBody:
191-
192-
def __init__(self, solar_system, name, mass, density, position, velocity, color, nr_pos, point_dist):
193-
194-
self.solar_system = solar_system
229+
"""Data storage and caclulation of a simulation body"""
230+
231+
def __init__(self, simulation, name, mass, density, position, velocity, color, nr_pos, point_dist):
232+
"""Method that sets up a simulation body.
233+
234+
Method sets up and declares variables for a simulation body.
235+
236+
Args:
237+
simulation: Simulation that stores all other bodies.
238+
name: Name/Unique identifier of body.
239+
mass: Mass of the body in kg.
240+
density: Density of the body in g/cm3.
241+
position: Position vector in meters.
242+
velocity: Velocity vector in m/s.
243+
color: Color of the body.
244+
nr_pos: Trail node number of the body.
245+
point_dist: Distance between nodes in the trail.
246+
"""
247+
self.simulation = simulation
195248
self.name = name
196249
self.mass = mass # mass in kg
197250
self.density = density # density in g/cm^3
@@ -207,16 +260,25 @@ def __init__(self, solar_system, name, mass, density, position, velocity, color,
207260
self.force = np.array([0, 0, 0])
208261
self.e_pot = 0
209262
self.e_kin = 0
210-
self.solar_system.add_body(self)
263+
self.simulation.add_body(self)
211264

212265
def update_position(self, timestep):
266+
"""Method that calculates the body position.
267+
268+
Method that calculates the body position follwing verlet velocity
269+
integration. Subtracts movement of focus_body if the view is movement
270+
is relative.
271+
272+
Args:
273+
timestep: Amount of seconds for the calculations.
274+
"""
213275
self.position += (timestep * (self.velocity + timestep * self.acceleration / 2))
214276

215-
if not self.solar_system.absolute_pos: # relative movement
216-
if self.solar_system.focus_type == "body":
217-
self.position -= (timestep * (self.solar_system.focused_body.velocity + timestep * self.solar_system.focused_body.acceleration / 2))
218-
elif self.solar_system.focus_type == "cm":
219-
self.position -= self.solar_system.cm_velo
277+
if not self.simulation.absolute_pos: # relative movement
278+
if self.simulation.focus_type == "body":
279+
self.position -= (timestep * (self.simulation.focused_body.velocity + timestep * self.simulation.focused_body.acceleration / 2))
280+
elif self.simulation.focus_type == "cm":
281+
self.position -= self.simulation.cm_velo
220282

221283
self.counter += 1
222284
if self.counter == self.point_dist:
@@ -227,6 +289,14 @@ def update_position(self, timestep):
227289
del self.trail[0]
228290

229291
def update_velocity(self, timestep):
292+
"""Method that calculates body velocity.
293+
294+
Method that calculates body velocity following velocity
295+
verlet ingtegration.
296+
297+
Args:
298+
timestep: Amount of seconds for the calculations.
299+
"""
230300
newacc = self.force / self.mass
231301
self.velocity += (self.acceleration + newacc) * timestep / 2
232302
self.acceleration = newacc
@@ -237,6 +307,18 @@ def calculate_radius(self):
237307
return ((3 * self.mass) / (4 * math.pi * density_kg))**(1 / 3)
238308

239309
def inelastic_collision(self, other, restitution_coefficient):
310+
"""Method that calculates an inelastic collision.
311+
312+
Method that calculates the collision between two objects based
313+
on the restitution coefficient. If that coefficient if 0 the planets
314+
merge. If the coefficient is higher than 0 to 1 the collision is more
315+
elastic. Destroys body that is smaller than the other.
316+
317+
Args:
318+
other: Other body that takes part in the collision.
319+
restitution_coefficient: Coefficient that determines type of
320+
collision.
321+
"""
240322
# only if restitution_coefficient is 0 bodies merge
241323
velo_u = ((self.mass * self.velocity) + (other.mass * other.velocity)) / (self.mass + other.mass)
242324

@@ -245,7 +327,7 @@ def inelastic_collision(self, other, restitution_coefficient):
245327
self.velocity = velo_u
246328
self.mass += other.mass
247329
self.radius = self.calculate_radius()
248-
self.solar_system.destroyed.append(other)
330+
self.simulation.destroyed.append(other)
249331
else:
250332
# somewhat elastic collision
251333
r = restitution_coefficient
@@ -256,35 +338,35 @@ def inelastic_collision(self, other, restitution_coefficient):
256338
# ---------------------------------------------------------------------------------------------------------------------
257339
class Sun(SimulationBody):
258340

259-
def __init__(self, solar_system, name="default sun", sol_mass=1, density=1.41, position_AU=(0, 0, 0), velocity_km=(0, 0, 0), color="yellow", nr_pos=50, point_dist=10):
341+
def __init__(self, simulation, name="default sun", sol_mass=1, density=1.41, position_AU=(0, 0, 0), velocity_km=(0, 0, 0), color="yellow", nr_pos=50, point_dist=10):
260342
au = 1.496 * 10**11
261343
km = 10**3
262344
solM = 1.989 * (10**30)
263345
mass = sol_mass * solM # multiple of solar mass
264346
position = tuple(au * pos for pos in position_AU)
265347
velocity = tuple(km * velo for velo in velocity_km)
266-
super(Sun, self).__init__(solar_system, name, mass, density, position, velocity, color, nr_pos, point_dist)
348+
super(Sun, self).__init__(simulation, name, mass, density, position, velocity, color, nr_pos, point_dist)
267349

268350

269351
class Planet(SimulationBody):
270352

271-
def __init__(self, solar_system, name="default planet", ea_mass=1, density=5.5, position_AU=(0, 0, 0), velocity_km=(0, 0, 0), color="blue", nr_pos=50, point_dist=10):
353+
def __init__(self, simulation, name="default planet", ea_mass=1, density=5.5, position_AU=(0, 0, 0), velocity_km=(0, 0, 0), color="blue", nr_pos=50, point_dist=10):
272354
au = 1.496 * 10**11
273355
km = 10**3
274356
eam = 5.972 * 10**24
275357
mass = ea_mass * eam # multiple of solar mass
276358
position = tuple(au * pos for pos in position_AU)
277359
velocity = tuple(km * velo for velo in velocity_km)
278-
super(Planet, self).__init__(solar_system, name, mass, density, position, velocity, color, nr_pos, point_dist)
360+
super(Planet, self).__init__(simulation, name, mass, density, position, velocity, color, nr_pos, point_dist)
279361

280362

281363
class Blackhole(SimulationBody):
282364

283-
def __init__(self, solar_system, name="default hole", sol_mass=10, density=4 * 10**14, position_AU=(0, 0, 0), velocity_km=(0, 0, 0), color="black", nr_pos=50, point_dist=10):
365+
def __init__(self, simulation, name="default hole", sol_mass=10, density=4 * 10**14, position_AU=(0, 0, 0), velocity_km=(0, 0, 0), color="black", nr_pos=50, point_dist=10):
284366
au = 1.496 * 10**11
285367
km = 10**3
286368
solM = 1.989 * (10**30)
287369
mass = sol_mass * solM # multiple of solar mass
288370
position = tuple(au * pos for pos in position_AU)
289371
velocity = tuple(km * velo for velo in velocity_km)
290-
super(Planet, self).__init__(solar_system, name, mass, density, position, velocity, color, nr_pos, point_dist)
372+
super(Planet, self).__init__(simulation, name, mass, density, position, velocity, color, nr_pos, point_dist)

main.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -450,9 +450,9 @@ def get_screen_xy(self, x, y, z):
450450
transformation and distance aswell as FOV.
451451
452452
Args:
453-
x: X component of position in 3d space
454-
y: Y component of position in 3d space
455-
z: Z component of position in 3d space
453+
x: X component of position in 3d space.
454+
y: Y component of position in 3d space.
455+
z: Z component of position in 3d space.
456456
457457
Returns:
458458
Either returns a xy position tuple and the point distance

0 commit comments

Comments
 (0)