Python/physics/n_body_simulation.py

351 lines
12 KiB
Python
Raw Normal View History

"""
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 )
2022-10-16 13:45:25 +00:00
(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]:
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.
2022-10-16 13:45:25 +00:00
(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 = []
Add flake8 pluin flake8 bugbear to pre-commit (#7132) * ci(pre-commit): Add ``flake8-builtins`` additional dependency to ``pre-commit`` (#7104) * refactor: Fix ``flake8-builtins`` (#7104) * fix(lru_cache): Fix naming conventions in docstrings (#7104) * ci(pre-commit): Order additional dependencies alphabetically (#7104) * fix(lfu_cache): Correct function name in docstring (#7104) * Update strings/snake_case_to_camel_pascal_case.py Co-authored-by: Christian Clauss <cclauss@me.com> * Update data_structures/stacks/next_greater_element.py Co-authored-by: Christian Clauss <cclauss@me.com> * Update digital_image_processing/index_calculation.py Co-authored-by: Christian Clauss <cclauss@me.com> * Update graphs/prim.py Co-authored-by: Christian Clauss <cclauss@me.com> * Update hashes/djb2.py Co-authored-by: Christian Clauss <cclauss@me.com> * refactor: Rename `_builtin` to `builtin_` ( #7104) * fix: Rename all instances (#7104) * refactor: Update variable names (#7104) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * ci: Create ``tox.ini`` and ignore ``A003`` (#7123) * revert: Remove function name changes (#7104) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Rename tox.ini to .flake8 * Update data_structures/heap/heap.py Co-authored-by: Dhruv Manilawala <dhruvmanila@gmail.com> * refactor: Rename `next_` to `next_item` (#7104) * ci(pre-commit): Add `flake8` plugin `flake8-bugbear` (#7127) * refactor: Follow `flake8-bugbear` plugin (#7127) * fix: Correct `knapsack` code (#7127) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Co-authored-by: Christian Clauss <cclauss@me.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Dhruv Manilawala <dhruvmanila@gmail.com>
2022-10-13 16:03:06 +00:00
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)