starting on the presentation

nbody/utils/solver_mesh.py

@@ -0,0 +1,89 @@
+## Implementation of a mesh based full solver with boundary conditions etc.
+import numpy as np
+from . import mesh_forces
+import logging
+logger = logging.getLogger(__name__)
+
+
+def mesh_solver(
+        particles: np.ndarray,
+        G: float,
+        mapping: callable,
+        n_grid: int,
+        bounds: tuple = (-1, 1),
+        boundary: str = "vanishing",
+    ) -> np.ndarray:
+    """
+    Computes the gravitational force acting on a set of particles using a mesh-based approach. The mesh is of fixed size: n_grid x n_grid x n_grid spanning the given bounds. Particles reaching the boundary are treated according to the boundary condition.
+    Args:
+    - particles: np.ndarray, shape=(n, 4). Assumes that the particles array has the following columns: x, y, z, m.
+    - G: float, the gravitational constant.
+    - mapping: callable, the mapping function to use.
+    - n_grid: int, the number of grid points in each direction.
+    - bounds: tuple, the bounds of the mesh.
+    - boundary: str, the boundary condition to apply.
+    """
+    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 [mapping={mapping.__name__}, {n_grid=}]")
+
+    # the function is fine, let's abuse it somewhat
+    
+    axis = np.linspace(bounds[0], bounds[1], n_grid)
+    mesh = np.zeros((n_grid, n_grid, n_grid))
+    spacing = np.abs(axis[1] - axis[0])
+    logger.debug(f"Using mesh spacing: {spacing}")
+
+
+    # Check that the boundary condition is fullfilled
+    if boundary == "periodic":
+        raise NotImplementedError("Periodic boundary conditions are not implemented yet")
+    elif boundary == "vanishing":
+        # remove the particles that are outside the mesh
+        outlier_mask = particles[:, :3] < bounds[0] | particles[:, :3] > bounds[1]
+
+        if np.any(outlier_mask):
+            logger.warning(f"Removing {np.sum(outlier_mask)} particles that are outside the mesh")
+            particles = particles[~outlier_mask]
+            logger.debug(f"New particles shape: {particles.shape}")
+    else:
+        raise ValueError(f"Unknown boundary condition: {boundary}")
+
+
+    if logger.isEnabledFor(logging.DEBUG):
+        mesh_forces.show_mesh_information(mesh, "Density mesh")
+
+    # compute the potential and its gradient
+    phi_grad = mesh_forces.mesh_poisson(rho, G, spacing)
+
+    if logger.isEnabledFor(logging.DEBUG):
+        logger.debug(f"Got phi_grad with: {phi_grad.shape}, {np.max(phi_grad)}")
+        mesh_forces.show_mesh_information(phi_grad[0], "Potential gradient (x-direction)")
+
+    # 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
+        logger.debug(f"Particle {p} maps to cell {ijk}")
+        # 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 particles_to_mesh(particles: np.ndarray, mesh: np.ndarray, axis: np.ndarray, mapping: callable) -> None:
+    """
+    Maps a list of particles to an existing mesh, filling it inplace
+    """
+    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
+    for p in particles:
+        m = p[-1]
+        # 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