Add an id to each octree node that is also exported as cell data

Author: w1th0utnam3

splashsurf/src/reconstruction.rs

@@ -7,7 +7,6 @@ use log::info;
 use rayon::prelude::*;
 use splashsurf_lib::mesh::PointCloud3d;
 use splashsurf_lib::profile;
-use splashsurf_lib::vtkio::model::UnstructuredGridPiece;
 use splashsurf_lib::{density_map, Index, Real};
 use std::convert::TryFrom;
 use std::path::PathBuf;
@@ -564,7 +563,11 @@ pub(crate) fn reconstruction_pipeline_generic<I: Index, R: Real>(
     if let Some(output_octree_file) = &paths.output_octree_file {
         info!("Writing octree to \"{}\"...", output_octree_file.display());
         io::vtk_format::write_vtk(
-            UnstructuredGridPiece::from(&reconstruction.octree().unwrap().hexmesh(grid, false)),
+            reconstruction
+                .octree()
+                .unwrap()
+                .hexmesh(grid, true)
+                .to_dataset(),
             output_octree_file,
             "mesh",
         )

splashsurf_lib/src/density_map.rs

@@ -20,7 +20,7 @@
 
 use crate::aabb::AxisAlignedBoundingBox3d;
 use crate::kernel::DiscreteSquaredDistanceCubicKernel;
-use crate::mesh::{HexMesh3d, MeshWithPointData};
+use crate::mesh::{HexMesh3d, MeshWithData};
 use crate::uniform_grid::{OwningSubdomainGrid, Subdomain, UniformGrid};
 use crate::utils::{ChunkSize, ParallelPolicy};
 use crate::{new_map, profile, HashState, Index, MapType, ParallelMapType, Real};
@@ -781,7 +781,7 @@ pub fn sparse_density_map_to_hex_mesh<I: Index, R: Real>(
     density_map: &DensityMap<I, R>,
     grid: &UniformGrid<I, R>,
     default_value: R,
-) -> MeshWithPointData<HexMesh3d<R>, R> {
+) -> MeshWithData<HexMesh3d<R>, R> {
     profile!("sparse_density_map_to_hex_mesh");
 
     let mut mesh = HexMesh3d {
@@ -859,5 +859,5 @@ pub fn sparse_density_map_to_hex_mesh<I: Index, R: Real>(
         ]);
     }
 
-    MeshWithPointData { mesh, data: values }
+    MeshWithData::with_point_data(mesh, values)
 }

splashsurf_lib/src/mesh.rs

@@ -230,28 +230,71 @@ pub struct PointCloud3d<R: Real> {
 
 /// A mesh with attached vertex or point data
 #[derive(Clone, Debug, Default)]
-pub struct MeshWithPointData<MeshT, DataT> {
+pub struct MeshWithData<MeshT, PointDataT, CellDataT = ()> {
     /// The mesh geometry itself
     pub mesh: MeshT,
     /// Data attached to each vertex or point of the mesh
-    pub data: Vec<DataT>,
+    pub point_data: Vec<PointDataT>,
+    /// Data attached to each cell of the mesh
+    pub cell_data: Vec<CellDataT>,
+}
+
+impl<MeshT, PointDataT, CellDataT> MeshWithData<MeshT, PointDataT, CellDataT> {
+    /// Creates a new mesh the given point data
+    pub fn with_point_data<PointData: Into<Vec<PointDataT>>>(
+        mesh: MeshT,
+        point_data: PointData,
+    ) -> Self {
+        Self {
+            mesh,
+            point_data: point_data.into(),
+            cell_data: vec![],
+        }
+    }
+
+    /// Creates a new mesh the given cell data
+    pub fn with_cell_data<CellData: Into<Vec<CellDataT>>>(
+        mesh: MeshT,
+        cell_data: CellData,
+    ) -> Self {
+        Self {
+            mesh,
+            point_data: vec![],
+            cell_data: cell_data.into(),
+        }
+    }
 }
 
 #[cfg(feature = "vtk_extras")]
 use vtkio::model::{Attribute, UnstructuredGridPiece};
 
 #[cfg(feature = "vtk_extras")]
-impl<'a, MeshT: 'a, DataT> MeshWithPointData<MeshT, DataT>
+impl<'a, MeshT: 'a, PointDataT> MeshWithData<MeshT, PointDataT, ()>
 where
     &'a MeshT: Into<UnstructuredGridPiece>,
-    DataT: Real,
+    PointDataT: Real,
 {
     pub fn to_dataset(&'a self) -> UnstructuredGridPiece {
         let mut grid_piece: UnstructuredGridPiece = (&self.mesh).into();
         grid_piece
             .data
             .point
-            .push(Attribute::scalars("density", 1).with_data(self.data.clone()));
+            .push(Attribute::scalars("density", 1).with_data(self.point_data.clone()));
+        grid_piece
+    }
+}
+
+#[cfg(feature = "vtk_extras")]
+impl<'a, MeshT: 'a> MeshWithData<MeshT, (), u64>
+where
+    &'a MeshT: Into<UnstructuredGridPiece>,
+{
+    pub fn to_dataset(&'a self) -> UnstructuredGridPiece {
+        let mut grid_piece: UnstructuredGridPiece = (&self.mesh).into();
+        grid_piece
+            .data
+            .cell
+            .push(Attribute::scalars("node_id", 1).with_data(self.cell_data.clone()));
         grid_piece
     }
 }

splashsurf_lib/src/octree.rs

@@ -1,7 +1,9 @@
 //! Octree for spatially partitioning particle sets
 
 use crate::generic_tree::*;
-use crate::mesh::{HexMesh3d, TriMesh3d};
+use crate::marching_cubes::SurfacePatch;
+use crate::mesh::{HexMesh3d, MeshWithData, TriMesh3d};
+use crate::topology::{Axis, Direction};
 use crate::uniform_grid::{PointIndex, UniformGrid};
 use crate::utils::{ChunkSize, ParallelPolicy};
 use crate::{
@@ -14,7 +16,9 @@ use nalgebra::Vector3;
 use octant_helper::{HalfspaceFlags, Octant, OctantAxisDirections};
 use rayon::prelude::*;
 use smallvec::SmallVec;
+use split_criterion::{default_split_criterion, LeafSplitCriterion};
 use std::cell::RefCell;
+use std::sync::atomic::{AtomicUsize, Ordering};
 use thread_local::ThreadLocal;
 
 // TODO: Make margin an Option
@@ -31,12 +35,17 @@ pub enum SubdivisionCriterion {
 /// Data structure for octree based spatial subdivision of particles sets, for tree iteration/visitation use the [`root`](Self::root) [`OctreeNode`]
 #[derive(Clone, Debug)]
 pub struct Octree<I: Index, R: Real> {
+    /// Root node of the tree
     root: OctreeNode<I, R>,
+    /// Counter for assigning ids to subdivided nodes
+    next_id: usize,
 }
 
 /// Represents a node in the octree hierarchy and stores child nodes, implements tree iteration/visitation from the [`generic_tree`](crate::generic_tree) module
 #[derive(Clone, Debug)]
 pub struct OctreeNode<I: Index, R: Real> {
+    /// Id of the node used to identify it for debugging
+    id: usize,
     /// All child nodes of this octree node
     children: ArrayVec<[Box<Self>; 8]>,
     /// Lower corner point of the octree node on the background grid
@@ -108,15 +117,12 @@ impl<I: Index, R: Real> SurfacePatchWrapper<I, R> {
 
 type OctreeNodeParticleStorage = SmallVec<[usize; 6]>;
 
-use crate::marching_cubes::SurfacePatch;
-use crate::topology::{Axis, Direction};
-use split_criterion::{default_split_criterion, LeafSplitCriterion};
-
 impl<I: Index, R: Real> Octree<I, R> {
     /// Creates a new octree with a single leaf node containing all vertices
     pub fn new(grid: &UniformGrid<I, R>, n_particles: usize) -> Self {
         Self {
             root: OctreeNode::new_root(grid, n_particles),
+            next_id: 0,
         }
     }
 
@@ -135,6 +141,8 @@ impl<I: Index, R: Real> Octree<I, R> {
     ) -> Self {
         let mut tree = Octree::new(&grid, particle_positions.len());
 
+        println!("Margin: {}", margin);
+
         if enable_multi_threading {
             tree.subdivide_recursively_margin_par(
                 grid,
@@ -156,11 +164,6 @@ impl<I: Index, R: Real> Octree<I, R> {
         tree
     }
 
-    /// Creates a new octree with the given node as root
-    pub fn with_root(root: OctreeNode<I, R>) -> Self {
-        Self { root }
-    }
-
     /// Returns a reference to the root node of the octree
     pub fn root(&self) -> &OctreeNode<I, R> {
         &self.root
@@ -188,15 +191,17 @@ impl<I: Index, R: Real> Octree<I, R> {
             enable_stitching,
         );
 
+        let next_id = AtomicUsize::new(0);
         self.root.visit_mut_bfs(|node| {
             // Stop recursion if split criterion is not fulfilled
             if !split_criterion.split_leaf(node) {
                 return;
             }
 
             // Perform one octree split on the node
-            node.subdivide_with_margin(grid, particle_positions, margin);
+            node.subdivide_with_margin(grid, particle_positions, margin, &next_id);
         });
+        self.next_id = next_id.into_inner();
     }
 
     /// Subdivide the octree recursively and in parallel using the given splitting criterion and a margin to add ghost particles
@@ -217,39 +222,55 @@ impl<I: Index, R: Real> Octree<I, R> {
         );
         let parallel_policy = ParallelPolicy::default();
 
-        let visitor = move |node: &mut OctreeNode<I, R>| {
-            // Stop recursion if split criterion is not fulfilled
-            if !split_criterion.split_leaf(node) {
-                return;
-            }
+        let next_id = AtomicUsize::new(0);
+        let visitor = {
+            let next_id = &next_id;
+            move |node: &mut OctreeNode<I, R>| {
+                // Stop recursion if split criterion is not fulfilled
+                if !split_criterion.split_leaf(node) {
+                    return;
+                }
 
-            // Perform one octree split on the leaf
-            if node
-                .data
-                .particle_set()
-                .expect("Node is not a leaf")
-                .particles
-                .len()
-                < parallel_policy.min_task_size
-            {
-                node.subdivide_with_margin(grid, particle_positions, margin);
-            } else {
-                node.subdivide_with_margin_par(grid, particle_positions, margin, &parallel_policy);
+                // Perform one octree split on the leaf
+                if node
+                    .data
+                    .particle_set()
+                    .expect("Node is not a leaf")
+                    .particles
+                    .len()
+                    < parallel_policy.min_task_size
+                {
+                    node.subdivide_with_margin(grid, particle_positions, margin, &next_id);
+                } else {
+                    node.subdivide_with_margin_par(
+                        grid,
+                        particle_positions,
+                        margin,
+                        &parallel_policy,
+                        &next_id,
+                    );
+                }
             }
         };
 
         self.root.par_visit_mut_bfs(visitor);
+        self.next_id = next_id.into_inner();
     }
 
     /// Constructs a hex mesh visualizing the cells of the octree, may contain hanging and duplicate vertices as cells are not connected
-    pub fn hexmesh(&self, grid: &UniformGrid<I, R>, only_non_empty: bool) -> HexMesh3d<R> {
+    pub fn hexmesh(
+        &self,
+        grid: &UniformGrid<I, R>,
+        only_non_empty: bool,
+    ) -> MeshWithData<HexMesh3d<R>, (), u64> {
         profile!("convert octree into hexmesh");
 
         let mut mesh = HexMesh3d {
             vertices: Vec::new(),
             cells: Vec::new(),
         };
 
+        let mut ids = Vec::new();
         self.root.dfs_iter().for_each(|node| {
             if node.children().is_empty() {
                 if only_non_empty
@@ -290,16 +311,18 @@ impl<I: Index, R: Real> Octree<I, R> {
 
                 mesh.vertices.extend(vertices);
                 mesh.cells.push(cell);
+                ids.push(node.id as u64);
             }
         });
 
-        mesh
+        assert_eq!(mesh.cells.len(), ids.len());
+        MeshWithData::with_cell_data(mesh, ids)
     }
 }
 
 impl<I: Index, R: Real> OctreeNode<I, R> {
-    pub fn new(min_corner: PointIndex<I>, max_corner: PointIndex<I>) -> Self {
-        Self::with_data(min_corner, max_corner, NodeData::None)
+    pub fn new(id: usize, min_corner: PointIndex<I>, max_corner: PointIndex<I>) -> Self {
+        Self::with_data(id, min_corner, max_corner, NodeData::None)
     }
 
     fn new_root(grid: &UniformGrid<I, R>, n_particles: usize) -> Self {
@@ -312,6 +335,7 @@ impl<I: Index, R: Real> OctreeNode<I, R> {
         ];
 
         Self::with_data(
+            0,
             grid.get_point(min_point)
                 .expect("Cannot get lower corner of grid"),
             grid.get_point(max_point)
@@ -321,18 +345,25 @@ impl<I: Index, R: Real> OctreeNode<I, R> {
     }
 
     fn with_data(
+        id: usize,
         min_corner: PointIndex<I>,
         max_corner: PointIndex<I>,
         data: NodeData<I, R>,
     ) -> Self {
         Self {
+            id,
             children: Default::default(),
             min_corner,
             max_corner,
             data,
         }
     }
 
+    /// Returns the id of the node
+    pub fn id(&self) -> usize {
+        self.id
+    }
+
     /// Returns a reference to the data stored in the node
     pub fn data(&self) -> &NodeData<I, R> {
         &self.data
@@ -385,6 +416,7 @@ impl<I: Index, R: Real> OctreeNode<I, R> {
         grid: &UniformGrid<I, R>,
         particle_positions: &[Vector3<R>],
         margin: R,
+        next_id: &AtomicUsize,
     ) {
         // Convert node body from Leaf to Children
         if let NodeData::ParticleSet(particle_set) = self.data.take() {
@@ -460,6 +492,7 @@ impl<I: Index, R: Real> OctreeNode<I, R> {
                 assert_eq!(octant_particles.len(), octant_particle_count);
 
                 let child = Box::new(OctreeNode::with_data(
+                    next_id.fetch_add(1, Ordering::SeqCst),
                     min_corner,
                     max_corner,
                     NodeData::new_particle_set(
@@ -485,6 +518,7 @@ impl<I: Index, R: Real> OctreeNode<I, R> {
         particle_positions: &[Vector3<R>],
         margin: R,
         parallel_policy: &ParallelPolicy,
+        next_id: &AtomicUsize,
     ) {
         // Convert node body from Leaf to Children
         if let NodeData::ParticleSet(particle_set) = self.data.take() {
@@ -594,6 +628,7 @@ impl<I: Index, R: Real> OctreeNode<I, R> {
                         assert_eq!(octant_particles.len(), octant_particle_count);
 
                         let child = Box::new(OctreeNode::with_data(
+                            next_id.fetch_add(1, Ordering::SeqCst),
                             min_corner,
                             max_corner,
                             NodeData::new_particle_set(

splashsurf_lib/src/profiling.rs

@@ -143,6 +143,7 @@ impl Profiler {
 
     /// Enter a scope with the given name and push it onto the stack of scopes. It will be a child scope of the current top of the stack.
     pub fn enter(&mut self, name: &'static str) -> (Guard, ScopeId) {
+        // TODO: If the scope on top of the stack has the same name, use the next one as parent, to avoid huge chains for recursive functions
         let id = self.new_id(name, self.scope_stack.last());
         self.enter_with_id(name, id)
     }