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

34
physics.py
View file

@ -17,38 +17,22 @@ def n_body(pos: np.ndarray, vel: np.ndarray, mass: np.ndarray):
pos += vel
def n_body_matrix(pos: np.ndarray, vel: np.ndarray, mass: np.ndarray):
dist = np.zeros((len(pos) - 1, len(pos), 2))
rot_mass = np.zeros((len(mass) - 1, len(mass), 1))
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))
pos2 = np.concatenate((pos, pos))
mass2 = np.concatenate((mass, mass))
for i in range(1, len(pos)):
dist[i - 1] = pos2[i: i + len(pos)] - pos
rot_mass[i - 1] = mass2[i: i + len(mass)]
vel += G * np.sum(
dist * rot_mass / (np.linalg.norm(dist, axis=2) ** 3)[:, :, np.newaxis],
axis=0
)
pos += vel
def n_body_matrix_constrained(pos: np.ndarray, vel: np.ndarray, mass: np.ndarray, close=2.):
dist = np.zeros((len(pos) - 1, len(pos), 2))
rot_mass = np.zeros((len(mass) - 1, len(mass), 1))
pos2 = np.concatenate((pos, pos))
mass2 = np.concatenate((mass, mass))
for i in range(1, len(pos)):
dist[i - 1] = pos2[i: i + len(pos)] - pos
rot_mass[i - 1] = mass2[i: i + len(mass)]
dist[i - 1] = pos2[i: i + n] - pos
rot_mass[i - 1] = mass2[i: i + n]
norms = np.linalg.norm(dist, axis=2)
norms[norms < close] = close
if constrain:
norms[norms < constrain] = constrain
vel += G * np.sum(
dist * rot_mass / (norms ** 3)[:, :, np.newaxis],
axis=0