## 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))