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

Support 3d simulations

Browse source
This commit is contained in:
Michael Bradley 2023-10-08 12:13:53 -04:00
parent 4a31fc1a7d
commit 424721ccfc
4 changed files with 70 additions and 46 deletions
Download patch file Download diff file Expand all files Collapse all files

69
data.py
View file

@ -2,20 +2,28 @@ import matplotlib.animation as animation
import matplotlib.cm as cm
import matplotlib.pyplot as plt
import numpy as np
import physics
def parse_csv(filename: str):
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), 2))
vel = np.zeros((len(lines), 2))
pos = np.zeros((len(lines), dimensions))
vel = np.zeros((len(lines), dimensions))
rad = np.zeros((len(lines), 1))
for i, [x, y, vx, vy, r] in enumerate(map(lambda l: map(float, l.split(',')), lines)):
pos[i] = [x, y]
vel[i] = [vx, vy]
rad[i] = r
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
@ -26,14 +34,20 @@ class Animator:
self.rad = rad
self.mass = np.pi * 4 / 3 * rad ** 3
self.scat = None
n, d = self.pos.shape
self.scat: plt.PathCollection = None
self.colours = cm.rainbow(
np.random.random(
(len(self.rad),)
(n,)
)
)
self.fig, self.ax = plt.subplots()
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,
@ -44,18 +58,33 @@ class Animator:
)
def setup_plot(self):
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])
_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
)
self.ax.axis([-500, 500, -500, 500, -500, 500])
return self.scat,
def update(self, *_args, **_kwargs):
physics.n_body_matrix_constrained(self.pos, self.vel, self.mass)
self.scat.set_offsets(self.pos)
_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.fig.canvas.draw()
return self.scat,
def show(self):