working nsquare integration?

nbody/checklist.md

@@ -30,7 +30,6 @@
 
 
 ### 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?
 - How to represent the time evolution of the system?
     - plot total energy vs time
     - plot particle positions?

nbody/task1.ipynb

@@ -17,6 +17,7 @@
     "\n",
     "import utils\n",
     "import utils.logging_config\n",
+    "utils.logging_config.set_log_level(\"info\")\n",
     "import logging\n",
     "logger = logging.getLogger(\"task1\")"
    ]
@@ -42,7 +43,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "08:59:02 - utils.load - Loaded 50010 rows and 10 columns from data/data.txt\n"
+      "12:52:17 - utils.load - Loaded 50010 rows and 10 columns from data/data.txt\n"
      ]
     }
    ],
@@ -104,8 +105,8 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "08:59:04 - task1 - Considering a globular cluster - total mass of particles: 4622219.258999999, maximum radius of particles: 724.689657812915\n",
+      "12:52:19 - task1 - Considering a globular cluster - total mass of particles: 4622219.258999999, maximum radius of particles: 724.689657812915\n",
-      "08:59:04 - utils.units - Set scales: M_SCALE = 0.022 solMass, R_SCALE = 0.028 pc\n"
+      "12:52:19 - utils.units - Set scales: M_SCALE = 0.022 solMass, R_SCALE = 0.028 pc\n"
      ]
     }
    ],
@@ -206,7 +207,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "08:59:05 - utils.particles - Found mean interparticle distance: 0.010402746349924056\n"
+      "12:52:19 - utils.particles - Found mean interparticle distance: 0.010402746349924056\n"
      ]
     }
    ],
@@ -304,8 +305,9 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "08:59:24 - task1 - Crossing time for half mass system: 1.7e-06 pc(3/2) / solMass(1/2)\n",
+      "12:54:14 - task1 - Central velocity @ HM 3098.502740746268 solMass(1/2) / pc(1/2)\n",
-      "08:59:24 - task1 - Direct estimate of the relaxation timescale: 0.00078 pc(3/2) / solMass(1/2)\n"
+      "12:54:14 - task1 - Crossing time for half mass system: 1.7e-06 pc(3/2) / solMass(1/2)\n",
+      "12:54:14 - task1 - Direct estimate of the relaxation timescale: 0.00078 pc(3/2) / solMass(1/2)\n"
      ]
     }
    ],
@@ -320,6 +322,7 @@
     "r_half = utils.apply_units(r_half, \"position\")\n",
     "\n",
     "v_c = np.sqrt(G * m_half / r_half)\n",
+    "logger.info(f\"Central velocity @ HM {v_c}\")\n",
     "\n",
     "t_c = r_half / v_c\n",
     "logger.info(f\"Crossing time for half mass system: {t_c:.2g}\")\n",
@@ -360,7 +363,9 @@
     "\n",
     "##### Effect\n",
     "- The relaxation time decreases with increasing softening length\n",
-    "- From the integration over all impact parameters $b$ even $b_{min}$ is chosen to be larger than $\\varepsilon$ $\\implies$ expect only a small effect on the relaxation time"
+    "- From the integration over all impact parameters $b$ even $b_{min}$ is chosen to be larger than $\\varepsilon$ $\\implies$ expect only a small effect on the relaxation time\n",
+    "\n",
+    "#TODO : The softening dampens the change of velocity => time to relax is longer"
    ]
   }
  ],

nbody/utils/forces_basic.py

@@ -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 % 5000 == 0:
+        if i!= 0 and i % 5000 == 0:
             logger.debug(f"Particle {i} done")
 
     return forces

nbody/utils/forces_mesh.py

