{{ message }}
forked from TheAlgorithms/Python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathn_body_simulation.py
More file actions
347 lines (290 loc) · 11.5 KB
/
Copy pathn_body_simulation.py
File metadata and controls
347 lines (290 loc) · 11.5 KB
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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
"""
In physics and astronomy, a gravitational N-body simulation is a simulation of a
dynamical system of particles under the influence of gravity. The system
consists of a number of bodies, each of which exerts a gravitational force on all
other bodies. These forces are calculated using Newton's law of universal
gravitation. The Euler method is used at each time-step to calculate the change in
velocity and position brought about by these forces. Softening is used to prevent
numerical divergences when a particle comes too close to another (and the force
goes to infinity).
(Description adapted from https://en.wikipedia.org/wiki/N-body_simulation )
(See also http://www.shodor.org/refdesk/Resources/Algorithms/EulersMethod/ )
"""
from __future__ import annotations
import random
from matplotlib import animation
from matplotlib import pyplot as plt
# Frame rate of the animation
INTERVAL = 20
# Time between time steps in seconds
DELTA_TIME = INTERVAL / 1000
class Body:
def __init__(
self,
position_x: float,
position_y: float,
velocity_x: float,
velocity_y: float,
mass: float = 1.0,
size: float = 1.0,
color: str = "blue",
) -> None:
"""
The parameters "size" & "color" are not relevant for the simulation itself,
they are only used for plotting.
"""
self.position_x = position_x
self.position_y = position_y
self.velocity_x = velocity_x
self.velocity_y = velocity_y
self.mass = mass
self.size = size
self.color = color
@property
def position(self) -> tuple[float, float]:
return self.position_x, self.position_y
@property
def velocity(self) -> tuple[float, float]:
return self.velocity_x, self.velocity_y
def update_velocity(
self, force_x: float, force_y: float, delta_time: float
) -> None:
"""
Euler algorithm for velocity
>>> body_1 = Body(0.,0.,0.,0.)
>>> body_1.update_velocity(1.,0.,1.)
>>> body_1.velocity
(1.0, 0.0)
>>> body_1.update_velocity(1.,0.,1.)
>>> body_1.velocity
(2.0, 0.0)
>>> body_2 = Body(0.,0.,5.,0.)
>>> body_2.update_velocity(0.,-10.,10.)
>>> body_2.velocity
(5.0, -100.0)
>>> body_2.update_velocity(0.,-10.,10.)
>>> body_2.velocity
(5.0, -200.0)
"""
self.velocity_x += force_x * delta_time
self.velocity_y += force_y * delta_time
def update_position(self, delta_time: float) -> None:
"""
Euler algorithm for position
>>> body_1 = Body(0.,0.,1.,0.)
>>> body_1.update_position(1.)
>>> body_1.position
(1.0, 0.0)
>>> body_1.update_position(1.)
>>> body_1.position
(2.0, 0.0)
>>> body_2 = Body(10.,10.,0.,-2.)
>>> body_2.update_position(1.)
>>> body_2.position
(10.0, 8.0)
>>> body_2.update_position(1.)
>>> body_2.position
(10.0, 6.0)
"""
self.position_x += self.velocity_x * delta_time
self.position_y += self.velocity_y * delta_time
class BodySystem:
"""
This class is used to hold the bodies, the gravitation constant, the time
factor and the softening factor. The time factor is used to control the speed
of the simulation. The softening factor is used for softening, a numerical
trick for N-body simulations to prevent numerical divergences when two bodies
get too close to each other.
"""
def __init__(
self,
bodies: list[Body],
gravitation_constant: float = 1.0,
time_factor: float = 1.0,
softening_factor: float = 0.0,
) -> None:
self.bodies = bodies
self.gravitation_constant = gravitation_constant
self.time_factor = time_factor
self.softening_factor = softening_factor
def __len__(self) -> int:
return len(self.bodies)
def update_system(self, delta_time: float) -> None:
"""
For each body, loop through all other bodies to calculate the total
force they exert on it. Use that force to update the body's velocity.
>>> body_system_1 = BodySystem([Body(0,0,0,0), Body(10,0,0,0)])
>>> len(body_system_1)
2
>>> body_system_1.update_system(1)
>>> body_system_1.bodies[0].position
(0.01, 0.0)
>>> body_system_1.bodies[0].velocity
(0.01, 0.0)
>>> body_system_2 = BodySystem([Body(-10,0,0,0), Body(10,0,0,0, mass=4)], 1, 10)
>>> body_system_2.update_system(1)
>>> body_system_2.bodies[0].position
(-9.0, 0.0)
>>> body_system_2.bodies[0].velocity
(0.1, 0.0)
"""
for body1 in self.bodies:
force_x = 0.0
force_y = 0.0
for body2 in self.bodies:
if body1 != body2:
dif_x = body2.position_x - body1.position_x
dif_y = body2.position_y - body1.position_y
# Calculation of the distance using Pythagoras's theorem
# Extra factor due to the softening technique
distance = (dif_x**2 + dif_y**2 + self.softening_factor) ** (1 / 2)
# Newton's law of universal gravitation.
force_x += (
self.gravitation_constant * body2.mass * dif_x / distance**3
)
force_y += (
self.gravitation_constant * body2.mass * dif_y / distance**3
)
# Update the body's velocity once all the force components have been added
body1.update_velocity(force_x, force_y, delta_time * self.time_factor)
# Update the positions only after all the velocities have been updated
for body in self.bodies:
body.update_position(delta_time * self.time_factor)
def update_step(
body_system: BodySystem, delta_time: float, patches: list[plt.Circle]
) -> None:
"""
Updates the body-system and applies the change to the patch-list used for plotting
>>> body_system_1 = BodySystem([Body(0,0,0,0), Body(10,0,0,0)])
>>> patches_1 = [plt.Circle((body.position_x, body.position_y), body.size,
... fc=body.color)for body in body_system_1.bodies] #doctest: +ELLIPSIS
>>> update_step(body_system_1, 1, patches_1)
>>> patches_1[0].center
(0.01, 0.0)
>>> body_system_2 = BodySystem([Body(-10,0,0,0), Body(10,0,0,0, mass=4)], 1, 10)
>>> patches_2 = [plt.Circle((body.position_x, body.position_y), body.size,
... fc=body.color)for body in body_system_2.bodies] #doctest: +ELLIPSIS
>>> update_step(body_system_2, 1, patches_2)
>>> patches_2[0].center
(-9.0, 0.0)
"""
# Update the positions of the bodies
body_system.update_system(delta_time)
# Update the positions of the patches
for patch, body in zip(patches, body_system.bodies):
patch.center = (body.position_x, body.position_y)
def plot(
title: str,
body_system: BodySystem,
x_start: float = -1,
x_end: float = 1,
y_start: float = -1,
y_end: float = 1,
) -> None:
"""
Utility function to plot how the given body-system evolves over time.
No doctest provided since this function does not have a return value.
"""
fig = plt.figure()
fig.canvas.manager.set_window_title(title)
ax = plt.axes(
xlim=(x_start, x_end), ylim=(y_start, y_end)
) # Set section to be plotted
plt.gca().set_aspect("equal") # Fix aspect ratio
# Each body is drawn as a patch by the plt-function
patches = [
plt.Circle((body.position_x, body.position_y), body.size, fc=body.color)
for body in body_system.bodies
]
for patch in patches:
ax.add_patch(patch)
# Function called at each step of the animation
def update(frame: int) -> list[plt.Circle]: # noqa: ARG001
update_step(body_system, DELTA_TIME, patches)
return patches
anim = animation.FuncAnimation( # noqa: F841
fig, update, interval=INTERVAL, blit=True
)
plt.show()
def example_1() -> BodySystem:
"""
Example 1: figure-8 solution to the 3-body-problem
This example can be seen as a test of the implementation: given the right
initial conditions, the bodies should move in a figure-8.
(initial conditions taken from http://www.artcompsci.org/vol_1/v1_web/node56.html)
>>> body_system = example_1()
>>> len(body_system)
3
"""
position_x = 0.9700436
position_y = -0.24308753
velocity_x = 0.466203685
velocity_y = 0.43236573
bodies1 = [
Body(position_x, position_y, velocity_x, velocity_y, size=0.2, color="red"),
Body(-position_x, -position_y, velocity_x, velocity_y, size=0.2, color="green"),
Body(0, 0, -2 * velocity_x, -2 * velocity_y, size=0.2, color="blue"),
]
return BodySystem(bodies1, time_factor=3)
def example_2() -> BodySystem:
"""
Example 2: Moon's orbit around the earth
This example can be seen as a test of the implementation: given the right
initial conditions, the moon should orbit around the earth as it actually does.
(mass, velocity and distance taken from https://en.wikipedia.org/wiki/Earth
and https://en.wikipedia.org/wiki/Moon)
No doctest provided since this function does not have a return value.
"""
moon_mass = 7.3476e22
earth_mass = 5.972e24
velocity_dif = 1022
earth_moon_distance = 384399000
gravitation_constant = 6.674e-11
# Calculation of the respective velocities so that total impulse is zero,
# i.e. the two bodies together don't move
moon_velocity = earth_mass * velocity_dif / (earth_mass + moon_mass)
earth_velocity = moon_velocity - velocity_dif
moon = Body(-earth_moon_distance, 0, 0, moon_velocity, moon_mass, 10000000, "grey")
earth = Body(0, 0, 0, earth_velocity, earth_mass, 50000000, "blue")
return BodySystem([earth, moon], gravitation_constant, time_factor=1000000)
def example_3() -> BodySystem:
"""
Example 3: Random system with many bodies.
No doctest provided since this function does not have a return value.
"""
bodies = []
for _ in range(10):
velocity_x = random.uniform(-0.5, 0.5)
velocity_y = random.uniform(-0.5, 0.5)
# Bodies are created pairwise with opposite velocities so that the
# total impulse remains zero
bodies.append(
Body(
random.uniform(-0.5, 0.5),
random.uniform(-0.5, 0.5),
velocity_x,
velocity_y,
size=0.05,
)
)
bodies.append(
Body(
random.uniform(-0.5, 0.5),
random.uniform(-0.5, 0.5),
-velocity_x,
-velocity_y,
size=0.05,
)
)
return BodySystem(bodies, 0.01, 10, 0.1)
if __name__ == "__main__":
plot("Figure-8 solution to the 3-body-problem", example_1(), -2, 2, -2, 2)
plot(
"Moon's orbit around the earth",
example_2(),
-430000000,
430000000,
-430000000,
430000000,
)
plot("Random system with many bodies", example_3(), -1.5, 1.5, -1.5, 1.5)
You can’t perform that action at this time.