better structure. most of task 1 done

nbody/checklist.md

@@ -4,8 +4,16 @@
 - [ ] Compute characteristic quantities/scales
 - [x] Compare analytical model and particle density distribution
 - [ ] Compute forces through nbody simulation
-- [ ] vary softening length and compare results
-- [ ] compare with the analytical expectation from Newtons 2nd law
+- [x] vary softening length and compare results
+- [x] compare with the analytical expectation from Newtons 2nd law
 - [ ] compute the relaxation time
 
 ### Task 2
+
+
+
+
+
+
+### Questions
+- Procedure for each time step of a mesh simulation? Potential on mesh -> forces on particles -> update particle positions -> new mesh potential? or skip the creation of particles in each time step?

nbody/utils/__init__.py

@@ -1,7 +1,14 @@
-# Import all functions in all the files in the current directory
+## Import all functions in all the files in the current directory
+# Basic helpers for interacting with the data 
 from .load import *
 from .mesh import *
 from .model import *
 from .particles import *
-from .forces import *
-from .integrate import *
+
+# Helpers for computing the forces
+from .forces_basic import *
+from .forces_tree import *
+from .forces_mesh import *
+
+# Helpers for solving the IVP and having time evolution
+from .integrate import *

nbody/utils/forces_basic.py

@@ -16,7 +16,7 @@ def n_body_forces(particles: np.ndarray, G: float, softening: float = 0):
 
     n = particles.shape[0]
     forces = np.zeros((n, 3))
-    logger.debug(f"Computing forces for {n} particles using n^2 algorithm")
+    logger.debug(f"Computing forces for {n} particles using n^2 algorithm (using {softening=})")
 
     for i in range(n):
         # the current particle is at x_current
@@ -41,7 +41,7 @@ def n_body_forces(particles: np.ndarray, G: float, softening: float = 0):
         f = np.sum((num.T / r_adjusted**1.5).T, axis=0) * m_current
         forces[i] = -f
 
-        if i % 1000 == 0:
+        if i % 5000 == 0:
             logger.debug(f"Particle {i} done")
 
     return forces
@@ -71,7 +71,7 @@ def analytical_forces(particles: np.ndarray):
         f = - m_current * m_enclosed / r_current**2
         forces[i] = f
 
-        if i % 1000 == 0:
+        if i % 5000 == 0:
             logger.debug(f"Particle {i} done")
 
     return forces

nbody/utils/forces_mesh.py

