Py: Add plain marching cubes functionality

Author: w1th0utnam3

pysplashsurf/pysplashsurf/pysplashsurf.pyi

@@ -353,6 +353,11 @@ 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]:
+    r"""
+    Performs a standard marching cubes triangulation of a 3D array of values
+    """
+
 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]:
     r"""
     Performs simplification on the given mesh inspired by the "Compact Contouring"/"Mesh displacement" approach by Doug Moore and Joe Warren

pysplashsurf/src/lib.rs

@@ -44,6 +44,7 @@ fn pysplashsurf(m: &Bound<'_, PyModule>) -> PyResult<()> {
 
     m.add_function(wrap!(reconstruction::reconstruct_surface, m)?)?;
     m.add_function(wrap!(marching_cubes::check_mesh_consistency, m)?)?;
+    m.add_function(wrap!(marching_cubes::marching_cubes, m)?)?;
     m.add_function(wrap!(postprocessing::marching_cubes_cleanup, m)?)?;
     m.add_function(wrap!(postprocessing::convert_tris_to_quads, m)?)?;
     m.add_function(wrap!(postprocessing::barnacle_decimation, m)?)?;

pysplashsurf/src/marching_cubes.rs

@@ -1,10 +1,14 @@
-use numpy::PyUntypedArray;
+use numpy::prelude::*;
+use numpy::{Element, PyArray3, PyUntypedArray};
 use pyo3::prelude::*;
 use pyo3_stub_gen::derive::*;
+use splashsurf_lib::nalgebra::Vector3;
+use splashsurf_lib::{DensityMap, Real, UniformGrid};
 
 use crate::mesh::{PyTriMesh3d, get_triangle_mesh_generic};
 use crate::uniform_grid::PyUniformGrid;
-use crate::utils::*;
+use crate::utils;
+use crate::utils::IndexT;
 
 /// Checks the consistency of a reconstructed surface mesh (watertightness, manifoldness), optionally returns a string with details if problems are found
 #[gen_stub_pyfunction]
@@ -22,7 +26,7 @@ pub fn check_mesh_consistency<'py>(
     let py = mesh.py();
 
     // Try to extract the triangle mesh;
-    let mesh = get_triangle_mesh_generic(&mesh).ok_or_else(pyerr_only_triangle_mesh)?;
+    let mesh = get_triangle_mesh_generic(&mesh).ok_or_else(utils::pyerr_only_triangle_mesh)?;
     let mesh = mesh.borrow(py);
 
     if let (Some(grid), Some(mesh)) = (grid.as_f32(), mesh.as_f32()) {
@@ -44,6 +48,66 @@ pub fn check_mesh_consistency<'py>(
         )
         .err())
     } else {
-        Err(pyerr_scalar_type_mismatch())
+        Err(utils::pyerr_scalar_type_mismatch())
+    }
+}
+
+/// Performs a standard marching cubes triangulation of a 3D array of values
+#[gen_stub_pyfunction]
+#[pyfunction]
+#[pyo3(name = "marching_cubes")]
+#[pyo3(signature = (values, *, cube_size, iso_surface_threshold, translation = None))]
+pub fn marching_cubes<'py>(
+    values: &Bound<'py, PyUntypedArray>,
+    cube_size: f64,
+    iso_surface_threshold: f64,
+    translation: Option<[f64; 3]>,
+) -> PyResult<(PyTriMesh3d, PyUniformGrid)> {
+    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,
+        translation: Option<[R; 3]>,
+    ) -> PyResult<(PyTriMesh3d, PyUniformGrid)> {
+        let shape = values.shape();
+        let translation = Vector3::from(translation.unwrap_or([R::zero(); 3]));
+        let n_cells_per_dim = [
+            shape[0] as IndexT - 1,
+            shape[1] as IndexT - 1,
+            shape[2] as IndexT - 1,
+        ];
+
+        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 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)?,
+        ))
+    }
+
+    if let Ok(values) = values.downcast::<PyArray3<f32>>() {
+        triangulate_density_map_generic(
+            &values,
+            cube_size as f32,
+            iso_surface_threshold as f32,
+            translation.map(|t| t.map(|t| t as f32)),
+        )
+    } else if let Ok(values) = values.downcast::<PyArray3<f64>>() {
+        triangulate_density_map_generic(&values, cube_size, iso_surface_threshold, translation)
+    } else {
+        Err(utils::pyerr_unsupported_scalar())
     }
 }

pysplashsurf/test_sdf.py

@@ -0,0 +1,45 @@
+import os
+import numpy as np
+import pysplashsurf
+import time
+
+devnull = open(os.devnull, 'w')
+
+bench_results = {}
+
+def bench_fn(name, iters, fn):
+    print(f'Running benchmark "{name}"... ', end=' ', flush=True)
+    start = time.perf_counter_ns()
+    result = None
+    for _ in range(iters):
+        result = fn()
+    elapsed = time.perf_counter_ns() - start
+    avg_ms = (elapsed / iters) / 1000.0 / 1000.0
+    bench_results[name] = (avg_ms, iters)
+    print(f"Done. Average runtime per iteration: {avg_ms:.3f} ms over {iters} iterations")
+    return result
+
+def print_bench_results():
+    for name, (avg_ms, iters) in bench_results.items():
+        print(f"Benchmark '{name}': {avg_ms:.3f} ms over {iters} iterations")
+
+# --- Check marching cubes
+
+radius = 1.0
+num_verts = 100
+
+grid_size = radius * 2.2
+dx = grid_size / (num_verts - 1)
+
+translation = -0.5 * grid_size
+
+def make_sdf():
+    coords = np.arange(num_verts, dtype=np.float32) * dx + translation
+    x, y, z = np.meshgrid(coords, coords, coords, indexing='ij')
+    sdf = np.sqrt(x**2 + y**2 + z**2) - radius
+    return sdf
+
+sdf = bench_fn("make_sdf", 1, make_sdf)
+
+mesh = pysplashsurf.marching_cubes(sdf, cube_size=dx, iso_surface_threshold=0.0, translation=[translation]*3)
+mesh.write_to_file("test_marching_cubes.vtk")

pysplashsurf/tests/test_basic.py

@@ -205,6 +205,8 @@ def test_check_consistency():
         pysplashsurf.check_mesh_consistency(mesh_with_data, reconstruction.grid) is None
     )
 
+    # TODO: Delete some triangles and check for failure
+
 
 def test_tris_to_quads():
     particles = np.array(meshio.read(VTK_PATH).points, dtype=np.float32)

splashsurf_lib/src/density_map.rs

@@ -217,11 +217,12 @@ pub fn parallel_compute_particle_densities<I: Index, R: Real>(
 /// A sparse density map
 ///
 /// 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 of the compact support of any particles).
+/// 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> {
     Standard(MapType<I, R>),
     DashMap(ReadDashMap<I, R, HashState>),
+    Dense(Vec<R>),
 }
 
 impl<I: Index, R: Real> Default for DensityMap<I, R> {
@@ -242,12 +243,24 @@ impl<I: Index, R: Real> From<ParallelMapType<I, R>> for DensityMap<I, R> {
     }
 }
 
+impl<I: Index, R: Real> From<Vec<R>> for DensityMap<I, R> {
+    fn from(values: Vec<R>) -> Self {
+        Self::Dense(values)
+    }
+}
+
 impl<I: Index, R: Real> DensityMap<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 {
             DensityMap::Standard(map) => map.iter().map(|(&i, &r)| (i, r)).collect(),
             DensityMap::DashMap(map) => map.iter().map(|(&i, &r)| (i, r)).collect(),
+            DensityMap::Dense(values) => values
+                .iter()
+                .copied()
+                .enumerate()
+                .map(|(i, r)| (I::from_usize(i).unwrap(), r))
+                .collect(),
         }
     }
 
@@ -256,6 +269,7 @@ impl<I: Index, R: Real> DensityMap<I, R> {
         match self {
             DensityMap::Standard(map) => map.len(),
             DensityMap::DashMap(map) => map.len(),
+            DensityMap::Dense(values) => values.len(),
         }
     }
 
@@ -264,6 +278,7 @@ impl<I: Index, R: Real> DensityMap<I, R> {
         match self {
             DensityMap::Standard(map) => map.get(&flat_point_index).copied(),
             DensityMap::DashMap(map) => map.get(&flat_point_index).copied(),
+            DensityMap::Dense(values) => values.get(flat_point_index.to_usize()?).copied(),
         }
     }
 
@@ -273,6 +288,9 @@ impl<I: Index, R: Real> DensityMap<I, R> {
         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)),
+            DensityMap::Dense(values) => values.iter().copied().enumerate().for_each(|(i, r)| {
+                f(I::from_usize(i).unwrap(), r);
+            }),
         }
     }
 }