~finished task 1, on to integration

fft/potentials.ipynb

@@ -223,7 +223,6 @@
    "source": [
     "def acceleration(phi, grid, center, end):\n",
     "    # line between center and end\n",
-    "    # todo center and end are wrong\n",
     "    direction = (end - center / np.linalg.norm(end - center)).astype(int)\n",
     "    print(direction)\n",
     "    points = [center + i * direction for i in range(1, phi.shape[0]//2)]\n",
@@ -253,7 +252,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.7"
+   "version": "3.13.1"
   }
  },
  "nbformat": 4,

nbody/checklist.md

@@ -35,3 +35,4 @@
     - plot total energy vs time
     - plot particle positions?
 - What is the parameter a of the Hernquist model?
+- How to handle zero values for the k vector in the mesh method?

nbody/utils/forces_basic.py

@@ -60,11 +60,11 @@ def analytical_forces(particles: np.ndarray):
 
     logger.debug(f"Computing forces for {n} particles using spherical approximation")
 
+    r_particles = np.linalg.norm(particles[:, :3], axis=1)
     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

nbody/utils/forces_mesh.py

@@ -53,6 +53,10 @@ def mesh_poisson(mesh: np.ndarray, G: float) -> np.ndarray:
 
 #### Version 2 - only storing the scalar potential
 def mesh_forces_v2(particles: np.ndarray, G: float, n_grid: int, mapping: callable) -> np.ndarray:
+    """
+    Computes the gravitational force acting on a set of particles using a mesh-based approach.
+    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")
 
@@ -62,9 +66,10 @@ def mesh_forces_v2(particles: np.ndarray, G: float, n_grid: int, mapping: callab
     if logger.level >= logging.DEBUG:
         show_mesh_information(mesh, "Density mesh")
 
+    # compute the potential and its gradient
     spacing = axis[1] - axis[0]
     logger.debug(f"Using mesh spacing: {spacing}")
-    phi = mesh_poisson_v2(mesh, G)
+    phi = mesh_poisson_v2(mesh, G, spacing)
     logger.debug(f"Got phi with: {phi.shape}, {np.max(phi)}")
     phi_grad = np.stack(np.gradient(phi, spacing), axis=0)
     if logger.level >= logging.DEBUG:
@@ -72,6 +77,7 @@ def mesh_forces_v2(particles: np.ndarray, G: float, n_grid: int, mapping: callab
         show_mesh_information(phi_grad[0], "Potential gradient")
     logger.debug(f"Got phi_grad with: {phi_grad.shape}, {np.max(phi_grad)}")
     
+    # 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
@@ -82,16 +88,23 @@ def mesh_forces_v2(particles: np.ndarray, G: float, n_grid: int, mapping: callab
     return forces
 
 
-def mesh_poisson_v2(mesh: np.ndarray, G: float) -> np.ndarray:
+def mesh_poisson_v2(mesh: np.ndarray, G: float, spacing: float) -> np.ndarray:
+    """
+    Solves the poisson equation for the mesh using the FFT.
+    Returns the scalar potential.
+    """
     rho_hat = fft.fftn(mesh)
-    k = np.fft.fftfreq(mesh.shape[0])
+    k = fft.fftfreq(mesh.shape[0], spacing)
     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
+    # avoid division by zero
+    # TODO: review this
+    k_sr[k_sr == 0] = np.inf
+    phi_hat = - 4 * np.pi * G * rho_hat / k_sr
+    # - comes from i squared
+    # TODO: 4pi stays since the backtransform removes the 1/2pi factor
     phi = np.real(fft.ifftn(phi_hat))
     return phi
 

nbody/utils/integrate.py

@@ -66,3 +66,12 @@ def to_particles(y: np.ndarray) -> np.ndarray:
     y = y.reshape((n, 7), copy=False, order='F')
     logger.debug(f"Unflattened array into {y.shape=}")
     return y
+
+
+
+def runge_kutta_4(y0 : np.ndarray, t : float, f, dt : float):
+    k1 = f(y0, t)
+    k2 = f(y0 + k1/2 * dt, t + dt/2)
+    k3 = f(y0 + k2/2 * dt, t + dt/2)
+    k4 = f(y0 + k3 * dt, t + dt)
+    return y0 + (k1 + 2*k2 + 2*k3 + k4)/6 * dt

nbody/utils/units.py

@@ -17,7 +17,7 @@ def seed_scales(r_scale: u.Quantity, m_scale: u.Quantity):
     global M_SCALE, R_SCALE
     M_SCALE = m_scale
     R_SCALE = r_scale
-    logger.info(f"Set scales: M_SCALE = {M_SCALE}, R_SCALE = {R_SCALE}")
+    logger.info(f"Set scales: M_SCALE = {M_SCALE:.2g}, R_SCALE = {R_SCALE:.2g}")
 
 
 def apply_units(columns: np.array, quantity: str):
@@ -25,17 +25,20 @@ def apply_units(columns: np.array, quantity: str):
         return columns * M_SCALE
     elif quantity == "position":
         return columns * R_SCALE
-    elif quantity == "velocity":
-        return columns * R_SCALE / u.s
-    elif quantity == "time":
-        return columns * u.s
-    elif quantity == "acceleration":
-        return columns * R_SCALE / u.s**2
-    elif quantity == "force":
-        return columns * M_SCALE * R_SCALE / u.s**2
     elif quantity == "volume":
         return columns * R_SCALE**3
     elif quantity == "density":
         return columns * M_SCALE / R_SCALE**3
+    
+    ## Derived quantities
+    elif quantity == "force":
+        # using F = GMm/R^2 => F = M_SCALE**2 / R_SCALE**2 (G = 1)
+        return columns * M_SCALE**2 / R_SCALE**2
+    elif quantity == "velocity":
+        # using the Virial theorem: v^2 = GM/R => v = sqrt(GM/R) => v = sqrt(M_SCALE / R_SCALE) (G = 1)
+        return columns * np.sqrt(M_SCALE / R_SCALE)
+    elif quantity == "time":
+        # using the dynamical time: t_dyn = 1/sqrt(G*rho) => t_dyn = sqrt(4/3 * pi * R_SCALE**3 / M_SCALE) (G = 1)
+        return columns * np.sqrt(4/3 * np.pi * R_SCALE**3 / M_SCALE)
     else:
         raise ValueError(f"Unknown quantity: {quantity}")