Py: Methods for UniformGrid, add tests for many functions

Author: w1th0utnam3

pysplashsurf/pysplashsurf/pysplashsurf.pyi

@@ -270,7 +270,26 @@ class UniformGrid:
     r"""
     Struct containing the parameters of the uniform grid used for the surface reconstruction
     """
-    ...
+    @property
+    def aabb(self) -> Aabb3d:
+        r"""
+        The AABB of the grid containing all marching cubes vertices influenced by the particle kernels
+        """
+    @property
+    def cell_size(self) -> builtins.float:
+        r"""
+        Returns the cell size of the uniform grid (the marching cubes voxel size)
+        """
+    @property
+    def npoints_per_dim(self) -> builtins.list[builtins.int]:
+        r"""
+        Returns the number of points (marching cubes vertices) per dimension in the uniform grid
+        """
+    @property
+    def ncells_per_dim(self) -> builtins.list[builtins.int]:
+        r"""
+        Returns the number of cells (marching cubes voxels) per dimension in the uniform grid
+        """
 
 class VertexVertexConnectivity:
     r"""
@@ -362,5 +381,3 @@ def reconstruction_pipeline(particles:numpy.typing.NDArray[typing.Any], *, attri
     
     Note that smoothing length and cube size are given in multiples of the particle radius.
     """
-
-def triangulate_density_map(values:numpy.typing.NDArray[typing.Any], grid:UniformGrid, *, iso_surface_threshold:builtins.float) -> TriMesh3d: ...

pysplashsurf/src/mesh.rs

@@ -339,7 +339,7 @@ where
 /// Enum specifying the type of mesh wrapped by a ``MeshWithData``
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 #[gen_stub_pyclass_enum]
-#[pyclass]
+#[pyclass(eq)]
 pub enum MeshType {
     /// 3D triangle mesh
     Tri3d,

pysplashsurf/src/uniform_grid.rs

@@ -1,11 +1,11 @@
+use crate::aabb::PyAabb3d;
+use crate::utils;
+use crate::utils::{IndexT, enum_wrapper_impl_from};
 use pyo3::exceptions::PyTypeError;
 use pyo3::prelude::*;
 use pyo3_stub_gen::derive::*;
 use splashsurf_lib::{Real, UniformGrid};
 
-use crate::utils;
-use crate::utils::{IndexT, enum_wrapper_impl_from};
-
 enum PyUniformGridData {
     F32(UniformGrid<IndexT, f32>),
     F64(UniformGrid<IndexT, f64>),
@@ -45,3 +45,43 @@ impl PyUniformGrid {
         }
     }
 }
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl PyUniformGrid {
+    /// The AABB of the grid containing all marching cubes vertices influenced by the particle kernels
+    #[getter]
+    pub fn aabb(&self) -> PyAabb3d {
+        match &self.inner {
+            PyUniformGridData::F32(grid) => PyAabb3d::from(grid.aabb().clone()),
+            PyUniformGridData::F64(grid) => PyAabb3d::from(grid.aabb().clone()),
+        }
+    }
+
+    /// Returns the cell size of the uniform grid (the marching cubes voxel size)
+    #[getter]
+    pub fn cell_size(&self) -> f64 {
+        match &self.inner {
+            PyUniformGridData::F32(grid) => grid.cell_size() as f64,
+            PyUniformGridData::F64(grid) => grid.cell_size(),
+        }
+    }
+
+    /// Returns the number of points (marching cubes vertices) per dimension in the uniform grid
+    #[getter]
+    pub fn npoints_per_dim(&self) -> [IndexT; 3] {
+        match &self.inner {
+            PyUniformGridData::F32(grid) => grid.points_per_dim().clone(),
+            PyUniformGridData::F64(grid) => grid.points_per_dim().clone(),
+        }
+    }
+
+    /// Returns the number of cells (marching cubes voxels) per dimension in the uniform grid
+    #[getter]
+    pub fn ncells_per_dim(&self) -> [IndexT; 3] {
+        match &self.inner {
+            PyUniformGridData::F32(grid) => grid.cells_per_dim().clone(),
+            PyUniformGridData::F64(grid) => grid.cells_per_dim().clone(),
+        }
+    }
+}

