78 lines
2.6 KiB
Python
78 lines
2.6 KiB
Python
import numpy as np
|
|
import logging
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def n_body_forces(particles: np.ndarray, G: float, softening: float = 0):
|
|
"""
|
|
Computes the gravitational forces between a set of particles.
|
|
Assumes that the particles array has the following columns: x, y, z, m.
|
|
"""
|
|
if particles.shape[1] != 4:
|
|
raise ValueError("Particles array must have 4 columns: x, y, z, m")
|
|
|
|
x_vec = particles[:, 0:3]
|
|
masses = particles[:, 3]
|
|
|
|
n = particles.shape[0]
|
|
forces = np.zeros((n, 3))
|
|
logger.debug(f"Computing forces for {n} particles using n^2 algorithm (using {softening=:.2g})")
|
|
|
|
for i in range(n):
|
|
# the current particle is at x_current
|
|
x_current = x_vec[i, :]
|
|
m_current = masses[i]
|
|
|
|
# first compute the displacement to all other particles
|
|
displacements = x_vec - x_current
|
|
# and its magnitude
|
|
r = np.linalg.norm(displacements, axis=1)
|
|
# add softening to the denominator
|
|
r_adjusted = r**2 + softening**2
|
|
# the numerator is tricky:
|
|
# m is a list of scalars and displacements is a list of vectors (2D array)
|
|
# we only want row_wise multiplication
|
|
num = G * (masses * displacements.T).T
|
|
|
|
# a zero value is expected where we have the same particle
|
|
r_adjusted[i] = 1
|
|
num[i] = 0
|
|
|
|
f = np.sum((num.T / r_adjusted**1.5).T, axis=0) * m_current
|
|
forces[i] = -f
|
|
|
|
if i % 5000 == 0:
|
|
logger.debug(f"Particle {i} done")
|
|
|
|
return forces
|
|
|
|
|
|
|
|
def analytical_forces(particles: np.ndarray):
|
|
"""
|
|
Computes the interparticle forces without computing the n^2 interactions.
|
|
This is done by using newton's second theorem for a spherical mass distribution.
|
|
The force on a particle at radius r is simply the force exerted by a point mass with the enclosed mass.
|
|
Assumes that the particles array has the following columns: x, y, z, m.
|
|
"""
|
|
n = particles.shape[0]
|
|
forces = np.zeros((n, 3))
|
|
|
|
logger.debug(f"Computing forces for {n} particles using spherical approximation")
|
|
|
|
for i in range(n):
|
|
r_current = np.linalg.norm(particles[i, 0:3])
|
|
m_current = particles[i, 3]
|
|
|
|
r_particles = np.linalg.norm(particles[:, :3], axis=1)
|
|
m_enclosed = np.sum(particles[r_particles < r_current, 3])
|
|
|
|
# the force is the same as the force exerted by a point mass at the center
|
|
f = - m_current * m_enclosed / r_current**2
|
|
forces[i] = f
|
|
|
|
if i % 5000 == 0:
|
|
logger.debug(f"Particle {i} done")
|
|
|
|
return forces
|