@@ -8,7 +8,7 @@ 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:
+'''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.
@@ -39,6 +39,8 @@ def mesh_poisson(mesh: np.ndarray, G: float) -> np.ndarray:
     rho_hat = fft.fftn(mesh)
     # the laplacian in fourier space takes the form of a multiplication
     k = np.fft.fftfreq(mesh.shape[0])
+    # shift the zero frequency to the center
+    k = np.fft.fftshift(k)
     # 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])
@@ -49,7 +51,7 @@ def mesh_poisson(mesh: np.ndarray, G: float) -> np.ndarray:
     # 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:
@@ -63,27 +65,34 @@ def mesh_forces_v2(particles: np.ndarray, G: float, n_grid: int, mapping: callab
     logger.debug(f"Computing forces for {particles.shape[0]} particles using mesh [mapping={mapping.__name__}, {n_grid=}]")
 
     mesh, axis = to_mesh(particles, n_grid, mapping)
+    spacing = axis[1] - axis[0]
+    logger.debug(f"Using mesh spacing: {spacing}")
+
+    # we want a density mesh:
+    cell_volume = spacing**3
+    rho = mesh / cell_volume
+
     if logger.level >= logging.DEBUG:
         show_mesh_information(mesh, "Density mesh")
 
     # compute the potential and its gradient
-    spacing = axis[1] - axis[0]
+    phi = mesh_poisson_v2(rho, G, spacing)
-    logger.debug(f"Using mesh spacing: {spacing}")
-    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:
         show_mesh_information(phi, "Potential mesh")
-        show_mesh_information(phi_grad[0], "Potential gradient")
+        show_mesh_information(phi_grad[0], "Potential gradient (x-direction)")
         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
         # 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]]
+        forces[i] = - p[3] * phi_grad[..., ijk[0], ijk[1], ijk[2]] / 10
+        # TODO remove factor of 10
+        # TODO could also index phi_grad the other way around?
 
     return forces
 
@@ -95,13 +104,19 @@ def mesh_poisson_v2(mesh: np.ndarray, G: float, spacing: float) -> np.ndarray:
     """
     rho_hat = fft.fftn(mesh)
     k = fft.fftfreq(mesh.shape[0], spacing)
+    # shift the zero frequency to the center
+    k = np.fft.fftshift(k)
+
     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)}")
+    if logger.level >= logging.DEBUG:
-    logger.debug(f"Count of zeros: {np.sum(k_sr == 0)}")
+        logger.debug(f"Got k_square with: {k_sr.shape}, {np.max(k_sr)} {np.min(k_sr)}")
+        logger.debug(f"Count of ksquare zeros: {np.sum(k_sr == 0)}")
+        show_mesh_information(np.abs(k_sr), "k_square")
     # avoid division by zero
     # TODO: review this
     k_sr[k_sr == 0] = np.inf
+    logger.debug(f"Proceeding to poisson equation with {rho_hat.shape=}, {k_sr.shape=}")
     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
@@ -116,6 +131,8 @@ def to_mesh(particles: np.ndarray, n_grid: int, mapping: callable) -> tuple[np.n
     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.
     """
+    if particles.shape[1] < 4:
+        raise ValueError("Particles array must have at least 4 columns: x, y, z, m")
     # 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)
@@ -176,11 +193,11 @@ def show_mesh_information(mesh: np.ndarray, name: str):
     logger.info(f"Max cell value: {np.max(mesh)}")
     logger.info(f"Min cell value: {np.min(mesh)}")
     logger.info(f"Mean cell value: {np.mean(mesh)}")
-    plot_3d(mesh, name)
+    mesh_plot_3d(mesh, name)
-    plot_2d(mesh, name)
+    mesh_plot_2d(mesh, name)
 
 
-def plot_3d(mesh: np.ndarray, name: str):
+def mesh_plot_3d(mesh: np.ndarray, name: str):
     fig = plt.figure()
     fig.suptitle(f"{name} - {mesh.shape}")
     ax = fig.add_subplot(111, projection='3d')
@@ -188,14 +205,17 @@ def plot_3d(mesh: np.ndarray, name: str):
     plt.show()
 
 
-def plot_2d(mesh: np.ndarray, name: str):
+def mesh_plot_2d(mesh: np.ndarray, name: str, only_z: bool = False):
     fig = plt.figure()
     fig.suptitle(f"{name} - {mesh.shape}")
-    axs = fig.subplots(1, 3)
+    if only_z:
-    axs[0].imshow(np.sum(mesh, axis=0))
+        plt.imshow(np.sum(mesh, axis=2), cmap='viridis', origin='lower')
-    axs[0].set_title("Flattened in x")
+    else:    
-    axs[1].imshow(np.sum(mesh, axis=1))
+        axs = fig.subplots(1, 3)
-    axs[1].set_title("Flattened in y")
+        axs[0].imshow(np.sum(mesh, axis=0), origin='lower')
-    axs[2].imshow(np.sum(mesh, axis=2))
+        axs[0].set_title("Flattened in x")
-    axs[2].set_title("Flattened in z")
+        axs[1].imshow(np.sum(mesh, axis=1), origin='lower')
+        axs[1].set_title("Flattened in y")
+        axs[2].imshow(np.sum(mesh, axis=2), origin='lower')
+        axs[2].set_title("Flattened in z")
     plt.show()

nbody/utils/logging_config.py

@@ -1,12 +1,12 @@
 import logging
 
-logging.basicConfig(
+
-    ## set logging level
+def set_log_level(level: str):
-    level = logging.DEBUG,
+    logging.basicConfig(
-    # level = logging.INFO,
+        level = logging.DEBUG if level == 'debug' else logging.INFO,
-    format = '%(asctime)s - %(name)s - %(message)s',
+        format = '%(asctime)s - %(name)s - %(message)s',
-    datefmt = '%H:%M:%S'
+        datefmt = '%H:%M:%S'
-)
+    )
 
 
 # silence some debug messages

nbody/utils/particles.py

@@ -1,8 +1,7 @@
 import numpy as np
 import logging
-from . import forces_basic
 logger = logging.getLogger(__name__)
-
+import matplotlib.pyplot as plt
 
 def density_distribution(r_bins: np.ndarray, particles: np.ndarray, ret_error: bool = False):
     """
@@ -144,3 +143,58 @@ def total_energy(particles: np.ndarray):
     # # TODO: i am pretty sure this is wrong
     pe = 0
     return ke + pe
+
+
+def particles_plot_3d(particles: np.ndarray, title: str = "Particle distribution (3D)"):
+    """
+    Plots a 3D scatter plot of a set of particles.
+    Assumes that the particles array has the shape:
+    - either 4 columns: x, y, z, m
+    - or 7 columns: x, y, z, vx, vy, vz, m
+    Colormap is the mass of the particles.
+    """
+    if particles.shape[1] == 4:
+        x, y, z, m = particles[:, 0], particles[:, 1], particles[:, 2], particles[:, 3]
+        c = m
+    elif particles.shape[1] == 7:
+        x, y, z, m = particles[:, 0], particles[:, 1], particles[:, 2], particles[:, 6]
+        c = m
+    else:
+        raise ValueError("Particles array must have 4 or 7 columns")
+    
+    fig = plt.figure()
+    plt.title(title)
+    ax = fig.add_subplot(111, projection='3d')
+    ax.scatter(particles[:,0], particles[:,1], particles[:,2], cmap='viridis', c=particles[:,3])
+    plt.show()
+    logger.debug("3D scatter plot with mass colormap")
+
+
+def particles_plot_2d(particles: np.ndarray, title: str = "Flattened distribution (along z)"):
+    """
+    Plots a 2 colormap of a set of particles, flattened in the z direction.
+    Assumes that the particles array has the shape:
+    - either 4 columns: x, y, z, m
+    - or 7 columns: x, y, z, vx, vy, vz, m
+    """
+    if particles.shape[1] == 4:
+        x, y, z, m = particles[:, 0], particles[:, 1], particles[:, 2], particles[:, 3]
+        c = m
+    elif particles.shape[1] == 7:
+        x, y, z, m = particles[:, 0], particles[:, 1], particles[:, 2], particles[:, 6]
+        c = m
+    else:
+        raise ValueError("Particles array must have 4 or 7 columns")
+    
+    # plt.figure()
+    # plt.title(title)
+    # plt.scatter(x, y, c=range(particles.shape[0]))
+    # plt.colorbar()
+    # plt.show()
+
+    # or as a discrete heatmap
+    plt.figure()
+    plt.title(title)
+    plt.hist2d(x, y, bins=100, cmap='viridis')
+    plt.colorbar()
+    plt.show()