202 lines
6.7 KiB
Python
202 lines
6.7 KiB
Python
import numpy as np
|
|
import logging
|
|
logger = logging.getLogger(__name__)
|
|
import matplotlib.pyplot as plt
|
|
|
|
def density_distribution(r_bins: np.ndarray, particles: np.ndarray, ret_error: bool = False):
|
|
"""
|
|
Computes the radial density distribution of a set of particles.
|
|
Assumes that the particles array has the following columns: x, y, z, m.
|
|
If ret_error is True, it will return the absolute error of the density.
|
|
"""
|
|
if particles.shape[1] != 4:
|
|
raise ValueError("Particles array must have 4 columns: x, y, z, m")
|
|
|
|
m = particles[:, 3]
|
|
r = np.linalg.norm(particles[:, :3], axis=1)
|
|
|
|
m_shells = np.zeros_like(r_bins)
|
|
v_shells = np.zeros_like(r_bins)
|
|
error_relative = np.zeros_like(r_bins)
|
|
r_bins = np.insert(r_bins, 0, 0)
|
|
|
|
for i in range(len(r_bins) - 1):
|
|
mask = (r >= r_bins[i]) & (r < r_bins[i + 1])
|
|
m_shells[i] = np.sum(m[mask])
|
|
v_shells[i] = 4/3 * np.pi * (r_bins[i + 1]**3 - r_bins[i]**3)
|
|
if ret_error:
|
|
count = np.count_nonzero(mask)
|
|
if count > 0:
|
|
# the absolute error is the square root of the number of particles
|
|
error_relative[i] = 1 / np.sqrt(count)
|
|
else:
|
|
error_relative[i] = 0
|
|
|
|
density = m_shells / v_shells
|
|
|
|
if ret_error:
|
|
return density, density * error_relative
|
|
else:
|
|
return density
|
|
|
|
|
|
|
|
def r_distribution(particles: np.ndarray):
|
|
"""
|
|
Computes the distribution of distances (to the origin) of a set of particles.
|
|
Assumes that the particles array has the following columns: x, y, z ...
|
|
"""
|
|
if particles.shape[1] < 3:
|
|
raise ValueError("Particles array must have at least 3 columns: x, y, z")
|
|
|
|
r = np.linalg.norm(particles[:, :3], axis=1)
|
|
return r
|
|
|
|
|
|
|
|
def remove_outliers(particles: np.ndarray, std_threshold: float = 3):
|
|
"""
|
|
Removes outliers from a set of particles.
|
|
Assumes that the particles array has the following columns: x, y, z ...
|
|
"""
|
|
if particles.shape[1] < 3:
|
|
raise ValueError("Particles array must have at least 3 columns: x, y, z")
|
|
|
|
r = np.linalg.norm(particles[:, :3], axis=1)
|
|
r_std = np.std(r)
|
|
r_mean = np.mean(r)
|
|
mask = np.abs(r - r_mean) < std_threshold * r_std
|
|
return particles[mask]
|
|
|
|
|
|
|
|
def mean_interparticle_distance(particles: np.ndarray):
|
|
"""
|
|
Computes the mean interparticle distance of a set of particles.
|
|
Assumes that the particles array has the following columns: x, y, z ...
|
|
"""
|
|
if particles.shape[1] < 3:
|
|
raise ValueError("Particles array must have at least 3 columns: x, y, z")
|
|
|
|
|
|
r_half_mass = half_mass_radius(particles)
|
|
r = np.linalg.norm(particles[:, :3], axis=1)
|
|
|
|
n_half_mass = np.sum(r < r_half_mass)
|
|
logger.debug(f"Number of particles within half mass radius: {n_half_mass} of {particles.shape[0]}")
|
|
|
|
rho = n_half_mass / (4/3 * np.pi * r_half_mass**3)
|
|
# the mean distance between particles is the inverse of the density
|
|
|
|
epsilon = (1 / rho)**(1/3)
|
|
logger.info(f"Found mean interparticle distance: {epsilon}")
|
|
return epsilon
|
|
# TODO: check if this is correct
|
|
|
|
|
|
|
|
def half_mass_radius(particles: np.ndarray):
|
|
"""
|
|
Computes the half mass radius of a set of particles.
|
|
Assumes that the particles array has the following columns: x, y, z ...
|
|
"""
|
|
if particles.shape[1] < 3:
|
|
raise ValueError("Particles array must have at least 3 columns: x, y, z")
|
|
|
|
# even though in the simple example, all the masses are the same, we will consider the general case
|
|
total_mass = np.sum(particles[:, 3])
|
|
half_mass = total_mass / 2
|
|
|
|
# sort the particles by distance
|
|
r = np.linalg.norm(particles[:, :3], axis=1)
|
|
indices = np.argsort(r)
|
|
r = r[indices]
|
|
masses = particles[indices, 3]
|
|
masses_cumsum = np.cumsum(masses)
|
|
|
|
i = np.argmin(np.abs(masses_cumsum - half_mass))
|
|
logger.debug(f"Half mass radius: {r[i]} for {i}th particle of {particles.shape[0]}")
|
|
r_hm = r[i]
|
|
|
|
return r_hm
|
|
|
|
|
|
|
|
def total_energy(particles: np.ndarray):
|
|
"""
|
|
Computes the total energy of a set of particles.
|
|
Assumes that the particles array has the following columns: x, y, z, vx, vy, vz, m.
|
|
Uses the approximation that the particles are in a central potential as computed in analytical.py
|
|
"""
|
|
if particles.shape[1] != 7:
|
|
raise ValueError("Particles array must have 7 columns: x, y, z, vx, vy, vz, m")
|
|
|
|
# compute the kinetic energy
|
|
v = particles[:, 3:6]
|
|
m = particles[:, 6]
|
|
ke = 0.5 * np.sum(m * np.linalg.norm(v, axis=1)**2)
|
|
|
|
# # compute the potential energy
|
|
# forces = forces_basic.analytical_forces(particles)
|
|
# r = np.linalg.norm(particles[:, :3], axis=1)
|
|
# pe_particles = -forces[:, 0] * particles[:, 0] - forces[:, 1] * particles[:, 1] - forces[:, 2] * particles[:, 2]
|
|
# pe = np.sum(pe_particles)
|
|
# # 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()
|