1use crate::aabb::Aabb3d;
27use crate::kernel::{CubicSplineKernel, SymmetricKernel3d};
28use crate::mesh::{HexMesh3d, MeshWithData, OwnedMeshAttribute};
29use crate::neighborhood_search::NeighborhoodList;
30use crate::uniform_grid::UniformGrid;
31use crate::utils::{ChunkSize, ParallelPolicy};
32use crate::{HashState, Index, MapType, ParallelMapType, Real, new_map, profile};
33use dashmap::ReadOnlyView as ReadDashMap;
34use log::{info, trace, warn};
35use nalgebra::{Scalar, Vector3};
36use rayon::prelude::*;
37use std::borrow::Cow;
38use std::cell::RefCell;
39use std::hash::Hash;
40use thiserror::Error as ThisError;
41use thread_local::ThreadLocal;
42
43#[derive(Debug, ThisError)]
48pub enum DensityMapError<R: Scalar> {
49 #[error("the adapted subdomain for the density map is inconsistent/degenerate")]
55 InvalidDomain {
56 margin: R,
58 domain: Aabb3d<R>,
60 },
61}
62
63#[inline(never)]
65pub fn compute_particle_densities<I: Index, R: Real>(
66 particle_positions: &[Vector3<R>],
67 particle_neighbor_lists: &[Vec<usize>],
68 compact_support_radius: R,
69 particle_rest_mass: R,
70 enable_multi_threading: bool,
71) -> Vec<R> {
72 let mut densities = Vec::new();
73 if enable_multi_threading {
74 parallel_compute_particle_densities::<I, R>(
75 particle_positions,
76 particle_neighbor_lists,
77 compact_support_radius,
78 particle_rest_mass,
79 &mut densities,
80 )
81 } else {
82 sequential_compute_particle_densities::<I, R, _>(
83 particle_positions,
84 particle_neighbor_lists,
85 compact_support_radius,
86 particle_rest_mass,
87 &mut densities,
88 )
89 }
90 densities
91}
92
93#[inline(never)]
95pub fn compute_particle_densities_inplace<I: Index, R: Real>(
96 particle_positions: &[Vector3<R>],
97 particle_neighbor_lists: &[Vec<usize>],
98 compact_support_radius: R,
99 particle_rest_mass: R,
100 enable_multi_threading: bool,
101 densities: &mut Vec<R>,
102) {
103 if enable_multi_threading {
104 parallel_compute_particle_densities::<I, R>(
105 particle_positions,
106 particle_neighbor_lists,
107 compact_support_radius,
108 particle_rest_mass,
109 densities,
110 )
111 } else {
112 sequential_compute_particle_densities::<I, R, _>(
113 particle_positions,
114 particle_neighbor_lists,
115 compact_support_radius,
116 particle_rest_mass,
117 densities,
118 )
119 }
120}
121
122fn init_density_storage<R: Real>(densities: &mut Vec<R>, new_len: usize) {
123 densities.resize(new_len, R::zero());
125 }
127
128#[inline(never)]
130pub fn sequential_compute_particle_densities<I: Index, R: Real, Nl: NeighborhoodList + ?Sized>(
131 particle_positions: &[Vector3<R>],
132 particle_neighbor_lists: &Nl,
133 compact_support_radius: R,
134 particle_rest_mass: R,
135 particle_densities: &mut Vec<R>,
136) {
137 profile!("sequential_compute_particle_densities");
138
139 sequential_compute_particle_densities_filtered::<I, R, Nl>(
140 particle_positions,
141 particle_neighbor_lists,
142 compact_support_radius,
143 particle_rest_mass,
144 particle_densities,
145 |_| true,
146 )
147}
148
149#[inline(never)]
150pub fn sequential_compute_particle_densities_filtered<
151 I: Index,
152 R: Real,
153 Nl: NeighborhoodList + ?Sized,
154>(
155 particle_positions: &[Vector3<R>],
156 particle_neighbor_lists: &Nl,
157 compact_support_radius: R,
158 particle_rest_mass: R,
159 particle_densities: &mut Vec<R>,
160 filter: impl Fn(usize) -> bool,
161) {
162 profile!("sequential_compute_particle_densities_filtered");
163
164 init_density_storage(particle_densities, particle_positions.len());
165
166 let kernel = CubicSplineKernel::new(compact_support_radius);
168
169 for (i, particle_i_position) in particle_positions
170 .iter()
171 .enumerate()
172 .filter(|(i, _)| filter(*i))
173 {
174 let mut particle_i_density = kernel.evaluate(R::zero());
175 for particle_j_position in particle_neighbor_lists
176 .neighbors(i)
177 .iter()
178 .map(|&j| &particle_positions[j])
179 {
180 let r = (particle_j_position - particle_i_position).norm();
181 particle_i_density += kernel.evaluate(r);
182 }
183 particle_i_density *= particle_rest_mass;
184 particle_densities[i] = particle_i_density;
185 }
186}
187
188#[inline(never)]
190pub fn parallel_compute_particle_densities<I: Index, R: Real>(
191 particle_positions: &[Vector3<R>],
192 particle_neighbor_lists: &[Vec<usize>],
193 compact_support_radius: R,
194 particle_rest_mass: R,
195 particle_densities: &mut Vec<R>,
196) {
197 profile!("parallel_compute_particle_densities");
198
199 init_density_storage(particle_densities, particle_positions.len());
200
201 let kernel = CubicSplineKernel::new(compact_support_radius);
203
204 particle_positions
205 .par_iter()
206 .with_min_len(8)
207 .zip_eq(particle_neighbor_lists.par_iter())
208 .zip_eq(particle_densities.par_iter_mut())
209 .for_each(
210 |((particle_i_position, particle_i_neighbors), particle_i_density)| {
211 let mut density = kernel.evaluate(R::zero());
212 for particle_j_position in
213 particle_i_neighbors.iter().map(|&j| &particle_positions[j])
214 {
215 let r = (particle_j_position - particle_i_position).norm();
216 density += kernel.evaluate(r);
217 }
218 density *= particle_rest_mass;
219 *particle_i_density = density;
220 },
221 );
222}
223
224#[derive(Clone, Debug)]
229pub enum DensityMap<'a, I: Scalar + Eq + Hash, R: Scalar> {
230 Standard(MapType<I, R>),
231 DashMap(ReadDashMap<I, R, HashState>),
232 Dense(Cow<'a, [R]>),
233}
234
235pub type OwnedDensityMap<I, R> = DensityMap<'static, I, R>;
237
238impl<I: Index, R: Real> Default for OwnedDensityMap<I, R> {
239 fn default() -> Self {
240 DensityMap::Standard(MapType::default())
241 }
242}
243
244impl<I: Index, R: Real> From<MapType<I, R>> for OwnedDensityMap<I, R> {
245 fn from(map: MapType<I, R>) -> Self {
246 Self::Standard(map)
247 }
248}
249
250impl<I: Index, R: Real> From<ParallelMapType<I, R>> for OwnedDensityMap<I, R> {
251 fn from(map: ParallelMapType<I, R>) -> Self {
252 Self::DashMap(map.into_read_only())
253 }
254}
255
256impl<I: Index, R: Real> From<Vec<R>> for DensityMap<'static, I, R> {
257 fn from(values: Vec<R>) -> Self {
258 Self::Dense(values.into())
259 }
260}
261
262impl<'a, I: Index, R: Real> From<&'a [R]> for DensityMap<'a, I, R> {
263 fn from(values: &'a [R]) -> Self {
264 Self::Dense(values.into())
265 }
266}
267
268impl<'a, I: Index, R: Real> DensityMap<'a, I, R> {
269 pub fn to_vec(&self) -> Vec<(I, R)> {
271 match self {
272 DensityMap::Standard(map) => map.iter().map(|(&i, &r)| (i, r)).collect(),
273 DensityMap::DashMap(map) => map.iter().map(|(&i, &r)| (i, r)).collect(),
274 DensityMap::Dense(values) => values
275 .iter()
276 .copied()
277 .enumerate()
278 .map(|(i, r)| (I::from_usize(i).unwrap(), r))
279 .collect(),
280 }
281 }
282
283 pub fn len(&self) -> usize {
285 match self {
286 DensityMap::Standard(map) => map.len(),
287 DensityMap::DashMap(map) => map.len(),
288 DensityMap::Dense(values) => values.len(),
289 }
290 }
291
292 pub fn get(&self, flat_point_index: I) -> Option<R> {
294 match self {
295 DensityMap::Standard(map) => map.get(&flat_point_index).copied(),
296 DensityMap::DashMap(map) => map.get(&flat_point_index).copied(),
297 DensityMap::Dense(values) => values.get(flat_point_index.to_usize()?).copied(),
298 }
299 }
300
301 pub fn for_each<F: FnMut(I, R)>(&self, f: F) {
303 let mut f = f;
304 match self {
305 DensityMap::Standard(map) => map.iter().for_each(|(&i, &r)| f(i, r)),
306 DensityMap::DashMap(map) => map.iter().for_each(|(&i, &r)| f(i, r)),
307 DensityMap::Dense(values) => values.iter().copied().enumerate().for_each(|(i, r)| {
308 f(I::from_usize(i).unwrap(), r);
309 }),
310 }
311 }
312}
313
314#[inline(never)]
316pub fn generate_sparse_density_map<I: Index, R: Real>(
317 grid: &UniformGrid<I, R>,
318 particle_positions: &[Vector3<R>],
319 particle_densities: &[R],
320 active_particles: Option<&[usize]>,
321 particle_rest_mass: R,
322 compact_support_radius: R,
323 cube_size: R,
324 allow_threading: bool,
325 density_map: &mut DensityMap<I, R>,
326) -> Result<(), DensityMapError<R>> {
327 trace!(
328 "Starting construction of sparse density map... (Input: {} particles)",
329 if let Some(active_particles) = active_particles {
330 active_particles.len()
331 } else {
332 particle_positions.len()
333 }
334 );
335
336 if allow_threading {
337 *density_map = parallel_generate_sparse_density_map(
338 grid,
339 particle_positions,
340 particle_densities,
341 active_particles,
342 particle_rest_mass,
343 compact_support_radius,
344 cube_size,
345 )?
346 } else {
347 *density_map = sequential_generate_sparse_density_map(
348 grid,
349 particle_positions,
350 particle_densities,
351 active_particles,
352 particle_rest_mass,
353 compact_support_radius,
354 cube_size,
355 )?
356 }
357
358 trace!(
359 "Sparse density map was constructed. (Output: density map with {} grid point data entries)",
360 density_map.len()
361 );
362
363 Ok(())
364}
365
366#[inline(never)]
368pub fn sequential_generate_sparse_density_map<I: Index, R: Real>(
369 grid: &UniformGrid<I, R>,
370 particle_positions: &[Vector3<R>],
371 particle_densities: &[R],
372 active_particles: Option<&[usize]>,
373 particle_rest_mass: R,
374 compact_support_radius: R,
375 cube_size: R,
376) -> Result<OwnedDensityMap<I, R>, DensityMapError<R>> {
377 profile!("sequential_generate_sparse_density_map");
378
379 let mut sparse_densities = new_map();
380
381 let density_map_generator = SparseDensityMapGenerator::try_new(
382 grid,
383 compact_support_radius,
384 cube_size,
385 particle_rest_mass,
386 )?;
387
388 let process_particle = |particle_data: (&Vector3<R>, R)| {
389 let (particle, particle_density) = particle_data;
390 density_map_generator.compute_particle_density_contribution(
391 grid,
392 &mut sparse_densities,
393 particle,
394 particle_density,
395 );
396 };
397
398 match active_particles {
399 None => particle_positions
400 .iter()
401 .zip(particle_densities.iter().copied())
402 .for_each(process_particle),
403 Some(indices) => indices
404 .iter()
405 .map(|&i| &particle_positions[i])
406 .zip(indices.iter().map(|&i| particle_densities[i]))
407 .for_each(process_particle),
408 }
409
410 Ok(sparse_densities.into())
411}
412
413#[inline(never)]
415pub fn parallel_generate_sparse_density_map<I: Index, R: Real>(
416 grid: &UniformGrid<I, R>,
417 particle_positions: &[Vector3<R>],
418 particle_densities: &[R],
419 active_particles: Option<&[usize]>,
420 particle_rest_mass: R,
421 compact_support_radius: R,
422 cube_size: R,
423) -> Result<OwnedDensityMap<I, R>, DensityMapError<R>> {
424 profile!("parallel_generate_sparse_density_map");
425
426 let sparse_densities: ThreadLocal<RefCell<MapType<I, R>>> = ThreadLocal::new();
428
429 {
431 let density_map_generator = SparseDensityMapGenerator::try_new(
432 grid,
433 compact_support_radius,
434 cube_size,
435 particle_rest_mass,
436 )?;
437
438 profile!("generate thread local maps");
439
440 match active_particles {
441 None => {
443 let chunk_size =
444 ChunkSize::new(&ParallelPolicy::default(), particle_positions.len())
445 .with_log("particles", "density map generation")
446 .chunk_size;
447
448 particle_positions
449 .par_chunks(chunk_size)
450 .zip(particle_densities.par_chunks(chunk_size))
451 .for_each(|(position_chunk, density_chunk)| {
452 let map = sparse_densities
454 .get_or(|| RefCell::new(MapType::with_hasher(HashState::default())));
455 let mut mut_map = map.borrow_mut();
456
457 let process_particle_map = |particle_data: (&Vector3<R>, R)| {
458 let (particle, particle_density) = particle_data;
459 density_map_generator.compute_particle_density_contribution(
460 grid,
461 &mut mut_map,
462 particle,
463 particle_density,
464 );
465 };
466
467 assert_eq!(position_chunk.len(), density_chunk.len());
468 position_chunk
469 .iter()
470 .zip(density_chunk.iter().copied())
471 .for_each(process_particle_map);
472 })
473 }
474 Some(indices) => {
476 let chunk_size = ChunkSize::new(&ParallelPolicy::default(), indices.len())
477 .with_log("active particles", "density map generation")
478 .chunk_size;
479
480 indices.par_chunks(chunk_size).for_each(|index_chunk| {
481 let map = sparse_densities
483 .get_or(|| RefCell::new(MapType::with_hasher(HashState::default())));
484 let mut mut_map = map.borrow_mut();
485
486 let process_particle_map = |particle_data: (&Vector3<R>, R)| {
487 let (particle, particle_density) = particle_data;
488 density_map_generator.compute_particle_density_contribution(
489 grid,
490 &mut mut_map,
491 particle,
492 particle_density,
493 );
494 };
495
496 index_chunk
497 .iter()
498 .map(|&i| (&particle_positions[i], particle_densities[i]))
499 .for_each(process_particle_map);
500 });
501 }
502 }
503 }
504
505 {
507 profile!("merge thread local maps to global map");
508
509 let mut local_density_maps = sparse_densities
511 .into_iter()
512 .map(|m| m.into_inner())
513 .collect::<Vec<_>>();
514
515 info!(
516 "Merging {} thread local density maps to a single global map...",
517 local_density_maps.len()
518 );
519
520 let global_density_map = ParallelMapType::with_hasher(HashState::default());
522 local_density_maps.par_iter_mut().for_each(|local_map| {
523 for (idx, density) in local_map.drain() {
524 *global_density_map.entry(idx).or_insert(R::zero()) += density;
525 }
526 });
527
528 Ok(global_density_map.into())
529 }
530}
531
532struct SparseDensityMapGenerator<I: Index, R: Real> {
534 particle_rest_mass: R,
535 half_supported_cells: I,
536 supported_points: I,
537 kernel_evaluation_radius_sq: R,
538 kernel: CubicSplineKernel<R>,
539 allowed_domain: Aabb3d<R>,
540}
541
542pub(crate) struct GridKernelExtents<I: Index, R: Real> {
543 pub half_supported_cells: I,
545 pub supported_points: I,
547 pub kernel_evaluation_radius: R,
549}
550
551pub(crate) fn compute_kernel_evaluation_radius<I: Index, R: Real>(
552 compact_support_radius: R,
553 cube_size: R,
554) -> GridKernelExtents<I, R> {
555 assert!(
556 compact_support_radius >= R::zero(),
557 "compact support radius must be non-negative"
558 );
559 assert!(cube_size > R::zero(), "cube size must be positive");
560
561 let half_supported_cells_real = (compact_support_radius / cube_size).ceil();
563 let half_supported_cells: I = half_supported_cells_real.to_index_unchecked();
565
566 let supported_cells: I = half_supported_cells.times(2) + I::one();
568 let supported_points: I = I::one() + supported_cells;
570
571 let kernel_evaluation_radius =
573 cube_size * half_supported_cells_real * (R::one() + R::default_epsilon().sqrt());
574
575 GridKernelExtents {
576 half_supported_cells,
577 supported_points,
578 kernel_evaluation_radius,
579 }
580}
581
582impl<I: Index, R: Real> SparseDensityMapGenerator<I, R> {
584 fn try_new(
585 grid: &UniformGrid<I, R>,
586 compact_support_radius: R,
587 cube_size: R,
588 particle_rest_mass: R,
589 ) -> Result<Self, DensityMapError<R>> {
590 let GridKernelExtents {
591 half_supported_cells,
592 supported_points,
593 kernel_evaluation_radius,
594 } = compute_kernel_evaluation_radius(compact_support_radius, cube_size);
595
596 let kernel_evaluation_radius_sq = kernel_evaluation_radius * kernel_evaluation_radius;
598 let kernel = CubicSplineKernel::new(compact_support_radius);
599
600 let allowed_domain = {
607 let mut aabb = grid.aabb().clone();
608 aabb.grow_uniformly(kernel_evaluation_radius.neg());
609 aabb
610 };
611
612 if allowed_domain.is_degenerate() || !allowed_domain.is_consistent() {
613 warn!(
614 "The allowed domain of particles for a subdomain is inconsistent/degenerate: {:?}",
615 allowed_domain
616 );
617 warn!(
618 "No particles can be found in this domain. Increase the domain of the surface reconstruction to avoid this."
619 );
620 Err(DensityMapError::InvalidDomain {
621 margin: kernel_evaluation_radius,
622 domain: allowed_domain,
623 })
624 } else {
625 Ok(Self {
626 half_supported_cells,
627 supported_points,
628 kernel_evaluation_radius_sq,
629 kernel,
630 allowed_domain,
631 particle_rest_mass,
632 })
633 }
634 }
635
636 fn compute_particle_density_contribution(
638 &self,
639 grid: &UniformGrid<I, R>,
640 sparse_densities: &mut MapType<I, R>,
641 particle: &Vector3<R>,
642 particle_density: R,
643 ) {
644 if !self.allowed_domain.contains_point(particle) {
646 return;
647 }
648
649 let min_supported_point_ijk = {
651 let cell_ijk = grid.enclosing_cell(particle);
652 [
653 cell_ijk[0] - self.half_supported_cells,
654 cell_ijk[1] - self.half_supported_cells,
655 cell_ijk[2] - self.half_supported_cells,
656 ]
657 };
658
659 let max_supported_point_ijk = [
660 min_supported_point_ijk[0] + self.supported_points,
661 min_supported_point_ijk[1] + self.supported_points,
662 min_supported_point_ijk[2] + self.supported_points,
663 ];
664
665 self.particle_support_loop(
666 sparse_densities,
667 grid,
668 &min_supported_point_ijk,
669 &max_supported_point_ijk,
670 particle,
671 particle_density,
672 );
673 }
674
675 #[inline(always)]
677 fn particle_support_loop(
678 &self,
679 sparse_densities: &mut MapType<I, R>,
680 grid: &UniformGrid<I, R>,
681 min_supported_point_ijk: &[I; 3],
682 max_supported_point_ijk: &[I; 3],
683 particle: &Vector3<R>,
684 particle_density: R,
685 ) {
686 let particle_volume = self.particle_rest_mass / particle_density;
688
689 let min_supported_point = grid.point_coordinates_array(min_supported_point_ijk);
691
692 let mut dx = min_supported_point[0] - particle[0]
694 - grid.cell_size();
697
698 let mut i = min_supported_point_ijk[0];
704 while i != max_supported_point_ijk[0] {
705 dx += grid.cell_size();
706 let dxdx = dx * dx;
707
708 let mut dy = min_supported_point[1] - particle[1] - grid.cell_size();
709 let mut j = min_supported_point_ijk[1];
710 while j != max_supported_point_ijk[1] {
711 dy += grid.cell_size();
712 let dydy = dy * dy;
713
714 let mut dz = min_supported_point[2] - particle[2] - grid.cell_size();
715 let mut k = min_supported_point_ijk[2];
716 while k != max_supported_point_ijk[2] {
717 dz += grid.cell_size();
718 let dzdz = dz * dz;
719
720 let r_squared = dxdx + dydy + dzdz;
721 if r_squared < self.kernel_evaluation_radius_sq {
722 let density_contribution =
723 particle_volume * self.kernel.evaluate(r_squared.sqrt());
724
725 let flat_point_index = grid.flatten_point_indices(i, j, k);
726 *sparse_densities
727 .entry(flat_point_index)
728 .or_insert(R::zero()) += density_contribution;
729 }
730 k += I::one();
731 }
732 j += I::one();
733 }
734 i += I::one();
735 }
736 }
737}
738
739#[inline(never)]
741pub fn sparse_density_map_to_hex_mesh<I: Index, R: Real>(
742 density_map: &DensityMap<I, R>,
743 grid: &UniformGrid<I, R>,
744 default_value: R,
745) -> MeshWithData<R, HexMesh3d<R>> {
746 profile!("sparse_density_map_to_hex_mesh");
747
748 let mut mesh = HexMesh3d {
749 vertices: Vec::new(),
750 cells: Vec::new(),
751 };
752 let mut values = Vec::new();
753 let mut cells = new_map();
754
755 density_map.for_each(|flat_point_index, point_value| {
757 let point = grid.try_unflatten_point_index(flat_point_index).unwrap();
758 let point_coords = grid.point_coordinates(&point);
759
760 let vertex_index = mesh.vertices.len();
762 mesh.vertices.push(point_coords);
763 values.push(point_value);
764
765 let neighborhood = grid.get_point_neighborhood(&point);
767 for cell in grid.cells_adjacent_to_point(&neighborhood).iter().flatten() {
768 let flat_cell_index = grid.flatten_cell_index(cell);
769
770 let cell_connectivity_entry = cells
771 .entry(flat_cell_index)
772 .or_insert_with(|| [None, None, None, None, None, None, None, None]);
773
774 let local_point_index = cell.local_point_index_of(point.index()).unwrap();
775 cell_connectivity_entry[local_point_index] = Some(vertex_index);
776 }
777 });
778
779 let mut additional_vertices = new_map();
781 for (flat_cell_index, cell_vertices) in cells.iter_mut() {
782 let cell = grid.try_unflatten_cell_index(*flat_cell_index).unwrap();
783
784 for (local_point_index, vertex) in cell_vertices.iter_mut().enumerate() {
785 if vertex.is_none() {
786 let point = cell.global_point_index_of(local_point_index).unwrap();
788 let flat_point_index = grid.flatten_point_index(&point);
789
790 let vertex_entry =
792 additional_vertices
793 .entry(flat_point_index)
794 .or_insert_with(|| {
795 let point_coords = grid.point_coordinates(&point);
796 let vertex_index = mesh.vertices.len();
797 mesh.vertices.push(point_coords);
798 values.push(default_value);
799
800 vertex_index
801 });
802
803 *vertex = Some(*vertex_entry);
804 }
805 }
806 }
807
808 mesh.cells.reserve(cells.len());
810 for (_, cell_vertices) in cells.iter() {
811 mesh.cells.push([
812 cell_vertices[0].unwrap(),
813 cell_vertices[1].unwrap(),
814 cell_vertices[2].unwrap(),
815 cell_vertices[3].unwrap(),
816 cell_vertices[4].unwrap(),
817 cell_vertices[5].unwrap(),
818 cell_vertices[6].unwrap(),
819 cell_vertices[7].unwrap(),
820 ]);
821 }
822
823 MeshWithData::new(mesh).with_point_data(OwnedMeshAttribute::new_real_scalar(
824 "density".to_string(),
825 values,
826 ))
827}