Correct radius dimensions

This commit is contained in:
Michael Bradley 2023-10-07 18:41:29 -04:00
parent 3ebf5b9baf
commit d54c95347f
2 changed files with 5 additions and 3 deletions

View file

@ -11,11 +11,11 @@ def parse_csv(filename: str):
lines = file.read().strip().splitlines() lines = file.read().strip().splitlines()
pos = np.zeros((len(lines), 2)) pos = np.zeros((len(lines), 2))
vel = np.zeros((len(lines), 2)) vel = np.zeros((len(lines), 2))
rad = np.zeros((len(lines),)) rad = np.zeros((len(lines), 1))
for i, [x, y, vx, vy, r] in enumerate(map(lambda l: map(float, l.split(',')), lines)): for i, [x, y, vx, vy, r] in enumerate(map(lambda l: map(float, l.split(',')), lines)):
pos[i] = [x, y] pos[i] = [x, y]
vel[i] = [vx, vy] vel[i] = [vx, vy]
rad[i] = r rad[i] = [r]
return pos, vel, rad return pos, vel, rad

View file

@ -8,10 +8,12 @@ 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 np.split(a2, [i, i + len(a)])[1] yield np.split(a2, [i, i + len(a)])[1]
# TODO: Compare performance
# yield np.roll(a, i)
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[:, np.newaxis])): for (o_pos, o_mass) in zip(rotations(pos), rotations(mass)):
dist = o_pos - pos dist = o_pos - pos
vel += G * (dist / np.linalg.norm(dist, axis=1)[:, np.newaxis]) * o_mass / np.sum(dist ** 2, axis=1)[:, np.newaxis] vel += G * (dist / np.linalg.norm(dist, axis=1)[:, np.newaxis]) * o_mass / np.sum(dist ** 2, axis=1)[:, np.newaxis]
pos += vel pos += vel