82 lines
3.1 KiB
Python
82 lines
3.1 KiB
Python
"""
|
|
Simulation tick algorithms
|
|
|
|
Both algorithms cancel out some terms. As you know, the force of gravity is $\frac{G * m_1 * m_2}{r^2}$. However, this
|
|
force is applied in the direction of the vector between the two masses. Because this direction vector needs to be
|
|
normalized, we can combine the normalization with the above equation to get $\frac{G * m_1 * m_2 * (p_2 - p_1)}{r^3}$
|
|
and skip out on a costly square root to calculate $r$ again. Finally, because this force is applied to one of the
|
|
masses (say $m_1$), the actual change in velocity is the force divided by the mass. This means that we can just drop
|
|
the $m_1$ term from the equation, and we have our change in velocity.
|
|
"""
|
|
from typing import Any, Generator
|
|
|
|
import numpy as np
|
|
from numpy.typing import NDArray
|
|
|
|
|
|
def _rotations(arr: NDArray) -> Generator[NDArray, Any, None]:
|
|
"""
|
|
"Rotates" through an array, returning the [i: n+i] range on the ith iteration
|
|
:param arr: Array to rotate through
|
|
"""
|
|
a2 = np.concatenate((arr, arr))
|
|
for i in range(1, len(arr)):
|
|
yield a2[i: i + len(arr)]
|
|
|
|
|
|
def n_body(pos: NDArray, vel: NDArray, mass: NDArray, gravity: float) -> None:
|
|
"""
|
|
Easier-to-understand but slower update algorithm that simulates a tick.
|
|
Unused, just for demonstration purposes.
|
|
:param pos: Previous positions
|
|
:param vel: Previous velocities
|
|
:param mass: Object masses
|
|
:param gravity: Simulation gravity
|
|
:return: None - updated in-place
|
|
"""
|
|
for (other_pos, other_mass) in zip(_rotations(pos), _rotations(mass)):
|
|
# For each combination of 2 objects
|
|
dist = other_pos - pos
|
|
# Calculate the force of gravity from the first to the second, and use it to update the velocity
|
|
vel += gravity * dist * other_mass / (np.linalg.norm(dist, axis=1) ** 3)[:, np.newaxis]
|
|
# Update positions
|
|
pos += vel
|
|
|
|
|
|
def n_body_matrix(pos: NDArray, vel: NDArray, mass: NDArray, gravity: float, constrain: float = 2.) -> None:
|
|
"""
|
|
Harder-to-understand but faster update algorithm that simulates a tick.
|
|
Basically does the simpler algorithm all at once using numpy parallelism.
|
|
:param pos: Previous positions
|
|
:param vel: Previous velocities
|
|
:param mass: Object masses
|
|
:param gravity: Simulation gravity
|
|
:param constrain: Numerical stability term
|
|
:return: None - updated in-place
|
|
"""
|
|
num_masses, dimensions = pos.shape
|
|
dist = np.zeros((num_masses - 1, num_masses, dimensions), dtype=np.float64)
|
|
rot_mass = np.zeros((num_masses - 1, num_masses, 1), dtype=np.float64)
|
|
|
|
pos2 = np.concatenate((pos, pos))
|
|
mass2 = np.concatenate((mass, mass))
|
|
# Generates a matrix using the rotated arrays, like the for loop in the above algorithm in one go
|
|
for i in range(1, len(pos)):
|
|
# The distance between two objects
|
|
dist[i - 1] = pos2[i: i + num_masses] - pos
|
|
# The mass of the other object
|
|
rot_mass[i - 1] = mass2[i: i + num_masses]
|
|
|
|
# Normalize direction vectors, and ensure the distances aren't too close for numerical stability
|
|
norms = np.linalg.norm(dist, axis=2)
|
|
if constrain:
|
|
norms[norms < constrain] = constrain
|
|
|
|
# Calculate all changes in velocity at once, using the same method described and implemented above
|
|
vel += gravity * np.sum(
|
|
dist * rot_mass / (norms ** 3)[:, :, np.newaxis],
|
|
axis=0
|
|
)
|
|
|
|
# Update positions
|
|
pos += vel
|