Switch to tabs
This commit is contained in:
parent
054b54c291
commit
8bb6ba19d5
SSH key fingerprint:
SHA256:cj/YZ5VT+QOKncqSkx+ibKTIn0Obg7OIzwzl9BL8EO8
5 changed files with 157 additions and 157 deletions
|
|
@ -1,7 +1,7 @@
|
||||||
# nbody
|
# nbody
|
||||||
|
|
||||||
Threw this together in a day or two cause I thought it would be fun to mess around with.
|
Threw this together to get more comfortable with Numpy.
|
||||||
Can simulate a few hundred bodies in 2 or 3 dimensions without much hassle (hardware dependent of course).
|
Can simulate a few hundred bodies in 2 or 3 dimensions without much hassle.
|
||||||
|
|
||||||
Comments are non-existent, sorry about that.
|
Comments are non-existent, sorry about that.
|
||||||
|
|
||||||
|
|
|
||||||
152
data.py
152
data.py
|
|
@ -6,88 +6,88 @@ import physics
|
||||||
|
|
||||||
|
|
||||||
def parse_csv(filename: str, dimensions=2):
|
def parse_csv(filename: str, dimensions=2):
|
||||||
if not (1 < dimensions < 4):
|
if not (1 < dimensions < 4):
|
||||||
raise ValueError(f"Can only show 2or 3 dimensional scenes, not {dimensions}")
|
raise ValueError(f"Can only show 2or 3 dimensional scenes, not {dimensions}")
|
||||||
with open(filename, 'r') as file:
|
with open(filename, 'r') as file:
|
||||||
lines = file.read().strip().splitlines()
|
lines = file.read().strip().splitlines()
|
||||||
pos = np.zeros((len(lines), dimensions))
|
pos = np.zeros((len(lines), dimensions))
|
||||||
vel = np.zeros((len(lines), dimensions))
|
vel = np.zeros((len(lines), dimensions))
|
||||||
rad = np.zeros((len(lines), 1))
|
rad = np.zeros((len(lines), 1))
|
||||||
for i, values in enumerate(map(lambda l: map(float, l.split(',')), lines)):
|
for i, values in enumerate(map(lambda l: map(float, l.split(',')), lines)):
|
||||||
if dimensions == 2:
|
if dimensions == 2:
|
||||||
[x, y, vx, vy, r] = values
|
[x, y, vx, vy, r] = values
|
||||||
pos[i] = [x, y]
|
pos[i] = [x, y]
|
||||||
vel[i] = [vx, vy]
|
vel[i] = [vx, vy]
|
||||||
rad[i] = r
|
rad[i] = r
|
||||||
elif dimensions == 3:
|
elif dimensions == 3:
|
||||||
[x, y, z, vx, vy, vz, r] = values
|
[x, y, z, vx, vy, vz, r] = values
|
||||||
pos[i] = [x, y, z]
|
pos[i] = [x, y, z]
|
||||||
vel[i] = [vx, vy, vz]
|
vel[i] = [vx, vy, vz]
|
||||||
rad[i] = r
|
rad[i] = r
|
||||||
return pos, vel, rad
|
return pos, vel, rad
|
||||||
|
|
||||||
|
|
||||||
class Animator:
|
class Animator:
|
||||||
def __init__(self, pos: np.ndarray, vel: np.ndarray, rad: np.ndarray):
|
def __init__(self, pos: np.ndarray, vel: np.ndarray, rad: np.ndarray):
|
||||||
self.pos = pos
|
self.pos = pos
|
||||||
self.vel = vel
|
self.vel = vel
|
||||||
self.rad = rad
|
self.rad = rad
|
||||||
self.mass = np.pi * 4 / 3 * rad ** 3
|
self.mass = np.pi * 4 / 3 * rad ** 3
|
||||||
|
|
||||||
n, d = self.pos.shape
|
n, d = self.pos.shape
|
||||||
|
|
||||||
self.scat: plt.PathCollection = None
|
self.scat: plt.PathCollection = None
|
||||||
self.colours = cm.rainbow(
|
self.colours = cm.rainbow(
|
||||||
np.random.random(
|
np.random.random(
|
||||||
(n,)
|
(n,)
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
self.fig = plt.figure()
|
self.fig = plt.figure()
|
||||||
if d == 2:
|
if d == 2:
|
||||||
self.ax = self.fig.add_subplot()
|
self.ax = self.fig.add_subplot()
|
||||||
else:
|
else:
|
||||||
self.ax = self.fig.add_subplot(projection="3d")
|
self.ax = self.fig.add_subplot(projection="3d")
|
||||||
self.ani = animation.FuncAnimation(
|
self.ani = animation.FuncAnimation(
|
||||||
self.fig,
|
self.fig,
|
||||||
self.update,
|
self.update,
|
||||||
interval=1000 / (15 * 2 ** 4),
|
interval=1000 / (15 * 2 ** 4),
|
||||||
init_func=self.setup_plot,
|
init_func=self.setup_plot,
|
||||||
blit=True,
|
blit=True,
|
||||||
cache_frame_data=False
|
cache_frame_data=False
|
||||||
)
|
)
|
||||||
|
|
||||||
def setup_plot(self):
|
def setup_plot(self):
|
||||||
_n, d = self.pos.shape
|
_n, d = self.pos.shape
|
||||||
if d == 2:
|
if d == 2:
|
||||||
self.scat = self.ax.scatter(
|
self.scat = self.ax.scatter(
|
||||||
self.pos[:, 0],
|
self.pos[:, 0],
|
||||||
self.pos[:, 1],
|
self.pos[:, 1],
|
||||||
c=self.colours,
|
c=self.colours,
|
||||||
s=self.rad * 10
|
s=self.rad * 10
|
||||||
)
|
)
|
||||||
self.ax.axis([-950, 950, -500, 500])
|
self.ax.axis([-950, 950, -500, 500])
|
||||||
else:
|
else:
|
||||||
self.scat = self.ax.scatter(
|
self.scat = self.ax.scatter(
|
||||||
self.pos[:, 0],
|
self.pos[:, 0],
|
||||||
self.pos[:, 1],
|
self.pos[:, 1],
|
||||||
self.pos[:, 2],
|
self.pos[:, 2],
|
||||||
c=self.colours,
|
c=self.colours,
|
||||||
s=self.rad * 10,
|
s=self.rad * 10,
|
||||||
depthshade=False
|
depthshade=False
|
||||||
)
|
)
|
||||||
self.ax.axis([-500, 500, -500, 500, -500, 500])
|
self.ax.axis([-500, 500, -500, 500, -500, 500])
|
||||||
return self.scat,
|
return self.scat,
|
||||||
|
|
||||||
def update(self, *_args, **_kwargs):
|
def update(self, *_args, **_kwargs):
|
||||||
_n, d = self.pos.shape
|
_n, d = self.pos.shape
|
||||||
physics.n_body_matrix(self.pos, self.vel, self.mass, constrain=2.)
|
physics.n_body_matrix(self.pos, self.vel, self.mass, constrain=2.)
|
||||||
self.scat.set_offsets(self.pos[:, :2])
|
self.scat.set_offsets(self.pos[:, :2])
|
||||||
if d == 3:
|
if d == 3:
|
||||||
self.scat.set_3d_properties(self.pos[:, 2], 'z')
|
self.scat.set_3d_properties(self.pos[:, 2], 'z')
|
||||||
self.scat.set_sizes(self.rad[:, 0] * 10)
|
self.scat.set_sizes(self.rad[:, 0] * 10)
|
||||||
self.fig.canvas.draw()
|
self.fig.canvas.draw()
|
||||||
return self.scat,
|
return self.scat,
|
||||||
|
|
||||||
def show(self):
|
def show(self):
|
||||||
plt.show()
|
plt.show()
|
||||||
|
|
|
||||||
52
gen_data.py
52
gen_data.py
|
|
@ -4,35 +4,35 @@ from random import uniform, randint
|
||||||
|
|
||||||
|
|
||||||
class Args:
|
class Args:
|
||||||
width: int
|
width: int
|
||||||
height: int
|
height: int
|
||||||
depth: int
|
depth: int
|
||||||
velocity: float
|
velocity: float
|
||||||
mass: float
|
mass: float
|
||||||
count: int
|
count: int
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
parser = argparse.ArgumentParser(
|
parser = argparse.ArgumentParser(
|
||||||
prog="n-body data generator",
|
prog="n-body data generator",
|
||||||
description="Generates data for the n-body simulator.",
|
description="Generates data for the n-body simulator.",
|
||||||
add_help=False
|
add_help=False
|
||||||
)
|
)
|
||||||
|
|
||||||
parser.add_argument("-w", "--width", type=int, default=1900)
|
parser.add_argument("-w", "--width", type=int, default=1900)
|
||||||
parser.add_argument("-h", "--height", type=int, default=1000)
|
parser.add_argument("-h", "--height", type=int, default=1000)
|
||||||
parser.add_argument("-d", "--depth", type=int, default=0)
|
parser.add_argument("-d", "--depth", type=int, default=0)
|
||||||
parser.add_argument("-v", "--velocity", type=float, default=1.)
|
parser.add_argument("-v", "--velocity", type=float, default=1.)
|
||||||
parser.add_argument("-m", "--mass", type=float, default=1.)
|
parser.add_argument("-m", "--mass", type=float, default=1.)
|
||||||
parser.add_argument("-c", "--count", type=int, default=500)
|
parser.add_argument("-c", "--count", type=int, default=500)
|
||||||
|
|
||||||
args: Args = parser.parse_args()
|
args: Args = parser.parse_args()
|
||||||
|
|
||||||
for _ in range(args.count):
|
for _ in range(args.count):
|
||||||
print(f"{randint(-args.width // 2, args.width // 2)},"
|
print(f"{randint(-args.width // 2, args.width // 2)},"
|
||||||
f"{randint(-args.height // 2, args.height // 2)},"
|
f"{randint(-args.height // 2, args.height // 2)},"
|
||||||
f"{f'{randint(-args.depth // 2, args.depth // 2)},' if args.depth else ''}"
|
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"{uniform(-args.velocity, args.velocity)},"
|
f"{uniform(-args.velocity, args.velocity)},"
|
||||||
f"{f'{uniform(-args.velocity, args.velocity)},' if args.depth else ''}"
|
f"{f'{uniform(-args.velocity, args.velocity)},' if args.depth else ''}"
|
||||||
f"{uniform(1e-2, args.mass)}")
|
f"{uniform(1e-2, args.mass)}")
|
||||||
|
|
|
||||||
60
main.py
60
main.py
|
|
@ -7,40 +7,40 @@ import physics
|
||||||
|
|
||||||
|
|
||||||
class Args:
|
class Args:
|
||||||
filename: str
|
filename: str
|
||||||
gravity: float
|
gravity: float
|
||||||
dimensions: typing.Literal[2, 3]
|
dimensions: typing.Literal[2, 3]
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
parser = argparse.ArgumentParser(
|
parser = argparse.ArgumentParser(
|
||||||
prog="n-body simulation",
|
prog="n-body simulation",
|
||||||
description="Simulating gravitational effects"
|
description="Simulating gravitational effects"
|
||||||
)
|
)
|
||||||
|
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"-f",
|
"-f",
|
||||||
"--filename",
|
"--filename",
|
||||||
default="data/2d/simple.csv"
|
default="data/2d/simple.csv"
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"-g",
|
"-g",
|
||||||
"--gravity",
|
"--gravity",
|
||||||
type=float,
|
type=float,
|
||||||
default=1.
|
default=1.
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"-d",
|
"-d",
|
||||||
"--dimensions",
|
"--dimensions",
|
||||||
type=int,
|
type=int,
|
||||||
choices=[2, 3],
|
choices=[2, 3],
|
||||||
default=2
|
default=2
|
||||||
)
|
)
|
||||||
|
|
||||||
args: Args = parser.parse_args()
|
args: Args = parser.parse_args()
|
||||||
|
|
||||||
physics.G = args.gravity
|
physics.G = args.gravity
|
||||||
|
|
||||||
objects = data.parse_csv(args.filename, dimensions=args.dimensions)
|
objects = data.parse_csv(args.filename, dimensions=args.dimensions)
|
||||||
a = data.Animator(*objects)
|
a = data.Animator(*objects)
|
||||||
a.show()
|
a.show()
|
||||||
|
|
|
||||||
46
physics.py
46
physics.py
|
|
@ -5,37 +5,37 @@ G = 6.674e-11
|
||||||
|
|
||||||
|
|
||||||
def rotations(a: np.ndarray):
|
def rotations(a: np.ndarray):
|
||||||
a2 = np.concatenate((a, a))
|
a2 = np.concatenate((a, a))
|
||||||
for i in range(1, len(a)):
|
for i in range(1, len(a)):
|
||||||
yield a2[i: i + len(a)]
|
yield a2[i: i + len(a)]
|
||||||
|
|
||||||
|
|
||||||
def n_body(pos: np.ndarray, vel: np.ndarray, mass: np.ndarray):
|
def n_body(pos: np.ndarray, vel: np.ndarray, mass: np.ndarray):
|
||||||
for (o_pos, o_mass) in zip(rotations(pos), rotations(mass)):
|
for (o_pos, o_mass) in zip(rotations(pos), rotations(mass)):
|
||||||
dist = o_pos - pos
|
dist = o_pos - pos
|
||||||
vel += G * dist * o_mass / (np.linalg.norm(dist, axis=1) ** 3)[:, np.newaxis]
|
vel += G * dist * o_mass / (np.linalg.norm(dist, axis=1) ** 3)[:, np.newaxis]
|
||||||
pos += vel
|
pos += vel
|
||||||
|
|
||||||
|
|
||||||
def n_body_matrix(pos: np.ndarray, vel: np.ndarray, mass: np.ndarray, constrain=2.):
|
def n_body_matrix(pos: np.ndarray, vel: np.ndarray, mass: np.ndarray, constrain=2.):
|
||||||
n, d = pos.shape
|
n, d = pos.shape
|
||||||
dist = np.zeros((n - 1, n, d))
|
dist = np.zeros((n - 1, n, d))
|
||||||
rot_mass = np.zeros((n - 1, n, 1))
|
rot_mass = np.zeros((n - 1, n, 1))
|
||||||
|
|
||||||
pos2 = np.concatenate((pos, pos))
|
pos2 = np.concatenate((pos, pos))
|
||||||
mass2 = np.concatenate((mass, mass))
|
mass2 = np.concatenate((mass, mass))
|
||||||
|
|
||||||
for i in range(1, len(pos)):
|
for i in range(1, len(pos)):
|
||||||
dist[i - 1] = pos2[i: i + n] - pos
|
dist[i - 1] = pos2[i: i + n] - pos
|
||||||
rot_mass[i - 1] = mass2[i: i + n]
|
rot_mass[i - 1] = mass2[i: i + n]
|
||||||
|
|
||||||
norms = np.linalg.norm(dist, axis=2)
|
norms = np.linalg.norm(dist, axis=2)
|
||||||
if constrain:
|
if constrain:
|
||||||
norms[norms < constrain] = constrain
|
norms[norms < constrain] = constrain
|
||||||
|
|
||||||
vel += G * np.sum(
|
vel += G * np.sum(
|
||||||
dist * rot_mass / (norms ** 3)[:, :, np.newaxis],
|
dist * rot_mass / (norms ** 3)[:, :, np.newaxis],
|
||||||
axis=0
|
axis=0
|
||||||
)
|
)
|
||||||
|
|
||||||
pos += vel
|
pos += vel
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue