use crate::aabb::Aabb3d;
use crate::kernel::DiscreteSquaredDistanceCubicKernel;
use crate::mesh::{HexMesh3d, MeshAttribute, MeshWithData};
use crate::neighborhood_search::NeighborhoodList;
use crate::uniform_grid::{OwningSubdomainGrid, Subdomain, UniformGrid};
use crate::utils::{ChunkSize, ParallelPolicy};
use crate::{new_map, profile, HashState, Index, MapType, ParallelMapType, Real};
use dashmap::ReadOnlyView as ReadDashMap;
use log::{info, trace, warn};
use nalgebra::Vector3;
use rayon::prelude::*;
use std::cell::RefCell;
use thiserror::Error as ThisError;
use thread_local::ThreadLocal;
#[derive(Debug, ThisError)]
pub enum DensityMapError<R: Real> {
#[error("the adapted subdomain for the density map is inconsistent/degenerate")]
InvalidDomain {
margin: R,
domain: Aabb3d<R>,
},
}
#[inline(never)]
pub fn compute_particle_densities<I: Index, R: Real>(
particle_positions: &[Vector3<R>],
particle_neighbor_lists: &[Vec<usize>],
compact_support_radius: R,
particle_rest_mass: R,
enable_multi_threading: bool,
) -> Vec<R> {
let mut densities = Vec::new();
if enable_multi_threading {
parallel_compute_particle_densities::<I, R>(
particle_positions,
particle_neighbor_lists,
compact_support_radius,
particle_rest_mass,
&mut densities,
)
} else {
sequential_compute_particle_densities::<I, R, _>(
particle_positions,
particle_neighbor_lists,
compact_support_radius,
particle_rest_mass,
&mut densities,
)
}
densities
}
#[inline(never)]
pub fn compute_particle_densities_inplace<I: Index, R: Real>(
particle_positions: &[Vector3<R>],
particle_neighbor_lists: &[Vec<usize>],
compact_support_radius: R,
particle_rest_mass: R,
enable_multi_threading: bool,
densities: &mut Vec<R>,
) {
if enable_multi_threading {
parallel_compute_particle_densities::<I, R>(
particle_positions,
particle_neighbor_lists,
compact_support_radius,
particle_rest_mass,
densities,
)
} else {
sequential_compute_particle_densities::<I, R, _>(
particle_positions,
particle_neighbor_lists,
compact_support_radius,
particle_rest_mass,
densities,
)
}
}
fn init_density_storage<R: Real>(densities: &mut Vec<R>, new_len: usize) {
densities.resize(new_len, R::zero());
}
#[inline(never)]
pub fn sequential_compute_particle_densities<I: Index, R: Real, Nl: NeighborhoodList + ?Sized>(
particle_positions: &[Vector3<R>],
particle_neighbor_lists: &Nl,
compact_support_radius: R,
particle_rest_mass: R,
particle_densities: &mut Vec<R>,
) {
profile!("sequential_compute_particle_densities");
sequential_compute_particle_densities_filtered::<I, R, Nl>(
particle_positions,
particle_neighbor_lists,
compact_support_radius,
particle_rest_mass,
particle_densities,
|_| true,
)
}
#[inline(never)]
pub fn sequential_compute_particle_densities_filtered<
I: Index,
R: Real,
Nl: NeighborhoodList + ?Sized,
>(
particle_positions: &[Vector3<R>],
particle_neighbor_lists: &Nl,
compact_support_radius: R,
particle_rest_mass: R,
particle_densities: &mut Vec<R>,
filter: impl Fn(usize) -> bool,
) {
profile!("sequential_compute_particle_densities_filtered");
init_density_storage(particle_densities, particle_positions.len());
let kernel = DiscreteSquaredDistanceCubicKernel::new::<f64>(1000, compact_support_radius);
for (i, particle_i_position) in particle_positions
.iter()
.enumerate()
.filter(|(i, _)| filter(*i))
{
let mut particle_i_density = kernel.evaluate(R::zero());
for particle_j_position in particle_neighbor_lists
.neighbors(i)
.iter()
.map(|&j| &particle_positions[j])
{
let r_squared = (particle_j_position - particle_i_position).norm_squared();
particle_i_density += kernel.evaluate(r_squared);
}
particle_i_density *= particle_rest_mass;
particle_densities[i] = particle_i_density;
}
}
#[inline(never)]
pub fn parallel_compute_particle_densities<I: Index, R: Real>(
particle_positions: &[Vector3<R>],
particle_neighbor_lists: &[Vec<usize>],
compact_support_radius: R,
particle_rest_mass: R,
particle_densities: &mut Vec<R>,
) {
profile!("parallel_compute_particle_densities");
init_density_storage(particle_densities, particle_positions.len());
let kernel = DiscreteSquaredDistanceCubicKernel::new::<f64>(1000, compact_support_radius);
particle_positions
.par_iter()
.with_min_len(8)
.zip_eq(particle_neighbor_lists.par_iter())
.zip_eq(particle_densities.par_iter_mut())
.for_each(
|((particle_i_position, particle_i_neighbors), particle_i_density)| {
let mut density = kernel.evaluate(R::zero());
for particle_j_position in
particle_i_neighbors.iter().map(|&j| &particle_positions[j])
{
let r_squared = (particle_j_position - particle_i_position).norm_squared();
density += kernel.evaluate(r_squared);
}
density *= particle_rest_mass;
*particle_i_density = density;
},
);
}
#[derive(Clone, Debug)]
pub enum DensityMap<I: Index, R: Real> {
Standard(MapType<I, R>),
DashMap(ReadDashMap<I, R, HashState>),
}
impl<I: Index, R: Real> From<MapType<I, R>> for DensityMap<I, R> {
fn from(map: MapType<I, R>) -> Self {
Self::Standard(map)
}
}
impl<I: Index, R: Real> From<ParallelMapType<I, R>> for DensityMap<I, R> {
fn from(map: ParallelMapType<I, R>) -> Self {
Self::DashMap(map.into_read_only())
}
}
impl<I: Index, R: Real> DensityMap<I, R> {
pub fn to_vec(&self) -> Vec<(I, R)> {
match self {
DensityMap::Standard(map) => map.iter().map(|(&i, &r)| (i, r)).collect(),
DensityMap::DashMap(map) => map.iter().map(|(&i, &r)| (i, r)).collect(),
}
}
pub fn len(&self) -> usize {
match self {
DensityMap::Standard(map) => map.len(),
DensityMap::DashMap(map) => map.len(),
}
}
pub fn get(&self, flat_point_index: I) -> Option<R> {
match self {
DensityMap::Standard(map) => map.get(&flat_point_index).copied(),
DensityMap::DashMap(map) => map.get(&flat_point_index).copied(),
}
}
fn standard_or_insert_mut(&mut self) -> &mut MapType<I, R> {
match self {
DensityMap::Standard(map) => return map,
_ => {}
}
*self = new_map().into();
self.standard_or_insert_mut()
}
pub fn for_each<F: FnMut(I, R)>(&self, f: F) {
let mut f = f;
match self {
DensityMap::Standard(map) => map.iter().for_each(|(&i, &r)| f(i, r)),
DensityMap::DashMap(map) => map.iter().for_each(|(&i, &r)| f(i, r)),
}
}
}
#[inline(never)]
pub fn generate_sparse_density_map<I: Index, R: Real>(
grid: &UniformGrid<I, R>,
subdomain: Option<&OwningSubdomainGrid<I, R>>,
particle_positions: &[Vector3<R>],
particle_densities: &[R],
active_particles: Option<&[usize]>,
particle_rest_mass: R,
compact_support_radius: R,
cube_size: R,
allow_threading: bool,
density_map: &mut DensityMap<I, R>,
) -> Result<(), DensityMapError<R>> {
trace!(
"Starting construction of sparse density map... (Input: {} particles)",
if let Some(active_particles) = active_particles {
active_particles.len()
} else {
particle_positions.len()
}
);
if let Some(subdomain) = subdomain {
if allow_threading {
panic!("Multi threading not implemented for density map with subdomain");
} else {
sequential_generate_sparse_density_map_subdomain(
subdomain,
particle_positions,
particle_densities,
active_particles,
particle_rest_mass,
compact_support_radius,
cube_size,
density_map,
)?;
}
} else {
if allow_threading {
*density_map = parallel_generate_sparse_density_map(
grid,
particle_positions,
particle_densities,
active_particles,
particle_rest_mass,
compact_support_radius,
cube_size,
)?
} else {
*density_map = sequential_generate_sparse_density_map(
grid,
particle_positions,
particle_densities,
active_particles,
particle_rest_mass,
compact_support_radius,
cube_size,
)?
}
};
trace!(
"Sparse density map was constructed. (Output: density map with {} grid point data entries)",
density_map.len()
);
Ok(())
}
#[inline(never)]
pub fn sequential_generate_sparse_density_map<I: Index, R: Real>(
grid: &UniformGrid<I, R>,
particle_positions: &[Vector3<R>],
particle_densities: &[R],
active_particles: Option<&[usize]>,
particle_rest_mass: R,
compact_support_radius: R,
cube_size: R,
) -> Result<DensityMap<I, R>, DensityMapError<R>> {
profile!("sequential_generate_sparse_density_map");
let mut sparse_densities = new_map();
let density_map_generator = SparseDensityMapGenerator::try_new(
grid,
compact_support_radius,
cube_size,
particle_rest_mass,
)?;
let process_particle = |particle_data: (&Vector3<R>, R)| {
let (particle, particle_density) = particle_data;
density_map_generator.compute_particle_density_contribution(
grid,
&mut sparse_densities,
particle,
particle_density,
);
};
match active_particles {
None => particle_positions
.iter()
.zip(particle_densities.iter().copied())
.for_each(process_particle),
Some(indices) => indices
.iter()
.map(|&i| &particle_positions[i])
.zip(indices.iter().map(|&i| particle_densities[i]))
.for_each(process_particle),
}
Ok(sparse_densities.into())
}
#[inline(never)]
pub fn sequential_generate_sparse_density_map_subdomain<I: Index, R: Real>(
subdomain: &OwningSubdomainGrid<I, R>,
particle_positions: &[Vector3<R>],
particle_densities: &[R],
active_particles: Option<&[usize]>,
particle_rest_mass: R,
compact_support_radius: R,
cube_size: R,
density_map: &mut DensityMap<I, R>,
) -> Result<(), DensityMapError<R>> {
profile!("sequential_generate_sparse_density_map_subdomain");
let mut sparse_densities = density_map.standard_or_insert_mut();
sparse_densities.clear();
let density_map_generator = SparseDensityMapGenerator::try_new(
&subdomain.global_grid(),
compact_support_radius,
cube_size,
particle_rest_mass,
)?;
let process_particle = |particle_data: (&Vector3<R>, R)| {
let (particle, particle_density) = particle_data;
density_map_generator.compute_particle_density_contribution_subdomain(
subdomain,
&mut sparse_densities,
particle,
particle_density,
);
};
match active_particles {
None => particle_positions
.iter()
.zip(particle_densities.iter().copied())
.for_each(process_particle),
Some(indices) => indices
.iter()
.map(|&i| &particle_positions[i])
.zip(indices.iter().map(|&i| particle_densities[i]))
.for_each(process_particle),
}
Ok(())
}
#[inline(never)]
pub fn parallel_generate_sparse_density_map<I: Index, R: Real>(
grid: &UniformGrid<I, R>,
particle_positions: &[Vector3<R>],
particle_densities: &[R],
active_particles: Option<&[usize]>,
particle_rest_mass: R,
compact_support_radius: R,
cube_size: R,
) -> Result<DensityMap<I, R>, DensityMapError<R>> {
profile!("parallel_generate_sparse_density_map");
let sparse_densities: ThreadLocal<RefCell<MapType<I, R>>> = ThreadLocal::new();
{
let density_map_generator = SparseDensityMapGenerator::try_new(
grid,
compact_support_radius,
cube_size,
particle_rest_mass,
)?;
profile!("generate thread local maps");
match active_particles {
None => {
let chunk_size =
ChunkSize::new(&ParallelPolicy::default(), particle_positions.len())
.with_log("particles", "density map generation")
.chunk_size;
particle_positions
.par_chunks(chunk_size)
.zip(particle_densities.par_chunks(chunk_size))
.for_each(|(position_chunk, density_chunk)| {
let map = sparse_densities
.get_or(|| RefCell::new(MapType::with_hasher(HashState::default())));
let mut mut_map = map.borrow_mut();
let process_particle_map = |particle_data: (&Vector3<R>, R)| {
let (particle, particle_density) = particle_data;
density_map_generator.compute_particle_density_contribution(
grid,
&mut mut_map,
particle,
particle_density,
);
};
assert_eq!(position_chunk.len(), density_chunk.len());
position_chunk
.iter()
.zip(density_chunk.iter().copied())
.for_each(process_particle_map);
})
}
Some(indices) => {
let chunk_size = ChunkSize::new(&ParallelPolicy::default(), indices.len())
.with_log("active particles", "density map generation")
.chunk_size;
indices.par_chunks(chunk_size).for_each(|index_chunk| {
let map = sparse_densities
.get_or(|| RefCell::new(MapType::with_hasher(HashState::default())));
let mut mut_map = map.borrow_mut();
let process_particle_map = |particle_data: (&Vector3<R>, R)| {
let (particle, particle_density) = particle_data;
density_map_generator.compute_particle_density_contribution(
grid,
&mut mut_map,
particle,
particle_density,
);
};
index_chunk
.iter()
.map(|&i| (&particle_positions[i], particle_densities[i]))
.for_each(process_particle_map);
});
}
}
}
{
profile!("merge thread local maps to global map");
let mut local_density_maps = sparse_densities
.into_iter()
.map(|m| m.into_inner())
.collect::<Vec<_>>();
info!(
"Merging {} thread local density maps to a single global map...",
local_density_maps.len()
);
let global_density_map = ParallelMapType::with_hasher(HashState::default());
local_density_maps.par_iter_mut().for_each(|local_map| {
for (idx, density) in local_map.drain() {
*global_density_map.entry(idx).or_insert(R::zero()) += density;
}
});
Ok(global_density_map.into())
}
}
struct SparseDensityMapGenerator<I: Index, R: Real> {
particle_rest_mass: R,
half_supported_cells: I,
supported_points: I,
kernel_evaluation_radius_sq: R,
kernel: DiscreteSquaredDistanceCubicKernel<R>,
allowed_domain: Aabb3d<R>,
}
pub(crate) struct GridKernelExtents<I: Index, R: Real> {
pub half_supported_cells: I,
pub supported_points: I,
pub kernel_evaluation_radius: R,
}
pub(crate) fn compute_kernel_evaluation_radius<I: Index, R: Real>(
compact_support_radius: R,
cube_size: R,
) -> GridKernelExtents<I, R> {
let half_supported_cells_real = (compact_support_radius / cube_size).ceil();
let half_supported_cells: I = half_supported_cells_real.to_index_unchecked();
let supported_cells: I = half_supported_cells.times(2) + I::one();
let supported_points: I = I::one() + supported_cells;
let kernel_evaluation_radius =
cube_size * half_supported_cells_real * (R::one() + R::default_epsilon().sqrt());
GridKernelExtents {
half_supported_cells,
supported_points,
kernel_evaluation_radius,
}
}
impl<I: Index, R: Real> SparseDensityMapGenerator<I, R> {
fn try_new(
grid: &UniformGrid<I, R>,
compact_support_radius: R,
cube_size: R,
particle_rest_mass: R,
) -> Result<Self, DensityMapError<R>> {
let GridKernelExtents {
half_supported_cells,
supported_points,
kernel_evaluation_radius,
} = compute_kernel_evaluation_radius(compact_support_radius, cube_size);
let kernel_evaluation_radius_sq = kernel_evaluation_radius * kernel_evaluation_radius;
let kernel = DiscreteSquaredDistanceCubicKernel::new::<f64>(1000, compact_support_radius);
let allowed_domain = {
let mut aabb = grid.aabb().clone();
aabb.grow_uniformly(kernel_evaluation_radius.neg());
aabb
};
if allowed_domain.is_degenerate() || !allowed_domain.is_consistent() {
warn!(
"The allowed domain of particles for a subdomain is inconsistent/degenerate: {:?}",
allowed_domain
);
warn!("No particles can be found in this domain. Increase the domain of the surface reconstruction to avoid this.");
Err(DensityMapError::InvalidDomain {
margin: kernel_evaluation_radius,
domain: allowed_domain,
})
} else {
Ok(Self {
half_supported_cells,
supported_points,
kernel_evaluation_radius_sq,
kernel,
allowed_domain,
particle_rest_mass,
})
}
}
fn compute_particle_density_contribution(
&self,
grid: &UniformGrid<I, R>,
sparse_densities: &mut MapType<I, R>,
particle: &Vector3<R>,
particle_density: R,
) {
if !self.allowed_domain.contains_point(particle) {
return;
}
let min_supported_point_ijk = {
let cell_ijk = grid.enclosing_cell(particle);
[
cell_ijk[0] - self.half_supported_cells,
cell_ijk[1] - self.half_supported_cells,
cell_ijk[2] - self.half_supported_cells,
]
};
let max_supported_point_ijk = [
min_supported_point_ijk[0] + self.supported_points,
min_supported_point_ijk[1] + self.supported_points,
min_supported_point_ijk[2] + self.supported_points,
];
self.particle_support_loop(
sparse_densities,
grid,
&min_supported_point_ijk,
&max_supported_point_ijk,
particle,
particle_density,
);
}
fn compute_particle_density_contribution_subdomain(
&self,
subdomain: &OwningSubdomainGrid<I, R>,
sparse_densities: &mut MapType<I, R>,
particle: &Vector3<R>,
particle_density: R,
) {
let grid = subdomain.global_grid();
let subdomain_grid = subdomain.subdomain_grid();
let subdomain_offset = subdomain.subdomain_offset();
if !self.allowed_domain.contains_point(particle) {
return;
}
let global_subdomain_min_point = subdomain_offset;
let global_subdomain_max_point = [
global_subdomain_min_point[0] + subdomain_grid.points_per_dim()[0],
global_subdomain_min_point[1] + subdomain_grid.points_per_dim()[1],
global_subdomain_min_point[2] + subdomain_grid.points_per_dim()[2],
];
let min_supported_point_ijk = {
let cell_ijk = grid.enclosing_cell(particle);
[
(cell_ijk[0] - self.half_supported_cells).max(global_subdomain_min_point[0]),
(cell_ijk[1] - self.half_supported_cells).max(global_subdomain_min_point[1]),
(cell_ijk[2] - self.half_supported_cells).max(global_subdomain_min_point[2]),
]
};
let max_supported_point_ijk = {
[
(min_supported_point_ijk[0] + self.supported_points)
.min(global_subdomain_max_point[0]),
(min_supported_point_ijk[1] + self.supported_points)
.min(global_subdomain_max_point[1]),
(min_supported_point_ijk[2] + self.supported_points)
.min(global_subdomain_max_point[2]),
]
};
if min_supported_point_ijk
.iter()
.copied()
.zip(global_subdomain_max_point.iter().copied())
.any(|(min_support, subdomain_max)| min_support > subdomain_max)
{
return;
}
if max_supported_point_ijk
.iter()
.copied()
.zip(global_subdomain_min_point.iter().copied())
.any(|(max_support, subdomain_min)| max_support < subdomain_min)
{
return;
}
self.particle_support_loop(
sparse_densities,
grid,
&min_supported_point_ijk,
&max_supported_point_ijk,
particle,
particle_density,
);
}
#[inline(always)]
fn particle_support_loop(
&self,
sparse_densities: &mut MapType<I, R>,
grid: &UniformGrid<I, R>,
min_supported_point_ijk: &[I; 3],
max_supported_point_ijk: &[I; 3],
particle: &Vector3<R>,
particle_density: R,
) {
let particle_volume = self.particle_rest_mass / particle_density;
let min_supported_point = grid.point_coordinates_array(&min_supported_point_ijk);
let mut dx = min_supported_point[0] - particle[0]
- grid.cell_size();
let mut i = min_supported_point_ijk[0];
while i != max_supported_point_ijk[0] {
dx += grid.cell_size();
let dxdx = dx * dx;
let mut dy = min_supported_point[1] - particle[1] - grid.cell_size();
let mut j = min_supported_point_ijk[1];
while j != max_supported_point_ijk[1] {
dy += grid.cell_size();
let dydy = dy * dy;
let mut dz = min_supported_point[2] - particle[2] - grid.cell_size();
let mut k = min_supported_point_ijk[2];
while k != max_supported_point_ijk[2] {
dz += grid.cell_size();
let dzdz = dz * dz;
let r_squared = dxdx + dydy + dzdz;
if r_squared < self.kernel_evaluation_radius_sq {
let density_contribution =
particle_volume * self.kernel.evaluate(r_squared);
let flat_point_index = grid.flatten_point_indices(i, j, k);
*sparse_densities
.entry(flat_point_index)
.or_insert(R::zero()) += density_contribution;
}
k = k + I::one();
}
j = j + I::one();
}
i = i + I::one();
}
}
}
#[inline(never)]
pub fn sparse_density_map_to_hex_mesh<I: Index, R: Real>(
density_map: &DensityMap<I, R>,
grid: &UniformGrid<I, R>,
default_value: R,
) -> MeshWithData<R, HexMesh3d<R>> {
profile!("sparse_density_map_to_hex_mesh");
let mut mesh = HexMesh3d {
vertices: Vec::new(),
cells: Vec::new(),
};
let mut values = Vec::new();
let mut cells = new_map();
density_map.for_each(|flat_point_index, point_value| {
let point = grid.try_unflatten_point_index(flat_point_index).unwrap();
let point_coords = grid.point_coordinates(&point);
let vertex_index = mesh.vertices.len();
mesh.vertices.push(point_coords);
values.push(point_value);
let neighborhood = grid.get_point_neighborhood(&point);
for cell in grid.cells_adjacent_to_point(&neighborhood).iter().flatten() {
let flat_cell_index = grid.flatten_cell_index(cell);
let cell_connectivity_entry = cells
.entry(flat_cell_index)
.or_insert_with(|| [None, None, None, None, None, None, None, None]);
let local_point_index = cell.local_point_index_of(point.index()).unwrap();
cell_connectivity_entry[local_point_index] = Some(vertex_index);
}
});
let mut additional_vertices = new_map();
for (flat_cell_index, cell_vertices) in cells.iter_mut() {
let cell = grid.try_unflatten_cell_index(*flat_cell_index).unwrap();
for (local_point_index, vertex) in cell_vertices.iter_mut().enumerate() {
if vertex.is_none() {
let point = cell.global_point_index_of(local_point_index).unwrap();
let flat_point_index = grid.flatten_point_index(&point);
let vertex_entry =
additional_vertices
.entry(flat_point_index)
.or_insert_with(|| {
let point_coords = grid.point_coordinates(&point);
let vertex_index = mesh.vertices.len();
mesh.vertices.push(point_coords);
values.push(default_value);
vertex_index
});
*vertex = Some(*vertex_entry);
}
}
}
mesh.cells.reserve(cells.len());
for (_, cell_vertices) in cells.iter() {
mesh.cells.push([
cell_vertices[0].unwrap(),
cell_vertices[1].unwrap(),
cell_vertices[2].unwrap(),
cell_vertices[3].unwrap(),
cell_vertices[4].unwrap(),
cell_vertices[5].unwrap(),
cell_vertices[6].unwrap(),
cell_vertices[7].unwrap(),
]);
}
MeshWithData::new(mesh).with_point_data(MeshAttribute::new_real_scalar(
"density".to_string(),
values,
))
}