MichaelBradley/nbody
1
0
Fork 0
Code Issues Pull requests Activity Actions

Compare commits

base: MichaelBradley:6822df04327bd6aa40aaae45b8af83012580a391
Branches Tags
MichaelBradley:main
..
compare: MichaelBradley:9651d1470a0034305c8dcd29d22351aa137b59b9
Branches Tags
MichaelBradley:main

No commits in common. "6822df04327bd6aa40aaae45b8af83012580a391" and "9651d1470a0034305c8dcd29d22351aa137b59b9" have entirely different histories.
6822df0432 ... 9651d1470a

6 changed files with 192 additions and 289 deletions
Download patch file Download diff file Expand all files Collapse all files

11
README.md
View file

@ -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:

211
data.py
View file

@ -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()

84
gen_data.py
View file

@ -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 <count> 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)}")

54
main.py
View file

@ -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()

101
physics.py
View file

@ -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

20
requirements.txt
View file

@ -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