@@ -0,0 +1,185 @@
+## Implementation of a mesh based force solver
+import numpy as np
+import matplotlib.pyplot as plt
+from scipy import fft
+
+import logging
+logger = logging.getLogger(__name__)
+
+#### Version 1 - keeping the derivative of phi
+
+def mesh_forces(particles: np.ndarray, G: float, n_grid: int, mapping: callable) -> np.ndarray:
+    """
+    Computes the gravitational forces between a set of particles using a mesh.
+    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")
+
+    mesh, axis = to_mesh(particles, n_grid, mapping)
+    # show_mesh_information(mesh, "Initial mesh")
+
+    grad_phi = mesh_poisson(mesh, G)
+    # show_mesh_information(mesh, "Mesh potential")
+
+    # compute the particle forces from the mesh potential
+    forces = np.zeros_like(particles[:, :3])
+    for i, p in enumerate(particles):
+        ijk = np.digitize(p, axis) - 1
+        forces[i] = -grad_phi[ijk[0], ijk[1], ijk[2]] * p[3]
+
+    return forces
+
+
+def mesh_poisson(mesh: np.ndarray, G: float) -> np.ndarray:
+    """
+    'Solves' the poisson equation for the mesh using the FFT.
+    Returns the gradient of the potential since this is required for the force computation.
+    """ 
+    rho_hat = fft.fftn(mesh)
+    # the laplacian in fourier space takes the form of a multiplication
+    k = np.fft.fftfreq(mesh.shape[0])
+    # TODO: probably need to take the actual mesh bounds into account
+    kx, ky, kz = np.meshgrid(k, k, k)
+    k_vec = np.array([kx, ky, kz])
+    logger.debug(f"Got k_square with: {k_vec.shape}, {np.max(k_vec)} {np.min(k_vec)}")
+
+    grad_phi_hat = 4 * np.pi * G * rho_hat / (1j * k_vec * 2 * np.pi)
+
+    # the inverse fourier transform gives the potential (or its gradient)
+    grad_phi = np.real(fft.ifftn(grad_phi_hat))
+    return grad_phi
+
+
+#### Version 2 - only storing the scalar potential
+def mesh_forces_v2(particles: np.ndarray, G: float, n_grid: int, mapping: callable) -> np.ndarray:
+    if particles.shape[1] != 4:
+        raise ValueError("Particles array must have 4 columns: x, y, z, m")
+
+    logger.debug(f"Computing forces for {particles.shape[0]} particles using mesh (using mapping={mapping.__name__})")
+
+    mesh, axis = to_mesh(particles, n_grid, mapping)
+    show_mesh_information(mesh, "Mesh")
+
+    spacing = axis[1] - axis[0]
+    logger.debug(f"Using mesh spacing: {spacing}")
+    phi = mesh_poisson_v2(mesh, G)
+    logger.debug(f"Got phi with: {phi.shape}, {np.max(phi)}")
+    phi_grad = np.stack(np.gradient(phi, spacing), axis=0)
+    show_mesh_information(phi, "Mesh potential")
+    show_mesh_information(phi_grad[0], "Mesh potential grad x")
+    logger.debug(f"Got phi_grad with: {phi_grad.shape}, {np.max(phi_grad)}")
+    forces = np.zeros_like(particles[:, :3])
+    for i, p in enumerate(particles):
+        ijk = np.digitize(p, axis) - 1
+        # this gives 4 entries since p[3] the mass is digitized as well -> this is meaningless and we discard it
+        # logger.debug(f"Particle {p} maps to cell {ijk}")
+        forces[i] = - p[3] * phi_grad[..., ijk[0], ijk[1], ijk[2]]
+
+    return forces
+
+
+def mesh_poisson_v2(mesh: np.ndarray, G: float) -> np.ndarray:
+    rho_hat = fft.fftn(mesh)
+    k = np.fft.fftfreq(mesh.shape[0])
+    kx, ky, kz = np.meshgrid(k, k, k)
+    k_sr = kx**2 + ky**2 + kz**2
+    logger.debug(f"Got k_square with: {k_sr.shape}, {np.max(k_sr)} {np.min(k_sr)}")
+    logger.debug(f"Count of zeros: {np.sum(k_sr == 0)}")
+    k_sr[k_sr == 0] = 1e-10  # Add a small epsilon to avoid division by zero
+    phi_hat = - G * rho_hat / k_sr
+    # 4pi cancels, - comes from i squared
+    phi = np.real(fft.ifftn(phi_hat))
+    return phi
+
+
+## Helper functions for star mapping
+def to_mesh(particles: np.ndarray, n_grid: int, mapping: callable) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Maps a list of particles to a of mesh of size n_grid x n_grid x n_grid.
+    Assumes that the particles array has the following columns: x, y, z, ..., m.
+    Uses the mass of the particles and a smoothing function to detemine the contribution to each cell.
+    """
+    # axis provide an easy way to map the particles to the mesh
+    max_pos = np.max(particles[:, :3])
+    axis = np.linspace(-max_pos, max_pos, n_grid)
+    mesh_grid = np.meshgrid(axis, axis, axis)
+    mesh = np.zeros_like(mesh_grid[0])
+
+    for p in particles:
+        m = p[-1]
+        if logger.level <= logging.DEBUG and  m <= 0:
+            logger.warning(f"Particle with negative mass: {p}")
+        # spread the star onto cells through the shape function, taking into account the mass
+        ijks, weights = mapping(p, axis)
+        for ijk, weight in zip(ijks, weights):
+            mesh[ijk[0], ijk[1], ijk[2]] += weight * m
+    
+    return mesh, axis
+
+
+def particle_to_cells_nn(particle, axis):
+    # find the single cell that contains the particle
+    ijk = np.digitize(particle, axis) - 1
+    # the weight is obviously 1
+    return [ijk], [1]
+
+
+def particle_to_cells_cic(particle, axis, width):
+    # create a virtual cell around the particle
+    cell_bounds = [
+        particle + np.array([1, 0, 0]) * width,
+        particle + np.array([-1, 0, 0]) * width,
+        particle + np.array([1, 1, 0]) * width,
+        particle + np.array([-1, -1, 0]) * width,
+        particle + np.array([1, 1, 1]) * width,
+        particle + np.array([-1, -1, 1]) * width,
+        particle + np.array([1, 1, -1]) * width,
+        particle + np.array([-1, -1, -1]) * width,
+    ]
+    # find all the cells that intersect with the virtual cell
+    ijks = []
+    weights = []
+    for b in cell_bounds:
+        w = np.linalg.norm(particle - b)
+        ijk = np.digitize(b, axis) - 1
+        # print(f"b: {b}, ijk: {ijk}")
+        ijks.append(ijk)
+        weights.append(w)
+    # ensure that the weights sum to 1
+    weights = np.array(weights)
+    weights /= np.sum(weights)
+    return ijks, weights
+
+
+
+## Helper functions for mesh plotting
+def show_mesh_information(mesh: np.ndarray, name: str):
+    print(f"Mesh information for {name}")
+    print(f"Total mapped mass: {np.sum(mesh):.0f}")
+    print(f"Max cell value: {np.max(mesh)}")
+    print(f"Min cell value: {np.min(mesh)}")
+    print(f"Mean cell value: {np.mean(mesh)}")
+    plot_3d(mesh, name)
+    plot_2d(mesh, name)
+
+
+def plot_3d(mesh: np.ndarray, name: str):
+    fig = plt.figure()
+    fig.suptitle(name)
+    ax = fig.add_subplot(111, projection='3d')
+    ax.scatter(*np.where(mesh), c=mesh[np.where(mesh)], cmap='viridis')
+    plt.show()
+
+
+def plot_2d(mesh: np.ndarray, name: str):
+    fig = plt.figure()
+    fig.suptitle(name)
+    axs = fig.subplots(1, 3)
+    axs[0].imshow(np.sum(mesh, axis=0))
+    axs[0].set_title("Flattened in x")
+    axs[1].imshow(np.sum(mesh, axis=1))
+    axs[1].set_title("Flattened in y")
+    axs[2].imshow(np.sum(mesh, axis=2))
+    axs[2].set_title("Flattened in z")
+    plt.show()

nbody/utils/logging_config.py

@@ -0,0 +1,15 @@
+import logging
+
+logging.basicConfig(
+    ## set logging level
+    # level=logging.INFO,
+    level=logging.DEBUG,
+    format='%(asctime)s - %(name)s - %(message)s',
+    datefmt='%H:%M:%S'
+)
+
+
+# silence some debug messages
+logging.getLogger('matplotlib.font_manager').setLevel(logging.WARNING)
+logging.getLogger('matplotlib.ticker').setLevel(logging.WARNING)
+logging.getLogger('matplotlib.pyplot').setLevel(logging.WARNING)