From 8bb6ba19d5f939e16800127f766a64983dc8a62f Mon Sep 17 00:00:00 2001 From: Michael Bradley Date: Thu, 2 Jan 2025 00:05:16 +1300 Subject: [PATCH] Switch to tabs --- README.md | 4 +- data.py | 152 ++++++++++++++++++++++++++-------------------------- gen_data.py | 52 +++++++++--------- main.py | 60 ++++++++++----------- physics.py | 46 ++++++++-------- 5 files changed, 157 insertions(+), 157 deletions(-) diff --git a/README.md b/README.md index 22c4f2d..85837d8 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ # nbody -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). +Threw this together to get more comfortable with Numpy. +Can simulate a few hundred bodies in 2 or 3 dimensions without much hassle. Comments are non-existent, sorry about that. diff --git a/data.py b/data.py index 4c2df2a..7293163 100644 --- a/data.py +++ b/data.py @@ -6,88 +6,88 @@ import physics 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 + 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: - 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 + 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 - n, d = self.pos.shape + n, d = self.pos.shape - self.scat: plt.PathCollection = None - self.colours = cm.rainbow( - np.random.random( - (n,) - ) - ) + self.scat: plt.PathCollection = None + self.colours = cm.rainbow( + np.random.random( + (n,) + ) + ) - 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 - ) + 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 + ) - 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, + 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, - 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, + 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, - def show(self): - plt.show() + def show(self): + plt.show() diff --git a/gen_data.py b/gen_data.py index 1d1e32d..8030612 100755 --- a/gen_data.py +++ b/gen_data.py @@ -4,35 +4,35 @@ from random import uniform, randint class Args: - width: int - height: int - depth: int - velocity: float - mass: float - count: int + width: int + height: int + depth: int + velocity: float + mass: float + count: int if __name__ == "__main__": - parser = argparse.ArgumentParser( - prog="n-body data generator", - description="Generates data for the n-body simulator.", - add_help=False - ) + 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=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) + 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) - args: 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)}") + 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 4e9cc2a..4a95d0d 100755 --- a/main.py +++ b/main.py @@ -7,40 +7,40 @@ import physics class Args: - filename: str - gravity: float - dimensions: typing.Literal[2, 3] + filename: str + gravity: float + dimensions: typing.Literal[2, 3] if __name__ == "__main__": - parser = argparse.ArgumentParser( - prog="n-body simulation", - description="Simulating gravitational effects" - ) + parser = argparse.ArgumentParser( + prog="n-body simulation", + description="Simulating gravitational effects" + ) - 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 - ) + 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: 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) - a = data.Animator(*objects) - a.show() + objects = data.parse_csv(args.filename, dimensions=args.dimensions) + a = data.Animator(*objects) + a.show() diff --git a/physics.py b/physics.py index 4965775..0903530 100644 --- a/physics.py +++ b/physics.py @@ -5,37 +5,37 @@ G = 6.674e-11 def rotations(a: np.ndarray): - a2 = np.concatenate((a, a)) - for i in range(1, len(a)): - yield a2[i: i + len(a)] + 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] - pos += vel + 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 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)) + n, d = pos.shape + dist = np.zeros((n - 1, n, d)) + rot_mass = np.zeros((n - 1, n, 1)) - pos2 = np.concatenate((pos, pos)) - mass2 = np.concatenate((mass, mass)) + pos2 = np.concatenate((pos, pos)) + mass2 = np.concatenate((mass, mass)) - for i in range(1, len(pos)): - dist[i - 1] = pos2[i: i + n] - pos - rot_mass[i - 1] = mass2[i: i + n] + 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 + 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 - ) + vel += G * np.sum( + dist * rot_mass / (norms ** 3)[:, :, np.newaxis], + axis=0 + ) - pos += vel + pos += vel