pysplashsurf/tests/test_basic.py

@@ -1,6 +1,12 @@
 import pysplashsurf
 import numpy as np
 import meshio
+import os.path
+import pathlib
+import tempfile
+
+DIR = pathlib.Path(__file__).parent.resolve()
+VTK_PATH = DIR.joinpath("ParticleData_Random_1000.vtk")
 
 
 def test_aabb_class():
@@ -30,3 +36,272 @@ def test_aabb_class():
     assert aabb.contains_point([0.0, 0.0, 0.0])
     assert not aabb.contains_point([2.0, 1.0, 4.2])
     assert not aabb.contains_point([1.0, -1.0, 5.0])
+
+
+def impl_basic_test(dtype):
+    particles = np.array(meshio.read(VTK_PATH).points, dtype=dtype)
+
+    mesh_with_data, reconstruction = pysplashsurf.reconstruction_pipeline(
+        particles,
+        particle_radius=0.025,
+        rest_density=1000.0,
+        smoothing_length=2.0,
+        cube_size=1.0,
+        iso_surface_threshold=0.6,
+        mesh_smoothing_iters=5,
+        output_mesh_smoothing_weights=True,
+    )
+
+    assert type(mesh_with_data) is pysplashsurf.MeshWithData
+    assert type(reconstruction) is pysplashsurf.SurfaceReconstruction
+    assert type(mesh_with_data.mesh) is pysplashsurf.TriMesh3d
+
+    mesh = mesh_with_data.mesh
+
+    assert mesh_with_data.dtype == mesh.dtype
+    assert mesh_with_data.dtype == dtype
+
+    assert type(mesh_with_data.mesh_type) is pysplashsurf.MeshType
+    assert mesh_with_data.mesh_type == pysplashsurf.MeshType.Tri3d
+
+    assert mesh.vertices.dtype == dtype
+    assert mesh.triangles.dtype in [np.uint32, np.uint64]
+
+    assert mesh_with_data.nvertices == len(mesh.vertices)
+    assert mesh_with_data.ncells == len(mesh.triangles)
+
+    assert mesh_with_data.nvertices in range(21000, 25000)
+    assert mesh_with_data.ncells in range(45000, 49000)
+
+    assert mesh.vertices.shape == (mesh_with_data.nvertices, 3)
+    assert mesh.triangles.shape == (mesh_with_data.ncells, 3)
+
+    assert len(mesh_with_data.point_attributes) == 2
+    assert len(mesh_with_data.cell_attributes) == 0
+
+    assert "sw" in mesh_with_data.point_attributes
+    assert "wnn" in mesh_with_data.point_attributes
+
+    sw = mesh_with_data.point_attributes["sw"]
+    wnn = mesh_with_data.point_attributes["wnn"]
+
+    assert len(sw) == mesh_with_data.nvertices
+    assert len(wnn) == mesh_with_data.nvertices
+
+    assert sw.dtype == dtype
+    assert wnn.dtype == dtype
+
+    assert sw.shape == (mesh_with_data.nvertices,)
+    assert wnn.shape == (mesh_with_data.nvertices,)
+
+    assert sw.min() >= 0.0
+    assert sw.max() <= 1.0
+
+    assert wnn.min() >= 0.0
+
+
+def test_pipeline_f32():
+    impl_basic_test(np.float32)
+
+
+def test_pipeline_f64():
+    impl_basic_test(np.float64)
+
+
+def test_reconstruct():
+    particles = np.array(meshio.read(VTK_PATH).points, dtype=np.float32)
+
+    reconstruction = pysplashsurf.reconstruct_surface(
+        particles,
+        particle_radius=0.025,
+        rest_density=1000.0,
+        smoothing_length=2.0 * 0.025,
+        cube_size=1.0 * 0.025,
+        iso_surface_threshold=0.6,
+        global_neighborhood_list=True,
+    )
+
+    assert type(reconstruction) is pysplashsurf.SurfaceReconstruction
+    assert type(reconstruction.mesh) is pysplashsurf.TriMesh3d
+    assert type(reconstruction.grid) is pysplashsurf.UniformGrid
+    assert type(reconstruction.particle_densities) is np.ndarray
+    assert type(reconstruction.particle_inside_aabb) is type(None)
+    assert type(reconstruction.particle_neighbors) is pysplashsurf.NeighborhoodLists
+
+    mesh = reconstruction.mesh
+
+    assert mesh.dtype == np.float32
+
+    assert reconstruction.particle_densities.dtype == np.float32
+    assert len(reconstruction.particle_densities) == len(particles)
+
+    assert len(mesh.vertices) in range(25000, 30000)
+    assert len(mesh.triangles) in range(49000, 53000)
+
+
+def test_neighborhood_search():
+    particles = np.array(meshio.read(VTK_PATH).points, dtype=np.float32)
+
+    reconstruction = pysplashsurf.reconstruct_surface(
+        particles,
+        particle_radius=0.025,
+        rest_density=1000.0,
+        smoothing_length=2.0 * 0.025,
+        cube_size=1.0 * 0.025,
+        iso_surface_threshold=0.6,
+        global_neighborhood_list=True,
+    )
+
+    neighbors_reconstruct = reconstruction.particle_neighbors.get_neighborhood_lists()
+
+    assert type(neighbors_reconstruct) is list
+    assert len(neighbors_reconstruct) == len(particles)
+
+    aabb = reconstruction.grid.aabb
+
+    neighbor_lists = pysplashsurf.neighborhood_search_spatial_hashing_parallel(
+        particles, domain=aabb, search_radius=4.0 * 0.025
+    )
+
+    assert type(neighbor_lists) is pysplashsurf.NeighborhoodLists
+
+    neighbors = neighbor_lists.get_neighborhood_lists()
+
+    assert type(neighbors) is list
+    assert len(neighbors) == len(particles)
+    assert len(neighbors) == len(neighbors_reconstruct)
+
+    # TODO: Compare with naive neighbor search
+
+
+def test_check_consistency():
+    particles = np.array(meshio.read(VTK_PATH).points, dtype=np.float32)
+
+    reconstruction = pysplashsurf.reconstruct_surface(
+        particles,
+        particle_radius=0.025,
+        rest_density=1000.0,
+        smoothing_length=2.0 * 0.025,
+        cube_size=1.0 * 0.025,
+        iso_surface_threshold=0.6,
+        global_neighborhood_list=True,
+    )
+    mesh = reconstruction.mesh
+
+    assert pysplashsurf.check_mesh_consistency(mesh, reconstruction.grid) is None
+
+    mesh_with_data, reconstruction = pysplashsurf.reconstruction_pipeline(
+        particles,
+        particle_radius=0.025,
+        rest_density=1000.0,
+        smoothing_length=2.0,
+        cube_size=1.0,
+        iso_surface_threshold=0.6,
+        mesh_smoothing_iters=5,
+        output_mesh_smoothing_weights=True,
+    )
+
+    assert (
+        pysplashsurf.check_mesh_consistency(mesh_with_data, reconstruction.grid) is None
+    )
+
+
+def test_tris_to_quads():
+    particles = np.array(meshio.read(VTK_PATH).points, dtype=np.float32)
+
+    mesh_with_data, reconstruction = pysplashsurf.reconstruction_pipeline(
+        particles,
+        particle_radius=0.025,
+        rest_density=1000.0,
+        smoothing_length=2.0,
+        cube_size=1.0,
+        iso_surface_threshold=0.6,
+        mesh_smoothing_iters=5,
+        output_mesh_smoothing_weights=True,
+    )
+
+    mesh_with_data_quads = pysplashsurf.convert_tris_to_quads(mesh_with_data)
+
+    assert type(mesh_with_data_quads.mesh) is pysplashsurf.MixedTriQuadMesh3d
+    assert mesh_with_data_quads.mesh_type == pysplashsurf.MeshType.MixedTriQuad3d
+
+    assert mesh_with_data_quads.nvertices == mesh_with_data.nvertices
+    assert mesh_with_data_quads.ncells < mesh_with_data.ncells
+
+    tris = mesh_with_data_quads.mesh.get_triangles()
+    quads = mesh_with_data_quads.mesh.get_quads()
+
+    assert tris.dtype in [np.uint32, np.uint64]
+    assert quads.dtype in [np.uint32, np.uint64]
+
+    assert len(tris) + len(quads) == mesh_with_data_quads.ncells
+
+    assert tris.shape == (len(tris), 3)
+    assert quads.shape == (len(quads), 4)
+
+    assert len(tris) in range(35000, 39000)
+    assert len(quads) in range(4600, 5000)
+
+    assert len(mesh_with_data.point_attributes) == 2
+    assert len(mesh_with_data.cell_attributes) == 0
+
+    assert "sw" in mesh_with_data.point_attributes
+    assert "wnn" in mesh_with_data.point_attributes
+
+
+def test_interpolator():
+    particles = np.array(meshio.read(VTK_PATH).points, dtype=np.float32)
+
+    mesh_with_data, reconstruction = pysplashsurf.reconstruction_pipeline(
+        particles,
+        particle_radius=0.025,
+        rest_density=1000.0,
+        smoothing_length=2.0,
+        cube_size=1.0,
+        iso_surface_threshold=0.6,
+        mesh_smoothing_iters=5,
+        output_mesh_smoothing_weights=True,
+    )
+
+    compact_support = 4.0 * 0.025
+    rest_mass = 1000.0 * 0.025**3
+
+    interpolator = pysplashsurf.SphInterpolator(
+        particles, reconstruction.particle_densities, rest_mass, compact_support
+    )
+
+    assert type(interpolator) is pysplashsurf.SphInterpolator
+
+    mesh = mesh_with_data.mesh
+    mesh_densities = interpolator.interpolate_quantity(
+        reconstruction.particle_densities, mesh.vertices
+    )
+
+    assert type(mesh_densities) is np.ndarray
+    assert mesh_densities.dtype == np.float32
+    assert mesh_densities.shape == (len(mesh.vertices),)
+    assert mesh_densities.min() >= 0.0
+
+    mesh_particles = interpolator.interpolate_quantity(particles, mesh.vertices)
+
+    assert type(mesh_particles) is np.ndarray
+    assert mesh_particles.dtype == np.float32
+    assert mesh_particles.shape == (len(mesh.vertices), 3)
+
+    mesh_sph_normals = interpolator.interpolate_normals(mesh.vertices)
+
+    assert type(mesh_sph_normals) is np.ndarray
+    assert mesh_sph_normals.dtype == np.float32
+    assert mesh_sph_normals.shape == (len(mesh.vertices), 3)
+
+    mesh_with_data.add_point_attribute("density", mesh_densities)
+    mesh_with_data.add_point_attribute("position", mesh_particles)
+    mesh_with_data.add_point_attribute("normal", mesh_sph_normals)
+
+    assert "density" in mesh_with_data.point_attributes
+    assert "position" in mesh_with_data.point_attributes
+    assert "normal" in mesh_with_data.point_attributes
+
+    assert np.array_equal(mesh_with_data.point_attributes["density"], mesh_densities)
+    assert np.array_equal(mesh_with_data.point_attributes["position"], mesh_particles)
+    assert np.array_equal(mesh_with_data.point_attributes["normal"], mesh_sph_normals)