diff --git a/README.md b/README.md index 9e0e0d9..22c4f2d 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,16 @@ # nbody -Threw this together to get more comfortable with Numpy. -Can simulate a few hundred bodies in 2 or 3 dimensions without much hassle. +Threw this together in a day or two cause I thought it would be fun to mess around with. +Can simulate a few hundred bodies in 2 or 3 dimensions without much hassle (hardware dependent of course). + +Comments are non-existent, sorry about that. + To set up: ```shell python -m venv venv source venv/bin/activate -pip install -r requirements.txt +pip -r requirements.txt ``` To run: @@ -16,7 +19,7 @@ To run: ./main.py -f 2d/simple.csv # 3d simulation, increased gravity -./main.py -f 3d/some.csv -g 30 +./main.py -f 3d/some.csv -d 3 -g 30 ``` To create a new start state: diff --git a/data.py b/data.py index a382516..4c2df2a 100644 --- a/data.py +++ b/data.py @@ -1,148 +1,93 @@ -from pathlib import Path - import matplotlib.animation as animation import matplotlib.cm as cm import matplotlib.pyplot as plt import numpy as np -from matplotlib.collections import PathCollection -from numpy.typing import NDArray - -from physics import n_body_matrix +import physics -def parse_csv(file: Path) -> tuple[NDArray, NDArray, NDArray]: - """ - Reads the starting conditions of a simulation from a CSV - :param file: The CSV file - :return: The position, velocity, and radius matrices, in 2 or 3 dimensions - """ - # Get text from file - lines = file.read_text().strip().splitlines() - field_count = len(lines[0].split(",")) - if field_count not in (5, 7): - raise RuntimeError("CSV format not recognized, can only show scenes in 2 or 3 dimensions") - - # Allocate matrices - dimensions = (field_count - 1) // 2 - pos = np.zeros((len(lines), dimensions), dtype=np.float64) - vel = np.zeros((len(lines), dimensions), dtype=np.float64) - rad = np.zeros((len(lines), 1), dtype=np.float64) - - # Parse CSV lines - for i, line in enumerate(lines): - values = tuple(float(field) for field in line.split(",")) - if dimensions == 2: - x, y, vx, vy, r = values - pos[i] = (x, y) - vel[i] = (vx, vy) - rad[i] = r - elif dimensions == 3: - x, y, z, vx, vy, vz, r = values - pos[i] = (x, y, z) - vel[i] = (vx, vy, vz) - rad[i] = r - - return pos, vel, rad +def parse_csv(filename: str, dimensions=2): + if not (1 < dimensions < 4): + raise ValueError(f"Can only show 2or 3 dimensional scenes, not {dimensions}") + with open(filename, 'r') as file: + lines = file.read().strip().splitlines() + pos = np.zeros((len(lines), dimensions)) + vel = np.zeros((len(lines), dimensions)) + rad = np.zeros((len(lines), 1)) + for i, values in enumerate(map(lambda l: map(float, l.split(',')), lines)): + if dimensions == 2: + [x, y, vx, vy, r] = values + pos[i] = [x, y] + vel[i] = [vx, vy] + rad[i] = r + elif dimensions == 3: + [x, y, z, vx, vy, vz, r] = values + pos[i] = [x, y, z] + vel[i] = [vx, vy, vz] + rad[i] = r + return pos, vel, rad class Animator: - """Runs the simulation and displays it in a plot""" - def __init__(self, pos: NDArray, vel: NDArray, rad: NDArray, gravity: float) -> None: - """ - Sets up the simulation using the given data - :param pos: Start positions of the objects - :param vel: Start velocities of the objects - :param rad: Radii of the objects - :param gravity: Strength of gravity in this simulation - """ - # We update our arrays in-place, so we'll make out own copies to be safe - self._pos = pos.copy() - self._vel = vel.copy() - self._rad = rad.copy() - # Calculate volume of a sphere of this radius - self._mass = np.pi * 4 / 3 * rad ** 3 + def __init__(self, pos: np.ndarray, vel: np.ndarray, rad: np.ndarray): + self.pos = pos + self.vel = vel + self.rad = rad + self.mass = np.pi * 4 / 3 * rad ** 3 - self._gravity = gravity + n, d = self.pos.shape - object_count, dimensions = self._pos.shape + self.scat: plt.PathCollection = None + self.colours = cm.rainbow( + np.random.random( + (n,) + ) + ) - # Objects will be represented using a scatter plot - self._scatter_plot: plt.PathCollection | None = None - # We'll give each objects a random colour to better differentiate them - self._object_colours: NDArray = cm.rainbow( - np.random.random( - (object_count,) - ) - ) + self.fig = plt.figure() + if d == 2: + self.ax = self.fig.add_subplot() + else: + self.ax = self.fig.add_subplot(projection="3d") + self.ani = animation.FuncAnimation( + self.fig, + self.update, + interval=1000 / (15 * 2 ** 4), + init_func=self.setup_plot, + blit=True, + cache_frame_data=False + ) - # Create plot in an appropriate number of dimensions - self._fig = plt.figure() - if dimensions == 2: - self._ax = self._fig.add_subplot() - else: - self._ax = self._fig.add_subplot(projection="3d") + def setup_plot(self): + _n, d = self.pos.shape + if d == 2: + self.scat = self.ax.scatter( + self.pos[:, 0], + self.pos[:, 1], + c=self.colours, + s=self.rad * 10 + ) + self.ax.axis([-950, 950, -500, 500]) + else: + self.scat = self.ax.scatter( + self.pos[:, 0], + self.pos[:, 1], + self.pos[:, 2], + c=self.colours, + s=self.rad * 10, + depthshade=False + ) + self.ax.axis([-500, 500, -500, 500, -500, 500]) + return self.scat, - # Set up animation loop - # This attribute never gets used again, but we'll keep a reference so that it doesn't get garbage collected - self._animation = animation.FuncAnimation( - self._fig, - self.update, - interval=1000 / (15 * 2 ** 4), - init_func=self.setup_plot, - blit=True, - cache_frame_data=False, - ) + def update(self, *_args, **_kwargs): + _n, d = self.pos.shape + physics.n_body_matrix(self.pos, self.vel, self.mass, constrain=2.) + self.scat.set_offsets(self.pos[:, :2]) + if d == 3: + self.scat.set_3d_properties(self.pos[:, 2], 'z') + self.scat.set_sizes(self.rad[:, 0] * 10) + self.fig.canvas.draw() + return self.scat, - # Display the animation - plt.show() - - def setup_plot(self) -> tuple[PathCollection]: - """ - This is a FuncAnimation initialization function in the form matplotlib expects - :return: The single scatter plot we're using - """ - _n, dimensions = self._pos.shape - # Set up the scatter plot in 2 or 3 dimensions - if dimensions == 2: - self._scatter_plot = self._ax.scatter( - self._pos[:, 0], - self._pos[:, 1], - c=self._object_colours, - s=self._rad * 10, # To make the objects more visible - ) - # These values work nicely for a landscape window - self._ax.axis([-950, 950, -500, 500]) - else: - self._scatter_plot = self._ax.scatter( - self._pos[:, 0], - self._pos[:, 1], - self._pos[:, 2], - c=self._object_colours, - s=self._rad * 10, # To make the objects more visible - depthshade=False, # I find it confusing, YMMV - ) - # These values work nicely for a square window - self._ax.axis([-500, 500, -500, 500, -500, 500]) - return self._scatter_plot, - - def update(self, *_args, **_kwargs) -> tuple[PathCollection]: - """ - This is a FuncAnimation update function in the form matplotlib expects - :arg _args: We don't need any of matplotlib's information for our implementation - "arg _kwargs: Again, not necessary for us - :return: The single scatter plot we're using - """ - _n, dimensions = self._pos.shape - # Update the state of our simulation - n_body_matrix(self._pos, self._vel, self._mass, self._gravity) - - # Set the X and Y values of the objects - self._scatter_plot.set_offsets(self._pos[:, :2]) - if dimensions == 3: - # Update the Z value if in 3D - self._scatter_plot.set_3d_properties(self._pos[:, 2], 'z') - # Use radius to represent mass - self._scatter_plot.set_sizes(self._rad[:, 0] * 10) - # Redraw the plot - self._fig.canvas.draw() - return self._scatter_plot, + def show(self): + plt.show() diff --git a/gen_data.py b/gen_data.py index fc01fc6..1d1e32d 100755 --- a/gen_data.py +++ b/gen_data.py @@ -1,62 +1,38 @@ -#!venv/bin/python -"""Generates random CSV data to be read by the simulator""" - -from argparse import ArgumentParser -from random import uniform -from typing import Any, cast +#!./venv/bin/python +import argparse +from random import uniform, randint class Args: - """ - The types of the arguments retrieved from the user - """ - width: int - length: int - depth: int - speed: float - radius: float - count: int - - -def print_part(data: Any) -> None: - """ - Prints a CSV field - :param data: The data to put in the field - """ - print(str(data), end=",") - - -def main(args: Args) -> None: - """ - Generates the starting data - :param args: Parameters for the random generation - """ - # Print objects - for _ in range(args.count): - # Object location - print_part(uniform(-args.width / 2, args.width / 2)) - print_part(uniform(-args.length / 2, args.length / 2)) - if args.depth: - print_part(uniform(-args.depth / 2, args.depth / 2)) - - # Object velocity - print_part(uniform(-args.speed, args.speed)) - print_part(uniform(-args.speed, args.speed)) - if args.depth: - print_part(uniform(-args.speed, args.speed)) - - # Finish line with a positive radius - print(uniform(1e-2, args.radius)) + width: int + height: int + depth: int + velocity: float + mass: float + count: int if __name__ == "__main__": - parser = ArgumentParser(description="Generates data for the n-body simulator", epilog="You should redirect the output to a file") + parser = argparse.ArgumentParser( + prog="n-body data generator", + description="Generates data for the n-body simulator.", + add_help=False + ) - parser.add_argument("-w", "--width", type=float, default=1900., help="The width of the spawning area") - parser.add_argument("-l", "--length", type=float, default=1000., help="The length of the spawning area") - parser.add_argument("-d", "--depth", type=float, default=0., help="The depth of the spawning area, where 0 implies only 2 dimensions") - parser.add_argument("-s", "--speed", type=float, default=1., help="The maximum initial starting speed of an object in any dimension") - parser.add_argument("-r", "--radius", type=float, default=1., help="The maximum radius of an object") - parser.add_argument("-c", "--count", type=int, default=500, help="How many objects to create") + parser.add_argument("-w", "--width", type=int, default=1900) + parser.add_argument("-h", "--height", type=int, default=1000) + parser.add_argument("-d", "--depth", type=int, default=0) + parser.add_argument("-v", "--velocity", type=float, default=1.) + parser.add_argument("-m", "--mass", type=float, default=1.) + parser.add_argument("-c", "--count", type=int, default=500) - main(cast(Args, parser.parse_args())) + args: Args = parser.parse_args() + + for _ in range(args.count): + print(f"{randint(-args.width // 2, args.width // 2)}," + f"{randint(-args.height // 2, args.height // 2)}," + f"{f'{randint(-args.depth // 2, args.depth // 2)},' if args.depth else ''}" + f"{uniform(-args.velocity, args.velocity)}," + f"{uniform(-args.velocity, args.velocity)}," + f"{f'{uniform(-args.velocity, args.velocity)},' if args.depth else ''}" + f"{uniform(1e-2, args.mass)}") diff --git a/main.py b/main.py index a9b974e..4e9cc2a 100755 --- a/main.py +++ b/main.py @@ -1,26 +1,46 @@ -#!venv/bin/python -from argparse import ArgumentParser -from pathlib import Path -from typing import cast +#!./venv/bin/python +import argparse +import typing -from data import parse_csv, Animator +import data +import physics class Args: - file: Path - gravity: float - - -def main(file: Path, gravity: float) -> None: - pos, vel, rad = parse_csv(file) - Animator(pos, vel, rad, gravity) + filename: str + gravity: float + dimensions: typing.Literal[2, 3] if __name__ == "__main__": - parser = ArgumentParser(description="Gravitation simulation") + parser = argparse.ArgumentParser( + prog="n-body simulation", + description="Simulating gravitational effects" + ) - parser.add_argument("-f", "--file", type=Path, default="data/2d/simple.csv", help="The starting state of the simulation") - parser.add_argument("-g", "--gravity", type=float, default=1., help="The strength of gravity") + parser.add_argument( + "-f", + "--filename", + default="data/2d/simple.csv" + ) + parser.add_argument( + "-g", + "--gravity", + type=float, + default=1. + ) + parser.add_argument( + "-d", + "--dimensions", + type=int, + choices=[2, 3], + default=2 + ) - args = cast(Args, parser.parse_args()) - main(args.file, args.gravity) + args: Args = parser.parse_args() + + physics.G = args.gravity + + objects = data.parse_csv(args.filename, dimensions=args.dimensions) + a = data.Animator(*objects) + a.show() diff --git a/physics.py b/physics.py index a26f7bb..4965775 100644 --- a/physics.py +++ b/physics.py @@ -1,82 +1,41 @@ -""" -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)] +G = 6.674e-11 -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 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_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) +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] + pos += vel - 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 +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)) - # 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 - ) + pos2 = np.concatenate((pos, pos)) + mass2 = np.concatenate((mass, mass)) - # Update positions - pos += vel + for i in range(1, len(pos)): + dist[i - 1] = pos2[i: i + n] - pos + rot_mass[i - 1] = mass2[i: i + n] + + norms = np.linalg.norm(dist, axis=2) + if constrain: + norms[norms < constrain] = constrain + + vel += G * np.sum( + dist * rot_mass / (norms ** 3)[:, :, np.newaxis], + axis=0 + ) + + pos += vel diff --git a/requirements.txt b/requirements.txt index e9094c5..2e5db94 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,11 +1,11 @@ -contourpy==1.3.1 +contourpy==1.1.1 cycler==0.12.1 -fonttools==4.55.3 -kiwisolver==1.4.8 -matplotlib==3.10.0 -numpy==2.2.1 -packaging==24.2 -pillow==11.0.0 -pyparsing==3.2.1 -python-dateutil==2.9.0.post0 -six==1.17.0 +fonttools==4.43.1 +kiwisolver==1.4.5 +matplotlib==3.8.0 +numpy==1.26.0 +packaging==23.2 +Pillow==10.0.1 +pyparsing==3.1.1 +python-dateutil==2.8.2 +six==1.16.0