most static helpers up and running

nbody/utils/__init__.py

@@ -0,0 +1,7 @@
+# Import all functions in all the files in the current directory
+from .load import *
+from .mesh import *
+from .model import *
+from .particles import *
+from .forces import *
+from .integrate import *

nbody/utils/forces.py

@@ -0,0 +1,77 @@
+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")
+
+    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 % 1000 == 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 % 1000 == 0:
+            logger.debug(f"Particle {i} done")
+
+    return forces

nbody/utils/integrate.py

@@ -0,0 +1,69 @@
+import numpy as np
+import scipy.integrate as spi
+
+import logging
+logger = logging.getLogger(__name__)
+
+
+def ode_setup(particles: np.ndarray, force_function: callable) -> tuple[np.ndarray, callable]:
+    """
+    Linearizes the ODE system for the particles interacting gravitationally.
+    Returns:
+        - the Y0 array corresponding to the initial conditions (x0 and v0)
+        - the function that computes the right hand side of the ODE with function signature f(t, y)
+    Assumes that the particles array has the following columns: x, y, z, vx, vy, vz, m.
+    """
+    if particles.shape[1] != 7:
+        raise ValueError("Particles array must have 7 columns: x, y, z, vx, vy, vz, m")
+    
+    n = particles.shape[0]
+    # for scipy integrators we need to flatten the n 3D positions and n 3D velocities
+    y0 = np.zeros(6*n) 
+    y0[:3*n] = particles[:, :3].flatten()
+    y0[3*n:] = particles[:, 3:6].flatten()
+
+    # the masses don't change we can define them once
+    masses = particles[:, 6]
+    logger.debug(f"Reshaped {particles.shape} to y0 with {y0.shape} and masses with {masses.shape}")
+
+
+    def f(y, t):
+        """
+        Computes the right hand side of the ODE system.
+        The ODE system is linearized around the current positions and velocities.
+        """
+        n = y.size // 6
+        logger.debug(f"y with shape {y.shape}")
+        # unsqueeze and unstack to extract the positions and velocities
+        y = y.reshape((2*n, 3))
+        x = y[:n, ...]
+        v = y[n:, ...]
+        logger.debug(f"Unstacked y into x with shape {x.shape} and v with shape {v.shape}")
+        
+        # compute the forces
+        x_with_m = np.zeros((n, 4))
+        x_with_m[:, :3] = x
+        x_with_m[:, 3] = masses
+        forces = force_function(x_with_m)
+        
+        # compute the accelerations
+        a = forces / masses[:, None]
+        a.flatten()
+        # the [:, None] is to force broadcasting in order to divide each row of forces by the corresponding mass
+        
+        # reshape into a 1D array
+        return np.vstack((v, a)).flatten()
+    
+    return y0, f
+
+
+def to_particles(y: np.ndarray) -> np.ndarray:
+    """
+    Converts the 1D array y into a 2D array with the shape (n, 6) where n is the number of particles.
+    The columns are x, y, z, vx, vy, vz
+    """
+    n = y.size // 6
+    y = y.reshape((2*n, 3))
+    x = y[:n, ...]
+    v = y[n:, ...]
+    return np.hstack((x, v))

nbody/utils/load.py

@@ -0,0 +1,18 @@
+import numpy as np
+from pathlib import Path
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def load_data(file: Path) -> tuple[np.ndarray, list]:
+    try:
+        data = np.loadtxt(file)
+        columns =  []
+    except ValueError:
+        data = np.loadtxt(file, skiprows=1)
+        header = file.read_text().splitlines()[0]
+        columns = header.split()
+    logger.info(f"Loaded {data.shape[0]} rows and {data.shape[1]} columns from {file}")
+    return data, columns
+

nbody/utils/model.py

@@ -0,0 +1,12 @@
+import numpy as np
+
+M = 5
+a = 5
+
+def model_density_distribution(r_bins: np.ndarray):
+    """
+    Generate a density distribution for a spherical galaxy model, as per the Hernquist model.
+    See https://doi.org/10.1086%2F168845 for more information.
+    """
+    rho = M / (2 * np.pi) * a / (r_bins * (r_bins + a)**3)
+    return rho

nbody/utils/particles.py

