finally accurate mesh forces
This commit is contained in:
parent
37a2687ffe
commit
da8a7d4574
@ -31,13 +31,13 @@
|
||||
// Finally - The real content
|
||||
= N-body forces and analytical solutions
|
||||
|
||||
== Objective
|
||||
Implement naive N-body force computation and get an intuition of the challenges:
|
||||
- accuracy
|
||||
- computation time
|
||||
- stability
|
||||
// == Objective
|
||||
// Implement naive N-body force computation and get an intuition of the challenges:
|
||||
// - accuracy
|
||||
// - computation time
|
||||
// - stability
|
||||
|
||||
$==>$ still useful to compute basic quantities of the system, but too limited for large systems or the dynamical evolution of the system
|
||||
// $==>$ still useful to compute basic quantities of the system, but too limited for large systems or the dynamical evolution of the system
|
||||
|
||||
|
||||
== Overview - the system
|
||||
@ -47,7 +47,7 @@ Get a feel for the particles and their distribution. [#link(<task1:plot_particle
|
||||
|
||||
#columns(2)[
|
||||
#helpers.image_cell(t1, "plot_particle_distribution")
|
||||
Note: for visibility the outer particles are not shown.
|
||||
// Note: for visibility the outer particles are not shown.
|
||||
#colbreak()
|
||||
The system at hand is characterized by:
|
||||
- $N ~ 10^4$ stars
|
||||
@ -81,12 +81,12 @@ We compare the computed density with the analytical model provided by the _Hernq
|
||||
#helpers.image_cell(t1, "plot_density_distribution")
|
||||
]
|
||||
)
|
||||
// Note that by construction, the first shell contains no particles
|
||||
// => the numerical density is zero there
|
||||
// Having more bins means to have shells that are nearly empty
|
||||
// => the error is large, NBINS = 30 is a good compromise
|
||||
|
||||
|
||||
#block(
|
||||
height: 1fr,
|
||||
)
|
||||
|
||||
|
||||
== Force computation
|
||||
// N Body and variations
|
||||
@ -110,13 +110,48 @@ We compare the computed density with the analytical model provided by the _Hernq
|
||||
]
|
||||
)
|
||||
|
||||
// basic $N^2$ matches analytical solution without dropoff. but: noisy data from "bad" samples
|
||||
// $N^2$ with softening matches analytical solution but has a dropoff. No noisy data.
|
||||
|
||||
// => softening $\approx 1 \varepsilon$ is a sweet spot since the dropoff is "late"
|
||||
|
||||
|
||||
== Relaxation
|
||||
Relaxation [#link(<task1:compute_relaxation_time>)[code]]:
|
||||
// #helpers.code_cell(t1, "compute_relaxation_time")
|
||||
|
||||
We express system relaxation in terms of the dynamical time of the system.
|
||||
$
|
||||
t_"relax" = overbrace(N / (8 log N), n_"relax") dot t_"crossing"
|
||||
$
|
||||
where the crossing time of the system can be estimated through the half-mass velocity $t_"crossing" = v(r_"hm")/r_"hm"$.
|
||||
|
||||
We find a relaxation of [#link(<task1:compute_relaxation_time>)[code]].
|
||||
|
||||
|
||||
Discussion!
|
||||
// === Discussion
|
||||
#grid(
|
||||
columns: (1fr, 1fr),
|
||||
inset: 0.5em,
|
||||
block[
|
||||
#image("relaxation.png")
|
||||
],
|
||||
block[
|
||||
- Each star-star interaction contributes $delta v approx (2 G m )/b$
|
||||
- Shifting by $epsilon$ *dampens* each contribution
|
||||
- $=>$ relaxation time increases
|
||||
]
|
||||
)
|
||||
// The estimate for $n_{relax}$ comes from the contribution of each star-star encounter to the velocity dispersion. This depends on the perpendicular force
|
||||
|
||||
// $\implies$ a bigger softening length leads to a smaller $\delta v$.
|
||||
|
||||
// Using $n_{relax} = \frac{v^2}{\delta v^2}$, and knowing that the value of $v^2$ is derived from the Virial theorem (i.e. unaffected by the softening length), we can see that $n_{relax}$ should increase with $\varepsilon$.
|
||||
|
||||
// === Effect
|
||||
// - The relaxation time **increases** with increasing softening length
|
||||
// - From the integration over all impact parameters $b$ even $b_{min}$ is chosen to be larger than $\varepsilon$ $\implies$ expect only a small effect on the relaxation time
|
||||
|
||||
// **In other words:**
|
||||
// The softening dampens the change of velocity => time to relax is longer
|
||||
|
||||
|
||||
|
||||
@ -128,10 +163,13 @@ Discussion!
|
||||
columns: 2
|
||||
)[
|
||||
#helpers.image_cell(t2, "plot_particle_distribution")
|
||||
|
||||
$=>$ use $M_"sys" approx 10^4 M_"sol" + M_"BH"$
|
||||
]
|
||||
|
||||
|
||||
|
||||
|
||||
== Force computation
|
||||
#helpers.code_reference_cell(t2, "function_mesh_force")
|
||||
|
||||
@ -156,9 +194,16 @@ Discussion!
|
||||
- very large grids have issues with overdiscretization
|
||||
|
||||
$==> 75 times 75 times 75$ as a good compromise
|
||||
// Some other comments:
|
||||
// - see the artifacts because of the even grid numbers (hence the switch to 75)
|
||||
// overdiscretization for large grids -> vertical spread even though r is constant
|
||||
// this becomes even more apparent when looking at the data without noise - the artifacts remain
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
#helpers.image_cell(t2, "plot_force_computation_time")
|
||||
|
||||
== Time integration
|
||||
=== Runge-Kutta
|
||||
#helpers.code_reference_cell(t2, "function_runge_kutta")
|
||||
|
||||
BIN
nbody/presentation/relaxation.png
Normal file
BIN
nbody/presentation/relaxation.png
Normal file
Binary file not shown.
|
After 
(image error) Size: 14 KiB |
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@ -9,6 +9,7 @@ from .units import *
|
||||
from .forces_basic import *
|
||||
from .forces_tree import *
|
||||
from .forces_mesh import *
|
||||
from .forces_cache import *
|
||||
|
||||
# Helpers for solving the IVP and having time evolution
|
||||
from .integrate import *
|
||||
|
||||
@ -3,7 +3,7 @@ import logging
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def n_body_forces(particles: np.ndarray, G: float, softening: float = 0):
|
||||
def n_body_forces(particles: np.ndarray, G: float = 1, softening: float = 0):
|
||||
"""
|
||||
Computes the gravitational forces between a set of particles.
|
||||
Assumes that the particles array has the following columns: x, y, z, m.
|
||||
|
||||
33
nbody/utils/forces_cache.py
Normal file
33
nbody/utils/forces_cache.py
Normal file
@ -0,0 +1,33 @@
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import timeit
|
||||
import logging
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def cached_forces(cache_path: Path, particles: np.ndarray, force_function:callable, func_kwargs: dict):
|
||||
"""
|
||||
Tries to load the forces from a cache file. If that fails, computes the forces using the provided function.
|
||||
"""
|
||||
cache_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
n_particles = particles.shape[0]
|
||||
|
||||
kwargs_str = "_".join([f"{k}_{v}" for k, v in func_kwargs.items()])
|
||||
|
||||
force_cache = cache_path / f"forces__{force_function.__name__}__n_{n_particles}__kwargs_{kwargs_str}.npy"
|
||||
time_cache = cache_path / f"time__{force_function.__name__}__n_{n_particles}__kwargs_{kwargs_str}.npy"
|
||||
|
||||
if force_cache.exists() and time_cache.exists():
|
||||
force = np.load(force_cache)
|
||||
logger.info(f"Loaded forces from {force_cache}")
|
||||
time = np.load(time_cache)
|
||||
else:
|
||||
force = force_function(particles, **func_kwargs)
|
||||
np.save(force_cache, force)
|
||||
time = 0
|
||||
np.info(f"Timing {force_function.__name__} for {n_particles} particles")
|
||||
time = timeit.timeit(lambda: force_function(particles, **func_kwargs), number=10)
|
||||
np.save(time_cache, time)
|
||||
|
||||
return force, time
|
||||
@ -6,8 +6,7 @@ from scipy import fft
|
||||
import logging
|
||||
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:
|
||||
"""
|
||||
Computes the gravitational force acting on a set of particles using a mesh-based approach.
|
||||
@ -18,10 +17,11 @@ def mesh_forces(particles: np.ndarray, G: float, n_grid: int, mapping: callable)
|
||||
|
||||
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 = np.abs(axis[1] - axis[0])
|
||||
logger.debug(f"Using mesh spacing: {spacing}")
|
||||
# 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
|
||||
@ -54,9 +54,8 @@ def mesh_poisson(mesh: np.ndarray, G: float, spacing: float) -> np.ndarray:
|
||||
Returns the derivative of the potential - grad phi
|
||||
"""
|
||||
rho_hat = fft.fftn(mesh)
|
||||
k = fft.fftfreq(mesh.shape[0], spacing)
|
||||
k = fft.fftfreq(mesh.shape[0], spacing) * (2 * np.pi)
|
||||
# shift the zero frequency to the center
|
||||
k = fft.fftshift(k)
|
||||
|
||||
kx, ky, kz = np.meshgrid(k, k, k)
|
||||
k_vec = np.stack([kx, ky, kz], axis=0)
|
||||
@ -72,13 +71,14 @@ def mesh_poisson(mesh: np.ndarray, G: float, spacing: float) -> np.ndarray:
|
||||
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
|
||||
# todo: check minus
|
||||
# TODO: check minus
|
||||
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:
|
||||
|
||||
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.
|
||||
@ -88,126 +88,120 @@ 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 = np.abs(axis[1] - axis[0])
|
||||
logger.debug(f"Using mesh spacing: {spacing}")
|
||||
# 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_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)
|
||||
phi = 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, "Potential mesh")
|
||||
show_mesh_information(phi_grad[0], "Potential gradient (x-direction)")
|
||||
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])
|
||||
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]]
|
||||
# TODO remove factor of 10
|
||||
# TODO could also index phi_grad the other way around?
|
||||
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_v2(mesh: np.ndarray, G: float, spacing: float) -> np.ndarray:
|
||||
def mesh_poisson(mesh: np.ndarray, G: float, spacing: float) -> np.ndarray:
|
||||
"""
|
||||
Solves the poisson equation for the mesh using the FFT.
|
||||
Returns the scalar potential.
|
||||
Returns the the potential - grad
|
||||
"""
|
||||
rho_hat = fft.fftn(mesh)
|
||||
k = fft.fftfreq(mesh.shape[0], spacing)
|
||||
# shift the zero frequency to the center
|
||||
k = fft.fftshift(k)
|
||||
|
||||
|
||||
# we also need the wave numbers
|
||||
spacing_3d = np.linalg.norm([spacing, spacing, spacing])
|
||||
k = fft.fftfreq(mesh.shape[0], spacing) * (2 * np.pi)
|
||||
# TODO: check if this is correct
|
||||
# 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")
|
||||
# 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=}")
|
||||
# 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=}")
|
||||
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
|
||||
# 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 to_mesh(particles: np.ndarray, n_grid: int, mapping: callable) -> tuple[np.ndarray, np.ndarray]:
|
||||
def create_mesh(min_pos: float, max_pos: float, n_grid: int) -> tuple[np.ndarray, np.ndarray, float]:
|
||||
"""
|
||||
Maps a list of particles to a of mesh of size n_grid x n_grid x n_grid.
|
||||
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.
|
||||
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")
|
||||
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)
|
||||
mesh_grid = np.meshgrid(axis, axis, axis)
|
||||
mesh = np.zeros_like(mesh_grid[0])
|
||||
|
||||
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
|
||||
|
||||
return mesh, axis
|
||||
# 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_to_cells_nn(particle, axis):
|
||||
# find the single cell that contains the particle
|
||||
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
|
||||
# the weight is obviously 1
|
||||
return [ijk], [1]
|
||||
mesh_to_fill[ijk[0], ijk[1], ijk[2]] += particle[3]
|
||||
|
||||
bbox = np.array([
|
||||
[1, 0, 0],
|
||||
[-1, 0, 0],
|
||||
[1, 1, 0],
|
||||
[-1, -1, 0],
|
||||
[1, 1, 1],
|
||||
[-1, -1, 1],
|
||||
[1, 1, -1],
|
||||
[-1, -1, -1]
|
||||
])
|
||||
|
||||
def particle_to_cells_cic(particle, axis, width):
|
||||
# create a virtual cell around the particle and check the intersections
|
||||
bounding_cell = particle + width * bbox
|
||||
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]
|
||||
|
||||
# find all the cells that intersect with the virtual cell
|
||||
ijks = []
|
||||
weights = []
|
||||
for b in bounding_cell:
|
||||
# TODO: this is not the correct weight
|
||||
w = np.linalg.norm(particle - b)
|
||||
ijk = np.digitize(b, axis) - 1
|
||||
ijks.append(ijk)
|
||||
weights.append(w)
|
||||
|
||||
# ensure that the weights sum to 1
|
||||
weights = np.array(weights)
|
||||
weights /= np.sum(weights)
|
||||
return ijks, weights
|
||||
# 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
|
||||
|
||||
|
||||
|
||||
|
||||
@ -163,8 +163,13 @@ def particles_plot_3d(positions: np.ndarray, masses: np.ndarray, title: str = "P
|
||||
fig = plt.figure()
|
||||
fig.suptitle(title)
|
||||
ax = fig.add_subplot(111, projection='3d')
|
||||
sc = ax.scatter(x, y, z, cmap='coolwarm', c=masses)
|
||||
cbar = plt.colorbar(sc, ax=ax, pad=0.1)
|
||||
|
||||
if np.all(masses == masses[0]):
|
||||
sc = ax.scatter(x, y, z, c='blue')
|
||||
cbar = plt.colorbar(sc, ax=ax, pad=0.1)
|
||||
else:
|
||||
sc = ax.scatter(x, y, z, cmap='coolwarm', c=masses)
|
||||
cbar = plt.colorbar(sc, ax=ax, pad=0.1)
|
||||
|
||||
try:
|
||||
cbar.set_label(f'Mass [{masses.unit:latex}]')
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
## Implementation of a mesh based full solver with boundary conditions etc.
|
||||
import numpy as np
|
||||
from . import mesh_forces
|
||||
from . import forces_mesh
|
||||
import logging
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -36,30 +36,38 @@ def mesh_solver(
|
||||
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]
|
||||
# # 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 = np.logical_or(particles[:, :3] < bounds[0], particles[:, :3] > bounds[1])
|
||||
# if np.any(outlier_mask):
|
||||
# idx = np.any(outlier_mask, axis=1)
|
||||
# logger.info(f"{idx.shape=}")
|
||||
# logger.warning(f"Removing {np.sum(idx)} particles that left the mesh")
|
||||
# # replace the particles by nan values
|
||||
# particles[idx, :] = np.nan
|
||||
# print(np.sum(np.isnan(particles)))
|
||||
# else:
|
||||
# raise ValueError(f"Unknown boundary condition: {boundary}")
|
||||
|
||||
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}")
|
||||
|
||||
# fill the mesh
|
||||
particles_to_mesh(particles, mesh, axis, mapping)
|
||||
# we want a density mesh:
|
||||
cell_volume = spacing**3
|
||||
rho = mesh / cell_volume
|
||||
|
||||
if logger.isEnabledFor(logging.DEBUG):
|
||||
mesh_forces.show_mesh_information(mesh, "Density mesh")
|
||||
forces_mesh.show_mesh_information(mesh, "Density mesh")
|
||||
|
||||
# compute the potential and its gradient
|
||||
phi_grad = mesh_forces.mesh_poisson(rho, G, spacing)
|
||||
phi_grad = forces_mesh.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)")
|
||||
forces_mesh.show_mesh_information(phi_grad[0], "Potential gradient (x-direction)")
|
||||
|
||||
# compute the particle forces from the mesh potential
|
||||
forces = np.zeros_like(particles[:, :3])
|
||||
@ -87,3 +95,56 @@ def particles_to_mesh(particles: np.ndarray, mesh: np.ndarray, axis: np.ndarray,
|
||||
ijks, weights = mapping(p, axis)
|
||||
for ijk, weight in zip(ijks, weights):
|
||||
mesh[ijk[0], ijk[1], ijk[2]] += weight * m
|
||||
|
||||
|
||||
'''
|
||||
#### Actually need to patch this
|
||||
def ode_setup(particles: np.ndarray, force_function: callable) -> tuple[np.ndarray, callable]:
|
||||
"""
|
||||
Linearizes the ODE system for the particles interacting gravitationally.
|
||||
Returns:
|
||||
- the Y0 array corresponding to the initial conditions (x0 and v0)
|
||||
- the function that computes the right hand side of the ODE with function signature f(t, y)
|
||||
Assumes that the particles array has the following columns: x, y, z, vx, vy, vz, m.
|
||||
"""
|
||||
if particles.shape[1] != 7:
|
||||
raise ValueError("Particles array must have 7 columns: x, y, z, vx, vy, vz, m")
|
||||
|
||||
# for the integrators we need to flatten array which contains 7 columns for now
|
||||
# we don't really care how we reshape as long as we unflatten consistently
|
||||
particles = particles.flatten()
|
||||
logger.debug(f"Reshaped 7 columns into {particles.shape=}")
|
||||
|
||||
def f(y, t):
|
||||
"""
|
||||
Computes the right hand side of the ODE system.
|
||||
The ODE system is linearized around the current positions and velocities.
|
||||
"""
|
||||
p = to_particles(y)
|
||||
# this is explicitly a copy, which has shape (n, 7)
|
||||
# columns x, y, z, vx, vy, vz, m
|
||||
# (need to keep y intact since integrators make multiple function calls)
|
||||
|
||||
forces = force_function(p[:, [0, 1, 2, -1]])
|
||||
|
||||
# compute the accelerations
|
||||
masses = p[:, -1]
|
||||
a = forces / masses[:, None]
|
||||
# the [:, None] is to force broadcasting in order to divide each row of forces by the corresponding mass
|
||||
|
||||
# the position columns become the velocities
|
||||
# the velocity columns become the accelerations
|
||||
p[:, 0:3] = p[:, 3:6]
|
||||
p[:, 3:6] = a
|
||||
# the masses remain unchanged
|
||||
# p[:, -1] = p[:, -1]
|
||||
|
||||
# flatten the array again
|
||||
# logger.debug(f"As particles: {y}")
|
||||
p = p.reshape(-1, copy=False)
|
||||
|
||||
# logger.debug(f"As column: {y}")
|
||||
return p
|
||||
|
||||
return particles, f
|
||||
'''
|
||||
Loading…
x
Reference in New Issue
Block a user