Skip to content

Commit d14de5a

Browse files
authored
update n_body_simulation.py
1 parent e7809db commit d14de5a

File tree

1 file changed

+15
-3
lines changed

1 file changed

+15
-3
lines changed

other/n_body_simulation.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,11 @@ def plot(
139139
y_start: float = -1,
140140
y_end: float = 1,
141141
) -> None:
142+
"""
143+
Utility function to plot how the given body-system evolves over time.
144+
No doctest provided since this function does not have a return value.
145+
"""
146+
142147
INTERVAL = 20 # Frame rate of the animation
143148
DELTA_TIME = INTERVAL / 1000 # Time between time steps in seconds
144149

@@ -156,7 +161,7 @@ def plot(
156161
)
157162

158163
# Function called once at the start of the animation
159-
def init() -> list[patches.Circle]: # type: ignore
164+
def init() -> list[patches.Circle]:
160165
axes = plt.gca()
161166
axes.set_aspect("equal")
162167

@@ -165,7 +170,7 @@ def init() -> list[patches.Circle]: # type: ignore
165170
return patches
166171

167172
# Function called at each step of the animation
168-
def animate(i: int) -> list[patches.Circle]: # type: ignore
173+
def update(frame: int) -> list[patches.Circle]:
169174
# Update the positions of the bodies
170175
body_system.update_system(DELTA_TIME)
171176

@@ -175,14 +180,17 @@ def animate(i: int) -> list[patches.Circle]: # type: ignore
175180
return patches
176181

177182
anim = animation.FuncAnimation( # noqa: F841
178-
fig, animate, init_func=init, interval=INTERVAL, blit=True
183+
fig, update, init_func=init, interval=INTERVAL, blit=True
179184
)
180185

181186
plt.show()
182187

183188

184189
if __name__ == "__main__":
185190
# Example 1: figure-8 solution to the 3-body-problem
191+
# This example can be seen as a test of the implementation: given the right
192+
# initial conditions, the bodies should move in a figure-8.
193+
# (initial conditions taken from http://www.artcompsci.org/vol_1/v1_web/node56.html)
186194
position_x = 0.9700436
187195
position_y = -0.24308753
188196
velocity_x = 0.466203685
@@ -197,6 +205,10 @@ def animate(i: int) -> list[patches.Circle]: # type: ignore
197205
plot("Figure-8 solution to the 3-body-problem", body_system1, -2, 2, -2, 2)
198206

199207
# Example 2: Moon's orbit around the earth
208+
# This example can be seen as a test of the implementation: given the right
209+
# initial conditions, the moon should orbit around the earth as it actually does.
210+
# (mass, velocity and distance taken from https://en.wikipedia.org/wiki/Earth
211+
# and https://en.wikipedia.org/wiki/Moon)
200212
moon_mass = 7.3476e22
201213
earth_mass = 5.972e24
202214
velocity_dif = 1022

0 commit comments

Comments
 (0)