@@ -0,0 +1,124 @@
+import numpy as np
+import logging
+logger = logging.getLogger(__name__)
+
+
+def density_distribution(r_bins: np.ndarray, particles: np.ndarray, ret_error: bool = False):
+    """
+    Computes the radial density distribution of 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")
+
+    m = particles[:, 3]
+    r = np.linalg.norm(particles[:, :3], axis=1)
+    density = [np.sum(m[(r >= r_bins[i]) & (r < r_bins[i + 1])]) for i in range(len(r_bins) - 1)]
+
+    # add the first volume which  should be wrt 0
+    volume = 4/3 * np.pi * (r_bins[1:]**3 - r_bins[:-1]**3)
+    volume = np.insert(volume, 0, 4/3 * np.pi * r_bins[0]**3)
+    density = r_bins / volume
+    if ret_error:
+        return density, density / np.sqrt(r_bins)
+    else:
+        return density
+
+
+
+def r_distribution(particles: np.ndarray):
+    """
+    Computes the distribution of distances (to the origin) of a set of particles.
+    Assumes that the particles array has the following columns: x, y, z ...
+    """
+    if particles.shape[1] < 3:
+        raise ValueError("Particles array must have at least 3 columns: x, y, z")
+
+    r = np.linalg.norm(particles[:, :3], axis=1)
+    return r
+
+
+
+def remove_outliers(particles: np.ndarray, std_threshold: float = 3):
+    """
+    Removes outliers from a set of particles.
+    Assumes that the particles array has the following columns: x, y, z ...
+    """
+    if particles.shape[1] < 3:
+        raise ValueError("Particles array must have at least 3 columns: x, y, z")
+
+    r = np.linalg.norm(particles[:, :3], axis=1)
+    r_std = np.std(r)
+    r_mean = np.mean(r)
+    mask = np.abs(r - r_mean) < std_threshold * r_std
+    return particles[mask]
+
+
+
+def mean_interparticle_distance(particles: np.ndarray):
+    """
+    Computes the mean interparticle distance of a set of particles.
+    Assumes that the particles array has the following columns: x, y, z ...
+    """
+    if particles.shape[1] < 3:
+        raise ValueError("Particles array must have at least 3 columns: x, y, z")
+    
+
+    r_half_mass = half_mass_radius(particles)
+    r = np.linalg.norm(particles[:, :3], axis=1)
+
+    n_half_mass = np.sum(r < r_half_mass)
+    logger.debug(f"Number of particles within half mass radius: {n_half_mass} of {particles.shape[0]}")
+
+    rho = n_half_mass / (4/3 * np.pi * r_half_mass**3)
+    # the mean distance between particles is the inverse of the density
+    return (1 / rho)**(1/3)
+    # TODO: check if this is correct
+
+
+
+def half_mass_radius(particles: np.ndarray):
+    """
+    Computes the half mass radius of a set of particles.
+    Assumes that the particles array has the following columns: x, y, z ...
+    """
+    if particles.shape[1] < 3:
+        raise ValueError("Particles array must have at least 3 columns: x, y, z")
+    
+    # even though in the simple example, all the masses are the same, we will consider the general case 
+    total_mass = np.sum(particles[:, 3])
+    half_mass = total_mass / 2
+    
+    # sort the particles by distance
+    r = np.linalg.norm(particles[:, :3], axis=1)
+    indices = np.argsort(r)
+    r = r[indices]
+    masses = particles[indices, 3]
+    masses_cumsum = np.cumsum(masses)
+
+    i = np.argmin(np.abs(masses_cumsum - half_mass))
+    logger.debug(f"Half mass radius: {r[i]} for {i}th particle of {particles.shape[0]}")
+    r_hm = r[i]
+
+    return r_hm
+
+
+
+def relaxation_timescale(particles: np.ndarray, G:float) -> float:
+    """
+    Computes the relaxation timescale of a set of particles using the velocity at the half mass radius.
+    Assumes that the particles array has the following columns: x, y, z ...
+    """
+    m_half = np.sum(particles[:, 3]) / 2 # enclosed mass at half mass radius
+    r_half = half_mass_radius(particles)
+    n_half = np.sum(np.linalg.norm(particles[:, :3], axis=1) < r_half) # number of enclosed particles
+    v_c = np.sqrt(G * m_half / r_half)
+
+    # the crossing time for the half mass system is
+    t_c = r_half / v_c
+    logger.debug(f"Crossing time for half mass system: {t_c}")
+
+    # the relaxation timescale is t_c * N/(10 * log(N))
+    t_rel = t_c * n_half / (10 * np.ln(n_half))
+
+    return t_rel