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
|
||||
|
||||
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.
|
||||
|
||||
|
|
|
|||
152
data.py
152
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()
|
||||
|
|
|
|||
52
gen_data.py
52
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)}")
|
||||
|
|
|
|||
60
main.py
60
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()
|
||||
|
|
|
|||
46
physics.py
46
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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue