remoll/computational_astrophysics_nbody_solver
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

working nsquare integration?

Browse Source
This commit is contained in:
Remy Moll
2025-01-22 13:23:19 +01:00
parent b130f68a44
commit 8737441fbd
7 changed files with 244 additions and 4412 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
nbody/utils/forces_basic.py
View File

@@ -41,7 +41,7 @@ def n_body_forces(particles: np.ndarray, G: float, softening: float = 0):
f = np.sum((num.T / r_adjusted**1.5).T, axis=0) * m_current
forces[i] = - f
if i % 5000 == 0:
if i!= 0 and i % 5000 == 0:
logger.debug(f"Particle {i} done")
return forces

62
nbody/utils/forces_mesh.py
View File

@@ -8,7 +8,7 @@ logger = logging.getLogger(__name__)
#### Version 1 - keeping the derivative of phi
def mesh_forces(particles: np.ndarray, G: float, n_grid: int, mapping: callable) -> np.ndarray:
'''def mesh_forces(particles: np.ndarray, G: float, n_grid: int, mapping: callable) -> np.ndarray:
"""
Computes the gravitational forces between a set of particles using a mesh.
Assumes that the particles array has the following columns: x, y, z, m.
@@ -39,6 +39,8 @@ def mesh_poisson(mesh: np.ndarray, G: float) -> np.ndarray:
rho_hat = fft.fftn(mesh)
# the laplacian in fourier space takes the form of a multiplication
k = np.fft.fftfreq(mesh.shape[0])
# shift the zero frequency to the center
k = np.fft.fftshift(k)
# TODO: probably need to take the actual mesh bounds into account
kx, ky, kz = np.meshgrid(k, k, k)
k_vec = np.array([kx, ky, kz])
@@ -49,7 +51,7 @@ def mesh_poisson(mesh: np.ndarray, G: float) -> np.ndarray:
# the inverse fourier transform gives the potential (or its gradient)
grad_phi = np.real(fft.ifftn(grad_phi_hat))
return grad_phi
'''
#### Version 2 - only storing the scalar potential
def mesh_forces_v2(particles: np.ndarray, G: float, n_grid: int, mapping: callable) -> np.ndarray:
@@ -63,27 +65,34 @@ def mesh_forces_v2(particles: np.ndarray, G: float, n_grid: int, mapping: callab
logger.debug(f"Computing forces for {particles.shape[0]} particles using mesh [mapping={mapping.__name__}, {n_grid=}]")
mesh, axis = to_mesh(particles, n_grid, mapping)
spacing = axis[1] - axis[0]
logger.debug(f"Using mesh spacing: {spacing}")
# we want a density mesh:
cell_volume = spacing**3
rho = mesh / cell_volume
if logger.level >= logging.DEBUG:
show_mesh_information(mesh, "Density mesh")
# compute the potential and its gradient
spacing = axis[1] - axis[0]
logger.debug(f"Using mesh spacing: {spacing}")
phi = mesh_poisson_v2(mesh, G, spacing)
phi = mesh_poisson_v2(rho, G, spacing)
logger.debug(f"Got phi with: {phi.shape}, {np.max(phi)}")
phi_grad = np.stack(np.gradient(phi, spacing), axis=0)
if logger.level >= logging.DEBUG:
show_mesh_information(phi, "Potential mesh")
show_mesh_information(phi_grad[0], "Potential gradient")
show_mesh_information(phi_grad[0], "Potential gradient (x-direction)")
logger.debug(f"Got phi_grad with: {phi_grad.shape}, {np.max(phi_grad)}")
# 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
# 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]]
forces[i] = - p[3] * phi_grad[..., ijk[0], ijk[1], ijk[2]] / 10
# TODO remove factor of 10
# TODO could also index phi_grad the other way around?
return forces
@@ -95,13 +104,19 @@ def mesh_poisson_v2(mesh: np.ndarray, G: float, spacing: float) -> np.ndarray:
"""
rho_hat = fft.fftn(mesh)
k = fft.fftfreq(mesh.shape[0], spacing)
# shift the zero frequency to the center
k = np.fft.fftshift(k)
kx, ky, kz = np.meshgrid(k, k, k)
k_sr = kx**2 + ky**2 + kz**2
logger.debug(f"Got k_square with: {k_sr.shape}, {np.max(k_sr)} {np.min(k_sr)}")
logger.debug(f"Count of zeros: {np.sum(k_sr == 0)}")
if logger.level >= 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")
# avoid division by zero
# TODO: review this
k_sr[k_sr == 0] = np.inf
logger.debug(f"Proceeding to poisson equation with {rho_hat.shape=}, {k_sr.shape=}")
phi_hat = - 4 * np.pi * G * rho_hat / k_sr
# - comes from i squared
# TODO: 4pi stays since the backtransform removes the 1/2pi factor
@@ -116,6 +131,8 @@ def to_mesh(particles: np.ndarray, n_grid: int, mapping: callable) -> tuple[np.n
Assumes that the particles array has the following columns: x, y, z, ..., m.
Uses the mass of the particles and a smoothing function to detemine the contribution to each cell.
"""
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
max_pos = np.max(particles[:, :3])
axis = np.linspace(-max_pos, max_pos, n_grid)
@@ -176,11 +193,11 @@ def show_mesh_information(mesh: np.ndarray, name: str):
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)}")
plot_3d(mesh, name)
plot_2d(mesh, name)
mesh_plot_3d(mesh, name)
mesh_plot_2d(mesh, name)
def plot_3d(mesh: np.ndarray, name: str):
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')
@@ -188,14 +205,17 @@ def plot_3d(mesh: np.ndarray, name: str):
plt.show()
def plot_2d(mesh: np.ndarray, name: str):
def mesh_plot_2d(mesh: np.ndarray, name: str, only_z: bool = False):
fig = plt.figure()
fig.suptitle(f"{name} - {mesh.shape}")
axs = fig.subplots(1, 3)
axs[0].imshow(np.sum(mesh, axis=0))
axs[0].set_title("Flattened in x")
axs[1].imshow(np.sum(mesh, axis=1))
axs[1].set_title("Flattened in y")
axs[2].imshow(np.sum(mesh, axis=2))
axs[2].set_title("Flattened in z")
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()

