probably faulty but running at least

nbody/utils/forces_basic.py

@@ -33,13 +33,13 @@ def n_body_forces(particles: np.ndarray, G: float, softening: float = 0):
         # 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
+        forces[i] = - f
 
         if i % 5000 == 0:
             logger.debug(f"Particle {i} done")
@@ -47,6 +47,24 @@ def n_body_forces(particles: np.ndarray, G: float, softening: float = 0):
     return forces
 
 
+def n_body_forces_basic(particles: np.ndarray, G: float, softening: float = 0):
+    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))
+    for i in range(n):
+        for j in range(n):
+            if i == j:
+                continue # keep the value at zero
+            r_vec = x_vec[j] - x_vec[i]
+            r = np.linalg.norm(r_vec)
+            f = - G * masses[i] * masses[j] * r_vec / (r**3 + softening**3)
+            forces[i] += f
+
+    return forces
 
 def analytical_forces(particles: np.ndarray):
     """

nbody/utils/forces_mesh.py

@@ -75,7 +75,7 @@ def mesh_forces_v2(particles: np.ndarray, G: float, n_grid: int, mapping: callab
     if logger.level >= logging.DEBUG:
         show_mesh_information(phi, "Potential mesh")
         show_mesh_information(phi_grad[0], "Potential gradient")
-    logger.debug(f"Got phi_grad with: {phi_grad.shape}, {np.max(phi_grad)}")
+        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])

nbody/utils/integrate.py

@@ -16,10 +16,9 @@ def ode_setup(particles: np.ndarray, force_function: callable) -> tuple[np.ndarr
     if particles.shape[1] != 7:
         raise ValueError("Particles array must have 7 columns: x, y, z, vx, vy, vz, m")
     
-    # for scipy integrators we need to flatten array which contains 7 columns for now
-    # we don't really care how we reshape as long as we unflatten consistently afterwards
+    # for the integrators we need to flatten array which contains 7 columns for now
+    # we don't really care how we reshape as long as we unflatten consistently
     particles = particles.reshape(-1, copy=False, order='A')
-    # this is consistent with the unflattening in to_particles()!
     logger.debug(f"Reshaped 7 columns into {particles.shape=}")
 
     def f(y, t):
@@ -29,33 +28,35 @@ def ode_setup(particles: np.ndarray, force_function: callable) -> tuple[np.ndarr
         """
         y = to_particles(y)
         # now y has shape (n, 7), with columns x, y, z, vx, vy, vz, m
-        
 
         forces = force_function(y[:, [0, 1, 2, -1]])
-        
+
         # compute the accelerations
         masses = y[:, -1]
         a = forces / masses[:, None]
         # the [:, None] is to force broadcasting in order to divide each row of forces by the corresponding mass
-        # a.flatten()
-        
-        # replace some values in y:
+
+        dydt = np.zeros_like(y)
         # the position columns become the velocities
         # the velocity columns become the accelerations
-        y[:, 0:3] = y[:, 3:6]
-        y[:, 3:6] = a
+        dydt[:, 0:3] = y[:, 3:6]
+        dydt[:, 3:6] = a
         # the masses remain unchanged
+        dydt[:, -1] = masses
 
         # flatten the array again
-        y = y.reshape(-1, copy=False, order='A')
-        return y
+        # logger.debug(f"As particles: {y}")
+        dydt = dydt.reshape(-1, copy=False, order='A')
+        
+        # logger.debug(f"As column: {y}")
+        return dydt
 
     return particles, f
 
 
 def to_particles(y: np.ndarray) -> np.ndarray:
     """
-    Converts the 1D array y into a 2D array IN PLACE
+    Converts the 1D array y into a 2D array
     The new shape is (n, 7) where n is the number of particles.
     The columns are x, y, z, vx, vy, vz, m
     """
@@ -64,10 +65,21 @@ def to_particles(y: np.ndarray) -> np.ndarray:
     
     n = y.size // 7
     y = y.reshape((n, 7), copy=False, order='F')
-    logger.debug(f"Unflattened array into {y.shape=}")
+    # logger.debug(f"Unflattened array into {y.shape=}")
     return y
 
 
+def to_particles_3d(y: np.ndarray) -> np.ndarray:
+    """
+    Converts the 2D sol array with one vector per timestep into a 3D array:
+    2d particles (nx7) x nsteps
+    """
+    n_steps = y.shape[0]
+    n_particles = y.shape[1] // 7
+    y = y.reshape((n_steps, n_particles, 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)

nbody/utils/logging_config.py

@@ -2,10 +2,10 @@ import logging
 
 logging.basicConfig(
     ## set logging level
-    # level=logging.INFO,
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(message)s',
-    datefmt='%H:%M:%S'
+    level = logging.DEBUG,
+    # level = logging.INFO,
+    format = '%(asctime)s - %(name)s - %(message)s',
+    datefmt = '%H:%M:%S'
 )
 
 
@@ -13,4 +13,4 @@ logging.basicConfig(
 logging.getLogger('matplotlib.font_manager').setLevel(logging.WARNING)
 logging.getLogger('matplotlib.ticker').setLevel(logging.WARNING)
 logging.getLogger('matplotlib.pyplot').setLevel(logging.WARNING)
-logging.getLogger('matplotlib.colorbar').setLevel(logging.WARNING)
+logging.getLogger('matplotlib.colorbar').setLevel(logging.WARNING)

nbody/utils/model.py

@@ -11,4 +11,4 @@ def model_density_distribution(r_bins: np.ndarray, M: float = 5, a: float = 5) -
     See https://doi.org/10.1086%2F168845 for more information.
     """
     rho = M / (2 * np.pi) * a / (r_bins * (r_bins + a)**3)
-    return rho
+    return rho