90 lines
3.7 KiB
Python
90 lines
3.7 KiB
Python
## 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
|