14
nbody/utils/logging_config.py
View File

@@ -1,12 +1,12 @@
import logging
logging.basicConfig(
## set logging level
level = logging.DEBUG,
# level = logging.INFO,
format = '%(asctime)s - %(name)s - %(message)s',
datefmt = '%H:%M:%S'
)
def set_log_level(level: str):
logging.basicConfig(
level = logging.DEBUG if level == 'debug' else logging.INFO,
format = '%(asctime)s - %(name)s - %(message)s',
datefmt = '%H:%M:%S'
)
# silence some debug messages

58
nbody/utils/particles.py
View File

@@ -1,8 +1,7 @@
import numpy as np
import logging
from . import forces_basic
logger = logging.getLogger(__name__)
import matplotlib.pyplot as plt
def density_distribution(r_bins: np.ndarray, particles: np.ndarray, ret_error: bool = False):
"""
@@ -144,3 +143,58 @@ def total_energy(particles: np.ndarray):
# # TODO: i am pretty sure this is wrong
pe = 0
return ke + pe
def particles_plot_3d(particles: np.ndarray, title: str = "Particle distribution (3D)"):
"""
Plots a 3D scatter plot of a set of particles.
Assumes that the particles array has the shape:
- either 4 columns: x, y, z, m
- or 7 columns: x, y, z, vx, vy, vz, m
Colormap is the mass of the particles.
"""
if particles.shape[1] == 4:
x, y, z, m = particles[:, 0], particles[:, 1], particles[:, 2], particles[:, 3]
c = m
elif particles.shape[1] == 7:
x, y, z, m = particles[:, 0], particles[:, 1], particles[:, 2], particles[:, 6]
c = m
else:
raise ValueError("Particles array must have 4 or 7 columns")
fig = plt.figure()
plt.title(title)
ax = fig.add_subplot(111, projection='3d')
ax.scatter(particles[:,0], particles[:,1], particles[:,2], cmap='viridis', c=particles[:,3])
plt.show()
logger.debug("3D scatter plot with mass colormap")
def particles_plot_2d(particles: np.ndarray, title: str = "Flattened distribution (along z)"):
"""
Plots a 2 colormap of a set of particles, flattened in the z direction.
Assumes that the particles array has the shape:
- either 4 columns: x, y, z, m
- or 7 columns: x, y, z, vx, vy, vz, m
"""
if particles.shape[1] == 4:
x, y, z, m = particles[:, 0], particles[:, 1], particles[:, 2], particles[:, 3]
c = m
elif particles.shape[1] == 7:
x, y, z, m = particles[:, 0], particles[:, 1], particles[:, 2], particles[:, 6]
c = m
else:
raise ValueError("Particles array must have 4 or 7 columns")
# plt.figure()
# plt.title(title)
# plt.scatter(x, y, c=range(particles.shape[0]))
# plt.colorbar()
# plt.show()
# or as a discrete heatmap
plt.figure()
plt.title(title)
plt.hist2d(x, y, bins=100, cmap='viridis')
plt.colorbar()
plt.show()