Support borrowed data for DensityMap

Author: w1th0utnam3

pysplashsurf/pysplashsurf/docs/source/api.rst

@@ -19,6 +19,7 @@ Functions
     convert_tris_to_quads
     laplacian_smoothing_normals_parallel
     laplacian_smoothing_parallel
+    marching_cubes
     marching_cubes_cleanup
     neighborhood_search_spatial_hashing_parallel
     reconstruct_surface

pysplashsurf/pysplashsurf/docs/source/functions.rst

@@ -15,6 +15,8 @@ All functions infer float precision based on the input (``np.float32`` or ``np.f
 
 .. autofunction:: laplacian_smoothing_parallel
 
+.. autofunction:: marching_cubes
+
 .. autofunction:: marching_cubes_cleanup
 
 .. autofunction:: neighborhood_search_spatial_hashing_parallel

pysplashsurf/pysplashsurf/pysplashsurf.pyi

@@ -353,9 +353,28 @@ def laplacian_smoothing_parallel(mesh:typing.Union[TriMesh3d, MeshWithData], ver
     The smoothing is performed inplace and modifies the vertices of the given mesh.
     """
 
-def marching_cubes(values:numpy.typing.NDArray[typing.Any], *, cube_size:builtins.float, iso_surface_threshold:builtins.float, translation:typing.Optional[typing.Sequence[builtins.float]]=None) -> tuple[TriMesh3d, UniformGrid]:
+def marching_cubes(values:numpy.typing.NDArray[typing.Any], *, iso_surface_threshold:builtins.float, cube_size:builtins.float, translation:typing.Optional[typing.Sequence[builtins.float]]=None, return_grid:builtins.bool=False) -> typing.Union[TriMesh3d, tuple[TriMesh3d, UniformGrid]]:
     r"""
     Performs a standard marching cubes triangulation of a 3D array of values
+    
+    The array of values has to be a contiguous array with shape ``(nx, ny, nz)``.
+    The iso-surface threshold defines which value is considered to be "on" the surface.
+    The cube size and translation parameters define the scaling and translation of the resulting
+    mesh. Without translation, the value ``values[0, 0, 0]`` is located at coordinates ``(0, 0, 0)``.
+    
+    The values are interpreted as a "density field", meaning that values higher than the iso-surface
+    threshold are considered to be "inside" the surface and values lower than the threshold are
+    considered to be "outside" the surface. This is the opposite convention to an SDF (signed distance field).
+    However, even if values of an SDF are provided as an input, the marching cubes algorithm
+    will still work and produce a watertight surface mesh (if the surface is fully contained in the
+    array).
+    
+    If ``return_grid`` is set to ``True``, the function will return a tuple of the mesh and the
+    uniform grid that was used for the triangulation. This can be used for other functions such as
+    :py:func:`check_mesh_consistency`. Otherwise, only the mesh is returned.
+    
+    The function is currently single-threaded. The SPH surface reconstruction functions :py:func:`reconstruction_pipeline`
+    and :py:func:`reconstruct_surface` improve performance by processing multiple patches in parallel.
     """
 
 def marching_cubes_cleanup(mesh:typing.Union[TriMesh3d, MeshWithData], grid:UniformGrid, *, max_rel_snap_dist:typing.Optional[builtins.float]=None, max_iter:builtins.int=5, keep_vertices:builtins.bool=False) -> typing.Union[TriMesh3d, MeshWithData]:

pysplashsurf/src/marching_cubes.rs

@@ -1,5 +1,6 @@
 use numpy::prelude::*;
 use numpy::{Element, PyArray3, PyUntypedArray};
+use pyo3::IntoPyObjectExt;
 use pyo3::prelude::*;
 use pyo3_stub_gen::derive::*;
 use splashsurf_lib::nalgebra::Vector3;
@@ -53,24 +54,47 @@ pub fn check_mesh_consistency<'py>(
 }
 
 /// Performs a standard marching cubes triangulation of a 3D array of values
+///
+/// The array of values has to be a contiguous array with shape ``(nx, ny, nz)``.
+/// The iso-surface threshold defines which value is considered to be "on" the surface.
+/// The cube size and translation parameters define the scaling and translation of the resulting
+/// mesh. Without translation, the value ``values[0, 0, 0]`` is located at coordinates ``(0, 0, 0)``.
+///
+/// The values are interpreted as a "density field", meaning that values higher than the iso-surface
+/// threshold are considered to be "inside" the surface and values lower than the threshold are
+/// considered to be "outside" the surface. This is the opposite convention to an SDF (signed distance field).
+/// However, even if values of an SDF are provided as an input, the marching cubes algorithm
+/// will still work and produce a watertight surface mesh (if the surface is fully contained in the
+/// array).
+///
+/// If ``return_grid`` is set to ``True``, the function will return a tuple of the mesh and the
+/// uniform grid that was used for the triangulation. This can be used for other functions such as
+/// :py:func:`check_mesh_consistency`. Otherwise, only the mesh is returned.
+///
+/// The function is currently single-threaded. The SPH surface reconstruction functions :py:func:`reconstruction_pipeline`
+/// and :py:func:`reconstruct_surface` improve performance by processing multiple patches in parallel.
 #[gen_stub_pyfunction]
 #[pyfunction]
 #[pyo3(name = "marching_cubes")]
-#[pyo3(signature = (values, *, cube_size, iso_surface_threshold, translation = None))]
+#[pyo3(signature = (values, *, iso_surface_threshold, cube_size, translation = None, return_grid = false))]
+#[gen_stub(override_return_type(type_repr="typing.Union[TriMesh3d, tuple[TriMesh3d, UniformGrid]]", imports=()))]
 pub fn marching_cubes<'py>(
     values: &Bound<'py, PyUntypedArray>,
-    cube_size: f64,
     iso_surface_threshold: f64,
+    cube_size: f64,
     translation: Option<[f64; 3]>,
-) -> PyResult<(PyTriMesh3d, PyUniformGrid)> {
+    return_grid: bool,
+) -> PyResult<Py<PyAny>> {
     assert_eq!(values.shape().len(), 3, "values must be a 3D array");
 
     fn triangulate_density_map_generic<'py, R: Real + Element>(
         values: &Bound<'py, PyArray3<R>>,
-        cube_size: R,
         iso_surface_threshold: R,
+        cube_size: R,
         translation: Option<[R; 3]>,
-    ) -> PyResult<(PyTriMesh3d, PyUniformGrid)> {
+        return_grid: bool,
+    ) -> PyResult<Py<PyAny>> {
+        let py = values.py();
         let shape = values.shape();
         let translation = Vector3::from(translation.unwrap_or([R::zero(); 3]));
         let n_cells_per_dim = [
@@ -82,31 +106,42 @@ pub fn marching_cubes<'py>(
         let grid = UniformGrid::new(&translation, &n_cells_per_dim, cube_size)
             .map_err(anyhow::Error::from)?;
 
-        // TODO: Replace with borrow
-        let values = values.try_readonly()?.as_slice()?.to_vec();
-        let density_map = DensityMap::from(values);
+        let values = values.try_readonly()?;
+        let density_map = DensityMap::from(values.as_slice()?);
 
         let mesh = splashsurf_lib::marching_cubes::triangulate_density_map(
             &grid,
             &density_map,
             iso_surface_threshold,
         )
         .map_err(anyhow::Error::from)?;
-        Ok((
-            PyTriMesh3d::try_from_generic(mesh)?,
-            PyUniformGrid::try_from_generic(grid)?,
-        ))
+
+        let mesh = PyTriMesh3d::try_from_generic(mesh)?;
+        let grid = PyUniformGrid::try_from_generic(grid)?;
+
+        if return_grid {
+            (mesh, grid).into_py_any(py)
+        } else {
+            mesh.into_py_any(py)
+        }
     }
 
     if let Ok(values) = values.downcast::<PyArray3<f32>>() {
         triangulate_density_map_generic(
             &values,
-            cube_size as f32,
             iso_surface_threshold as f32,
+            cube_size as f32,
             translation.map(|t| t.map(|t| t as f32)),
+            return_grid,
         )
     } else if let Ok(values) = values.downcast::<PyArray3<f64>>() {
-        triangulate_density_map_generic(&values, cube_size, iso_surface_threshold, translation)
+        triangulate_density_map_generic(
+            &values,
+            iso_surface_threshold,
+            cube_size,
+            translation,
+            return_grid,
+        )
     } else {
         Err(utils::pyerr_unsupported_scalar())
     }

pysplashsurf/tests/test_sdf.py

@@ -21,7 +21,7 @@ def make_sdf():
 
     # Note: Currently this reconstruction assumes that inside the surface values get bigger (like a density function)
     mesh, grid = pysplashsurf.marching_cubes(
-        sdf, cube_size=dx, iso_surface_threshold=0.0, translation=[translation] * 3
+        sdf, iso_surface_threshold=0.0, cube_size=dx, translation=[translation] * 3, return_grid=True
     )
 
     assert len(mesh.vertices) > 0

splashsurf_lib/src/density_map.rs

@@ -12,7 +12,7 @@
 //! In case of a sparse density map, the values are stored in a hashmap. The keys are so called
 //! "flat point indices". These are computed from the background grid point coordinates `(i,j,k)`
 //! analogous to multidimensional array index flattening. That means for a grid with dimensions
-//! `[n_x, n_y, n_z]`, the flat point index is given by the expression `i*n_x + j*n_y + k*n_z`.
+//! `[n_x, n_y, n_z]`, the flat point index is given by the expression `i*n_y*n_z + j*n_z + k`.
 //! For these point index operations, the [`UniformGrid`] is used.
 //!
 //! Note that all density mapping functions always use the global background grid for flat point
@@ -29,6 +29,7 @@ use dashmap::ReadOnlyView as ReadDashMap;
 use log::{info, trace, warn};
 use nalgebra::Vector3;
 use rayon::prelude::*;
+use std::borrow::Cow;
 use std::cell::RefCell;
 use thiserror::Error as ThisError;
 use thread_local::ThreadLocal;
@@ -219,37 +220,46 @@ pub fn parallel_compute_particle_densities<I: Index, R: Real>(
 /// The density map contains values for all points of the background grid where the density is not
 /// trivially zero (which is the case when a point is outside the compact support of any particles).
 #[derive(Clone, Debug)]
-pub enum DensityMap<I: Index, R: Real> {
+pub enum DensityMap<'a, I: Index, R: Real> {
     Standard(MapType<I, R>),
     DashMap(ReadDashMap<I, R, HashState>),
-    Dense(Vec<R>),
+    Dense(Cow<'a, [R]>),
 }
 
-impl<I: Index, R: Real> Default for DensityMap<I, R> {
+/// Owned version of [`DensityMap`] (with static lifetime)
+pub type OwnedDensityMap<I, R> = DensityMap<'static, I, R>;
+
+impl<I: Index, R: Real> Default for OwnedDensityMap<I, R> {
     fn default() -> Self {
         DensityMap::Standard(MapType::default())
     }
 }
 
-impl<I: Index, R: Real> From<MapType<I, R>> for DensityMap<I, R> {
+impl<I: Index, R: Real> From<MapType<I, R>> for OwnedDensityMap<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> {
+impl<I: Index, R: Real> From<ParallelMapType<I, R>> for OwnedDensityMap<I, R> {
     fn from(map: ParallelMapType<I, R>) -> Self {
         Self::DashMap(map.into_read_only())
     }
 }
 
-impl<I: Index, R: Real> From<Vec<R>> for DensityMap<I, R> {
+impl<I: Index, R: Real> From<Vec<R>> for DensityMap<'static, I, R> {
     fn from(values: Vec<R>) -> Self {
-        Self::Dense(values)
+        Self::Dense(values.into())
+    }
+}
+
+impl<'a, I: Index, R: Real> From<&'a [R]> for DensityMap<'a, I, R> {
+    fn from(values: &'a [R]) -> Self {
+        Self::Dense(values.into())
     }
 }
 
-impl<I: Index, R: Real> DensityMap<I, R> {
+impl<'a, I: Index, R: Real> DensityMap<'a, I, R> {
     /// Converts the contained map into a vector of tuples of (flat_point_index, density)
     pub fn to_vec(&self) -> Vec<(I, R)> {
         match self {
@@ -357,7 +367,7 @@ pub fn sequential_generate_sparse_density_map<I: Index, R: Real>(
     particle_rest_mass: R,
     compact_support_radius: R,
     cube_size: R,
-) -> Result<DensityMap<I, R>, DensityMapError<R>> {
+) -> Result<OwnedDensityMap<I, R>, DensityMapError<R>> {
     profile!("sequential_generate_sparse_density_map");
 
     let mut sparse_densities = new_map();
@@ -404,7 +414,7 @@ pub fn parallel_generate_sparse_density_map<I: Index, R: Real>(
     particle_rest_mass: R,
     compact_support_radius: R,
     cube_size: R,
-) -> Result<DensityMap<I, R>, DensityMapError<R>> {
+) -> Result<OwnedDensityMap<I, R>, DensityMapError<R>> {
     profile!("parallel_generate_sparse_density_map");
 
     // Each thread will write to its own local density map