293 lines
11 KiB
Python
293 lines
11 KiB
Python
## Implementation of a mesh based force solver
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from scipy import fft
|
|
|
|
import logging
|
|
logger = logging.getLogger(__name__)
|
|
|
|
'''
|
|
def mesh_forces(particles: np.ndarray, G: float, n_grid: int, mapping: callable) -> np.ndarray:
|
|
"""
|
|
Computes the gravitational force acting on a set of particles using a mesh-based approach.
|
|
Assumes that the particles array has the following columns: x, y, z, m.
|
|
"""
|
|
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=}]")
|
|
|
|
# in this case we create an adaptively sized mesh containing all particles
|
|
max_pos = np.max(np.abs(particles[:, :3]))
|
|
mesh, axis, spacing = create_mesh(-max_pos, max_pos, n_grid)
|
|
|
|
fill_mesh(particles, mesh, axis, mapping)
|
|
# we want a density mesh:
|
|
cell_volume = spacing**3
|
|
rho = mesh / cell_volume
|
|
|
|
if logger.isEnabledFor(logging.DEBUG):
|
|
show_mesh_information(mesh, "Density mesh")
|
|
|
|
# compute the potential and its gradient
|
|
phi_grad = mesh_poisson(rho, G, spacing)
|
|
|
|
if logger.isEnabledFor(logging.DEBUG):
|
|
logger.debug(f"Got phi_grad with: {phi_grad.shape}, {np.max(phi_grad)}")
|
|
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 mesh_poisson(mesh: np.ndarray, G: float, spacing: float) -> np.ndarray:
|
|
"""
|
|
Solves the poisson equation for the mesh using the FFT.
|
|
Returns the derivative of the potential - grad phi
|
|
"""
|
|
rho_hat = fft.fftn(mesh)
|
|
k = fft.fftfreq(mesh.shape[0], spacing) * (2 * np.pi)
|
|
# shift the zero frequency to the center
|
|
|
|
kx, ky, kz = np.meshgrid(k, k, k)
|
|
k_vec = np.stack([kx, ky, kz], axis=0)
|
|
k_sr = kx**2 + ky**2 + kz**2
|
|
if logger.isEnabledFor(logging.DEBUG):
|
|
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")
|
|
|
|
k_sr[k_sr == 0] = np.inf
|
|
k_inv = k_vec / k_sr # allows for element-wise division
|
|
|
|
logger.debug(f"Proceeding to poisson equation with {rho_hat.shape=}, {k_inv.shape=}")
|
|
grad_phi_hat = - 4 * np.pi * G * rho_hat * k_inv * 1j
|
|
# nabla^2 phi => -i * k * nabla phi = 4 pi G rho => nabla phi = - i * rho * k / k^2
|
|
grad_phi = np.real(fft.ifftn(grad_phi_hat))
|
|
return grad_phi
|
|
|
|
'''
|
|
|
|
|
|
def mesh_forces(particles: np.ndarray, G: float = 1, n_grid: int = 50, mapping: callable = None) -> np.ndarray:
|
|
"""
|
|
Computes the gravitational force acting on a set of particles using a mesh-based approach.
|
|
Assumes that the particles array has the following columns: x, y, z, m.
|
|
"""
|
|
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=}]")
|
|
|
|
# in this case we create an adaptively sized mesh containing all particles
|
|
max_pos = np.max(np.abs(particles[:, :3]))
|
|
mesh, axis, spacing = create_mesh(-max_pos, max_pos, n_grid)
|
|
|
|
fill_mesh(particles, mesh, axis, mapping)
|
|
# we want a density mesh:
|
|
cell_volume = spacing**3
|
|
rho = mesh / cell_volume
|
|
|
|
if logger.isEnabledFor(logging.DEBUG):
|
|
show_mesh_information(mesh, "Density mesh")
|
|
|
|
# compute the potential and its gradient
|
|
phi = mesh_poisson(rho, G, spacing)
|
|
if logger.isEnabledFor(logging.DEBUG):
|
|
logger.debug(f"Got phi with: {phi.shape}, {np.max(phi)}")
|
|
show_mesh_information(phi, "Potential")
|
|
|
|
# get the acceleration from finite differences of the potential
|
|
# a = - grad phi
|
|
ax, ay, az = np.gradient(phi, spacing)
|
|
a_vec = - np.stack([ax, ay, az], axis=0)
|
|
|
|
|
|
# compute the particle forces from the mesh potential
|
|
forces = np.zeros_like(particles[:, :3])
|
|
ijks = np.digitize(particles[:, :3], axis) - 1
|
|
|
|
for i in range(particles.shape[0]):
|
|
m = particles[i, 3]
|
|
idx = ijks[i]
|
|
# f = m * a
|
|
forces[i] = m * a_vec[..., idx[0], idx[1], idx[2]]
|
|
|
|
return forces
|
|
|
|
|
|
def mesh_poisson(mesh: np.ndarray, G: float, spacing: float) -> np.ndarray:
|
|
"""
|
|
Solves the poisson equation for the mesh using the FFT.
|
|
Returns the the potential - grad
|
|
"""
|
|
rho_hat = fft.fftn(mesh)
|
|
|
|
# we also need the wave numbers
|
|
k = fft.fftfreq(mesh.shape[0], spacing) * (2 * np.pi)
|
|
# assuming the grid is cubic
|
|
kx, ky, kz = np.meshgrid(k, k, k)
|
|
k_sr = kx**2 + ky**2 + kz**2
|
|
|
|
if logger.isEnabledFor(logging.DEBUG):
|
|
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")
|
|
|
|
k_sr[k_sr == 0] = 1e-10
|
|
# k_inv = k_vec / k_sr # allows for element-wise division
|
|
|
|
phi_hat = - 4 * np.pi * G * rho_hat / k_sr
|
|
# nabla^2 phi becomes -i * k * nabla phi_hat = 4 pi G rho_hat
|
|
# => nabla phi = - i * rho * k / k^2
|
|
phi = np.real(fft.ifftn(phi_hat))
|
|
return phi
|
|
|
|
|
|
|
|
#### Helper functions for star mapping
|
|
def create_mesh(min_pos: float, max_pos: float, n_grid: int) -> tuple[np.ndarray, np.ndarray, float]:
|
|
"""
|
|
Creates an empty 3D mesh with the given dimensions.
|
|
Returns the mesh, the axis and the spacing between the cells.
|
|
"""
|
|
axis = np.linspace(min_pos, max_pos, n_grid)
|
|
mesh = np.zeros((n_grid, n_grid, n_grid))
|
|
spacing = np.diff(axis)[0]
|
|
logger.debug(f"Using mesh spacing: {spacing}")
|
|
return mesh, axis, spacing
|
|
|
|
|
|
def fill_mesh(particles: np.ndarray, mesh: np.ndarray, axis: np.ndarray, mapping: callable):
|
|
"""
|
|
Maps a list of particles to a the mesh (in place)
|
|
Assumes that the particles array has the following columns: x, y, z, ..., m.
|
|
Uses the mapping function to detemine the contribution to each cell. The mapped density should be normalized to 1.
|
|
"""
|
|
if particles.shape[1] < 4:
|
|
raise ValueError("Particles array must have at least 4 columns: x, y, z, ..., m")
|
|
|
|
# each particle will have its particular contirbution (determined through a weight function, mapping)
|
|
for i in range(particles.shape[0]):
|
|
p = particles[i]
|
|
mapping(mesh, p, axis) # this directly adds to the mesh
|
|
|
|
|
|
def particle_mapping_nn(mesh_to_fill: np.ndarray, particle: np.ndarray, axis: np.ndarray):
|
|
# fills the mesh in place with the particle mass
|
|
ijk = np.digitize(particle, axis) - 1
|
|
mesh_to_fill[ijk[0], ijk[1], ijk[2]] += particle[3]
|
|
|
|
|
|
def particle_mapping_cic(mesh_to_fill: np.ndarray, particle: np.ndarray, axis: np.ndarray):
|
|
# fills the mesh in place with the particle mass
|
|
ijk = np.digitize(particle, axis) - 1
|
|
spacing = axis[1] - axis[0]
|
|
|
|
# generate a 3D map of all the distances to the particle
|
|
px, py, pz = np.meshgrid(axis, axis, axis, indexing='ij')
|
|
dist = np.linalg.norm([px - particle[0], py - particle[1], pz - particle[2]], axis=0)
|
|
|
|
# the weights are the inverse of the distance, cut off at the cell size
|
|
weights = np.maximum(0, 1 - dist / spacing)
|
|
mesh_to_fill += particle[3] * weights
|
|
|
|
|
|
|
|
#### Helper functions for mesh plotting
|
|
def show_mesh_information(mesh: np.ndarray, name: str):
|
|
logger.info(f"Mesh information for {name}")
|
|
logger.info(f"Total mapped mass: {np.sum(mesh):.0f}")
|
|
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)}")
|
|
mesh_plot_3d(mesh, name)
|
|
mesh_plot_2d(mesh, name)
|
|
|
|
|
|
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')
|
|
sc = ax.scatter(*np.where(mesh), c=mesh[np.where(mesh)], cmap='viridis')
|
|
plt.colorbar(sc, ax=ax, label='Density')
|
|
plt.show()
|
|
|
|
|
|
def mesh_plot_2d(mesh: np.ndarray, name: str, only_z: bool = False):
|
|
fig = plt.figure()
|
|
fig.suptitle(f"{name} - {mesh.shape}")
|
|
if only_z:
|
|
plt.imshow(np.sum(mesh, axis=2), cmap='viridis', origin='lower')
|
|
else:
|
|
axs = fig.subplots(1, 3)
|
|
axs[0].imshow(np.sum(mesh, axis=0), origin='lower')
|
|
axs[0].set_title("Flattened in x")
|
|
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()
|
|
|
|
|
|
|
|
##################################
|
|
# For the presentation - without logging
|
|
def mesh__forces(particles: np.ndarray, G: float = 1, n_grid: int = 50, mapping: callable = None) -> np.ndarray:
|
|
"""
|
|
Computes the gravitational force acting on a set of particles using a mesh-based approach.
|
|
Assumes that the particles array has the following columns: x, y, z, m.
|
|
"""
|
|
max_pos = np.max(np.abs(particles[:, :3]))
|
|
mesh, axis, spacing = create_mesh(-max_pos, max_pos, n_grid)
|
|
|
|
fill_mesh(particles, mesh, axis, mapping)
|
|
# we want a density mesh:
|
|
cell_volume = spacing**3
|
|
rho = mesh / cell_volume
|
|
|
|
# compute the potential and its gradient
|
|
phi = mesh_poisson(rho, G, spacing)
|
|
|
|
# get the acceleration from finite differences of the potential
|
|
ax, ay, az = np.gradient(phi, spacing)
|
|
a_vec = - np.stack([ax, ay, az], axis=0)
|
|
|
|
# compute the particle forces from the mesh potential
|
|
forces = np.zeros_like(particles[:, :3])
|
|
ijks = np.digitize(particles[:, :3], axis) - 1
|
|
|
|
for i in range(particles.shape[0]):
|
|
m = particles[i, 3]
|
|
idx = ijks[i]
|
|
forces[i] = m * a_vec[..., idx[0], idx[1], idx[2]]
|
|
|
|
return forces
|
|
|
|
|
|
def mesh__poisson(mesh: np.ndarray, G: float, spacing: float) -> np.ndarray:
|
|
"""
|
|
Solves the poisson equation for the mesh using the FFT.
|
|
Returns the the potential - phi
|
|
"""
|
|
rho_hat = fft.fftn(mesh)
|
|
|
|
# we also need the wave numbers
|
|
k = fft.fftfreq(mesh.shape[0], spacing) * (2 * np.pi)
|
|
# assuming the grid is cubic
|
|
kx, ky, kz = np.meshgrid(k, k, k)
|
|
k_sr = kx**2 + ky**2 + kz**2
|
|
|
|
k_sr[k_sr == 0] = np.inf
|
|
|
|
phi_hat = - 4 * np.pi * G * rho_hat / k_sr
|
|
return np.real(fft.ifftn(phi_hat))
|