Clean up and comment

physics.py

@@ -1,41 +1,82 @@
+"""
+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
 
 
-G = 6.674e-11
+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 rotations(a: np.ndarray):
-	a2 = np.concatenate((a, a))
-	for i in range(1, len(a)):
-		yield a2[i: i + len(a)]
-
-
-def n_body(pos: np.ndarray, vel: np.ndarray, mass: np.ndarray):
-	for (o_pos, o_mass) in zip(rotations(pos), rotations(mass)):
-		dist = o_pos - pos
-		vel += G * dist * o_mass / (np.linalg.norm(dist, axis=1) ** 3)[:, np.newaxis]
+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: np.ndarray, vel: np.ndarray, mass: np.ndarray, constrain=2.):
-	n, d = pos.shape
-	dist = np.zeros((n - 1, n, d))
-	rot_mass = np.zeros((n - 1, n, 1))
+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)):
-		dist[i - 1] = pos2[i: i + n] - pos
-		rot_mass[i - 1] = mass2[i: i + n]
+		# 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
 
-	vel += G * np.sum(
+	# 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