n square solver orking

nbody/utils/forces_basic.py

@@ -23,23 +23,24 @@ def n_body_forces(particles: np.ndarray, G: float, softening: float = 0):
         x_current = x_vec[i, :]
         m_current = masses[i]
 
-        # first compute the displacement to all other particles
+        # first compute the displacement to all other particles (and its magnitude)
-        displacements = x_vec - x_current
+        r_vec = x_vec - x_current
-        # and its magnitude
+        r = np.linalg.norm(r_vec, axis=1)
-        r = np.linalg.norm(displacements, axis=1)
+
         # add softening to the denominator
         r_adjusted = r**2 + softening**2
-        # the numerator is tricky:
+        # usually with a square root: r' = sqrt(r^2 + softening^2) and then cubed, but we combine that below
-        # 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
 
+        # the numerator is tricky:
+        # m is a list of scalars and r_vec is a list of vectors (2D array)
+        # we only want row_wise multiplication
+        num = G * (masses * r_vec.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
+        f = - np.sum((num.T / r_adjusted**1.5).T, axis=0) * m_current
-        forces[i] = - f
+        forces[i] = f
 
         if i!= 0 and i % 5000 == 0:
             logger.debug(f"Particle {i} done")
@@ -66,6 +67,7 @@ def n_body_forces_basic(particles: np.ndarray, G: float, softening: float = 0):
 
     return forces
 
+
 def analytical_forces(particles: np.ndarray):
     """
     Computes the interparticle forces without computing the n^2 interactions.

nbody/utils/forces_mesh.py

@@ -90,7 +90,7 @@ def mesh_forces_v2(particles: np.ndarray, G: float, n_grid: int, mapping: callab
         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]] / 10
+        forces[i] = - p[3] * phi_grad[..., ijk[0], ijk[1], ijk[2]]
         # TODO remove factor of 10
         # TODO could also index phi_grad the other way around?
 
@@ -141,13 +141,11 @@ def to_mesh(particles: np.ndarray, n_grid: int, mapping: callable) -> tuple[np.n
 
     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
 
 
@@ -157,28 +155,31 @@ def particle_to_cells_nn(particle, axis):
     # the weight is obviously 1
     return [ijk], [1]
 
+bbox = np.array([
+    [1, 0, 0],
+    [-1, 0, 0],
+    [1, 1, 0],
+    [-1, -1, 0],
+    [1, 1, 1],
+    [-1, -1, 1],
+    [1, 1, -1],
+    [-1, -1, -1]
+])
 
 def particle_to_cells_cic(particle, axis, width):
-    # create a virtual cell around the particle
+    # create a virtual cell around the particle and check the intersections
-    cell_bounds = [
+    bounding_cell = particle + width * bbox
-        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:
+    for b in bounding_cell:
+        # TODO: this is not the correct weight
         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)

nbody/utils/integrate.py

@@ -15,10 +15,10 @@ 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 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')
+    particles = particles.flatten()
     logger.debug(f"Reshaped 7 columns into {particles.shape=}")
 
     def f(y, t):
@@ -26,45 +26,45 @@ def ode_setup(particles: np.ndarray, force_function: callable) -> tuple[np.ndarr
         Computes the right hand side of the ODE system.
         The ODE system is linearized around the current positions and velocities.
         """
-        y = to_particles(y)
+        p = to_particles(y)
-        # now y has shape (n, 7), with columns x, y, z, vx, vy, vz, m
+        # this is explicitly a copy, which has shape (n, 7)
+        # columns x, y, z, vx, vy, vz, m
+        # (need to keep y intact since integrators make multiple function calls)
 
-        forces = force_function(y[:, [0, 1, 2, -1]])
+        forces = force_function(p[:, [0, 1, 2, -1]])
 
         # compute the accelerations
-        masses = y[:, -1]
+        masses = p[:, -1]
         a = forces / masses[:, None]
         # the [:, None] is to force broadcasting in order to divide each row of forces by the corresponding mass
 
-        dydt = np.zeros_like(y)
         # the position columns become the velocities
         # the velocity columns become the accelerations
-        dydt[:, 0:3] = y[:, 3:6]
+        p[:, 0:3] = p[:, 3:6]
-        dydt[:, 3:6] = a
+        p[:, 3:6] = a
         # the masses remain unchanged
-        dydt[:, -1] = masses
+        # p[:, -1] = p[:, -1]
 
         # flatten the array again
         # logger.debug(f"As particles: {y}")
-        dydt = dydt.reshape(-1, copy=False, order='A')
+        p = p.reshape(-1, copy=False)
         
         # logger.debug(f"As column: {y}")
-        return dydt
+        return p
 
     return particles, f
 
 
 def to_particles(y: np.ndarray) -> np.ndarray:
     """
-    Converts the 1D array y into a 2D array
+    Converts the 1D array y into a 2D array, by creating a copy
     The new shape is (n, 7) where n is the number of particles.
     The columns are x, y, z, vx, vy, vz, m
     """
     if y.size % 7 != 0:
         raise ValueError("The array y should be inflatable to 7 columns")
-    
+
-    n = y.size // 7
+    y = y.reshape((-1, 7), copy=True)
-    y = y.reshape((n, 7), copy=False, order='F')
     # logger.debug(f"Unflattened array into {y.shape=}")
     return y
 
@@ -76,7 +76,7 @@ def to_particles_3d(y: np.ndarray) -> np.ndarray:
     """
     n_steps = y.shape[0]
     n_particles = y.shape[1] // 7
-    y = y.reshape((n_steps, n_particles, 7), copy=False, order='F')
+    y = y.reshape((n_steps, n_particles, 7))
     # logger.debug(f"Unflattened array into {y.shape=}")
     return y
 

nbody/utils/particles.py

@@ -14,7 +14,7 @@ def density_distribution(r_bins: np.ndarray, particles: np.ndarray, ret_error: b
 
     m = particles[:, 3]
     r = np.linalg.norm(particles[:, :3], axis=1)
-    
+
     m_shells = np.zeros_like(r_bins)
     v_shells = np.zeros_like(r_bins)
     error_relative = np.zeros_like(r_bins)
@@ -27,6 +27,7 @@ def density_distribution(r_bins: np.ndarray, particles: np.ndarray, ret_error: b
         if ret_error:
             count = np.count_nonzero(mask)
             if count > 0:
+                # the absolute error is the square root of the number of particles
                 error_relative[i] = 1 / np.sqrt(count)
             else:
                 error_relative[i] = 0