Implemented kernel switching in pysplashsurf

Author: Fek04

pysplashsurf/pysplashsurf/pysplashsurf.pyi

@@ -185,11 +185,11 @@ class NeighborhoodLists:
 
 class SphInterpolator:
     r"""
-    Interpolator of per-particle quantities to arbitrary points using SPH interpolation (with cubic kernel)
+    Interpolator of per-particle quantities to arbitrary points using SPH interpolation
     """
-    def __new__(cls, particle_positions:numpy.typing.NDArray[typing.Any], particle_densities:numpy.typing.NDArray[typing.Any], particle_rest_mass:builtins.float, compact_support_radius:builtins.float) -> SphInterpolator:
+    def __new__(cls, particle_positions:numpy.typing.NDArray[typing.Any], particle_densities:numpy.typing.NDArray[typing.Any], particle_rest_mass:builtins.float, compact_support_radius:builtins.float, kernel_type:KernelType) -> SphInterpolator:
         r"""
-        Constructs an SPH interpolator (with cubic kernels) for the given particles
+        Constructs an SPH interpolator for the given particles
         
         Parameters
         ----------
@@ -200,7 +200,9 @@ class SphInterpolator:
         particle_rest_mass
             The rest mass of each particle (assumed to be the same for all particles).
         compact_support_radius
-            The compact support radius of the cubic spline kernel used for interpolation.
+            The compact support radius of the kernel used for interpolation.
+        kernel_type
+            The kernel function used for interpolation
         """
     def interpolate_quantity(self, particle_quantity:numpy.typing.NDArray[typing.Any], interpolation_points:numpy.typing.NDArray[typing.Any], *, first_order_correction:builtins.bool=False) -> numpy.typing.NDArray[typing.Any]:
         r"""
@@ -315,6 +317,15 @@ class VertexVertexConnectivity:
         Returns the wrapped connectivity data by moving it out of this object (zero copy)
         """
 
+class KernelType(Enum):
+    r"""
+    Enum for specifying the Kernel function used for the reconstruction
+    """
+    CubicSpline = ...
+    Poly6 = ...
+    Spiky = ...
+    WendlandQuinticC2 = ...
+
 class MeshType(Enum):
     r"""
     Enum specifying the type of mesh wrapped by a ``MeshWithData``
@@ -497,7 +508,7 @@ def neighborhood_search_spatial_hashing_parallel(particle_positions:numpy.typing
         The radius per particle where other particles are considered neighbors.
     """
 
-def reconstruct_surface(particles:numpy.typing.NDArray[typing.Any], *, particle_radius:builtins.float, rest_density:builtins.float=1000.0, smoothing_length:builtins.float, cube_size:builtins.float, iso_surface_threshold:builtins.float=0.6, aabb_min:typing.Optional[typing.Sequence[builtins.float]]=None, aabb_max:typing.Optional[typing.Sequence[builtins.float]]=None, multi_threading:builtins.bool=True, global_neighborhood_list:builtins.bool=False, subdomain_grid:builtins.bool=True, subdomain_grid_auto_disable:builtins.bool=True, subdomain_num_cubes_per_dim:builtins.int=64) -> SurfaceReconstruction:
+def reconstruct_surface(particles:numpy.typing.NDArray[typing.Any], *, particle_radius:builtins.float, rest_density:builtins.float=1000.0, smoothing_length:builtins.float, cube_size:builtins.float, iso_surface_threshold:builtins.float=0.6, aabb_min:typing.Optional[typing.Sequence[builtins.float]]=None, aabb_max:typing.Optional[typing.Sequence[builtins.float]]=None, multi_threading:builtins.bool=True, simd:builtins.bool=True, kernel_type:KernelType=..., global_neighborhood_list:builtins.bool=False, subdomain_grid:builtins.bool=True, subdomain_grid_auto_disable:builtins.bool=True, subdomain_num_cubes_per_dim:builtins.int=64) -> SurfaceReconstruction:
     r"""
     Performs a surface reconstruction from the given particles without additional post-processing
     
@@ -523,6 +534,8 @@ def reconstruct_surface(particles:numpy.typing.NDArray[typing.Any], *, particle_
         Upper corner of the AABB of particles to consider in the reconstruction.
     multi_threading
         Flag to enable multi-threading for the reconstruction and post-processing steps.
+    simd
+       Flag to enable SIMD vectorization for the reconstruction if supported by the CPU architecture.
     subdomain_grid
         Flag to enable spatial decomposition by dividing the domain into subdomains with dense marching cube grids for efficient multi-threading.
     subdomain_grid_auto_disable
@@ -531,7 +544,7 @@ def reconstruct_surface(particles:numpy.typing.NDArray[typing.Any], *, particle_
         Number of marching cubes voxels along each coordinate axis in each subdomain if the subdomain grid is enabled.
     """
 
-def reconstruction_pipeline(particles:numpy.typing.NDArray[typing.Any], *, attributes_to_interpolate:typing.Optional[dict]=None, particle_radius:builtins.float, rest_density:builtins.float=1000.0, smoothing_length:builtins.float, cube_size:builtins.float, iso_surface_threshold:builtins.float=0.6, aabb_min:typing.Optional[typing.Sequence[builtins.float]]=None, aabb_max:typing.Optional[typing.Sequence[builtins.float]]=None, multi_threading:builtins.bool=True, subdomain_grid:builtins.bool=True, subdomain_grid_auto_disable:builtins.bool=True, subdomain_num_cubes_per_dim:builtins.int=64, check_mesh_closed:builtins.bool=False, check_mesh_manifold:builtins.bool=False, check_mesh_orientation:builtins.bool=False, check_mesh_debug:builtins.bool=False, mesh_cleanup:builtins.bool=False, mesh_cleanup_snap_dist:typing.Optional[builtins.float]=None, decimate_barnacles:builtins.bool=False, keep_vertices:builtins.bool=False, compute_normals:builtins.bool=False, sph_normals:builtins.bool=False, normals_smoothing_iters:typing.Optional[builtins.int]=None, mesh_smoothing_iters:typing.Optional[builtins.int]=None, mesh_smoothing_weights:builtins.bool=True, mesh_smoothing_weights_normalization:builtins.float=13.0, generate_quads:builtins.bool=False, quad_max_edge_diag_ratio:builtins.float=1.75, quad_max_normal_angle:builtins.float=10.0, quad_max_interior_angle:builtins.float=135.0, output_mesh_smoothing_weights:builtins.bool=False, output_raw_normals:builtins.bool=False, output_raw_mesh:builtins.bool=False, mesh_aabb_min:typing.Optional[typing.Sequence[builtins.float]]=None, mesh_aabb_max:typing.Optional[typing.Sequence[builtins.float]]=None, mesh_aabb_clamp_vertices:builtins.bool=True) -> tuple[MeshWithData, SurfaceReconstruction]:
+def reconstruction_pipeline(particles:numpy.typing.NDArray[typing.Any], *, attributes_to_interpolate:typing.Optional[dict]=None, particle_radius:builtins.float, rest_density:builtins.float=1000.0, smoothing_length:builtins.float, cube_size:builtins.float, iso_surface_threshold:builtins.float=0.6, aabb_min:typing.Optional[typing.Sequence[builtins.float]]=None, aabb_max:typing.Optional[typing.Sequence[builtins.float]]=None, multi_threading:builtins.bool=True, simd:builtins.bool=True, kernel_type:KernelType=..., subdomain_grid:builtins.bool=True, subdomain_grid_auto_disable:builtins.bool=True, subdomain_num_cubes_per_dim:builtins.int=64, check_mesh_closed:builtins.bool=False, check_mesh_manifold:builtins.bool=False, check_mesh_orientation:builtins.bool=False, check_mesh_debug:builtins.bool=False, mesh_cleanup:builtins.bool=False, mesh_cleanup_snap_dist:typing.Optional[builtins.float]=None, decimate_barnacles:builtins.bool=False, keep_vertices:builtins.bool=False, compute_normals:builtins.bool=False, sph_normals:builtins.bool=False, normals_smoothing_iters:typing.Optional[builtins.int]=None, mesh_smoothing_iters:typing.Optional[builtins.int]=None, mesh_smoothing_weights:builtins.bool=True, mesh_smoothing_weights_normalization:builtins.float=13.0, generate_quads:builtins.bool=False, quad_max_edge_diag_ratio:builtins.float=1.75, quad_max_normal_angle:builtins.float=10.0, quad_max_interior_angle:builtins.float=135.0, output_mesh_smoothing_weights:builtins.bool=False, output_raw_normals:builtins.bool=False, output_raw_mesh:builtins.bool=False, mesh_aabb_min:typing.Optional[typing.Sequence[builtins.float]]=None, mesh_aabb_max:typing.Optional[typing.Sequence[builtins.float]]=None, mesh_aabb_clamp_vertices:builtins.bool=True) -> tuple[MeshWithData, SurfaceReconstruction]:
     r"""
     Runs the surface reconstruction pipeline for the given particle positions with optional post-processing
     
@@ -561,6 +574,10 @@ def reconstruction_pipeline(particles:numpy.typing.NDArray[typing.Any], *, attri
         Upper corner [x,y,z] of the AABB of particles to consider in the reconstruction.
     multi_threading
         Flag to enable multi-threading for the reconstruction and post-processing steps.
+    simd
+        Flag to enable SIMD vectorization for the reconstruction if supported by the CPU architecture.
+    kernel_type
+        Kernel function to use for the reconstruction.
     subdomain_grid
         Flag to enable spatial decomposition by dividing the domain into subdomains with dense marching cube grids for efficient multi-threading.
     subdomain_grid_auto_disable

pysplashsurf/src/lib.rs

@@ -39,6 +39,7 @@ fn pysplashsurf(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<uniform_grid::PyUniformGrid>()?;
     m.add_class::<reconstruction::PySurfaceReconstruction>()?;
     m.add_class::<sph_interpolation::PySphInterpolator>()?;
+    m.add_class::<utils::KernelType>()?;
 
     use wrap_pyfunction as wrap;
 

pysplashsurf/src/pipeline.rs

@@ -18,7 +18,7 @@ use std::borrow::Cow;
 
 use crate::mesh::PyMeshWithData;
 use crate::reconstruction::PySurfaceReconstruction;
-use crate::utils::{IndexT, pyerr_unsupported_scalar};
+use crate::utils::{IndexT, KernelType, pyerr_unsupported_scalar};
 
 /// Runs the surface reconstruction pipeline for the given particle positions with optional post-processing
 ///
@@ -49,7 +49,9 @@ use crate::utils::{IndexT, pyerr_unsupported_scalar};
 /// multi_threading
 ///     Flag to enable multi-threading for the reconstruction and post-processing steps.
 /// simd
-///    Flag to enable SIMD vectorization for the reconstruction if supported by the CPU architecture.
+///     Flag to enable SIMD vectorization for the reconstruction if supported by the CPU architecture.
+/// kernel_type
+///     Kernel function to use for the reconstruction.
 /// subdomain_grid
 ///     Flag to enable spatial decomposition by dividing the domain into subdomains with dense marching cube grids for efficient multi-threading.
 /// subdomain_grid_auto_disable
@@ -110,7 +112,7 @@ use crate::utils::{IndexT, pyerr_unsupported_scalar};
 #[pyo3(name = "reconstruction_pipeline")]
 #[pyo3(signature = (particles, *, attributes_to_interpolate = None,
     particle_radius, rest_density = 1000.0, smoothing_length, cube_size, iso_surface_threshold = 0.6,
-    aabb_min = None, aabb_max = None, multi_threading = true, simd = true,
+    aabb_min = None, aabb_max = None, multi_threading = true, simd = true, kernel_type = KernelType::CubicSpline,
     subdomain_grid = true, subdomain_grid_auto_disable = true, subdomain_num_cubes_per_dim = 64,
     check_mesh_closed = false, check_mesh_manifold = false, check_mesh_orientation = false, check_mesh_debug = false,
     mesh_cleanup = false, mesh_cleanup_snap_dist = None, decimate_barnacles = false, keep_vertices = false, compute_normals = false, sph_normals = false,
@@ -131,6 +133,7 @@ pub fn reconstruction_pipeline<'py>(
     aabb_max: Option<[f64; 3]>,
     multi_threading: bool,
     simd: bool,
+    kernel_type: KernelType,
     subdomain_grid: bool,
     subdomain_grid_auto_disable: bool,
     subdomain_num_cubes_per_dim: u32,
@@ -198,7 +201,7 @@ pub fn reconstruction_pipeline<'py>(
         enable_simd: simd,
         spatial_decomposition,
         global_neighborhood_list: false,
-        kernel_type: splashsurf_lib::kernel::KernelType::CubicSpline,
+        kernel_type: kernel_type.into_lib_enum(),
     };
 
     let postprocessing_args = splashsurf::reconstruct::ReconstructionPostprocessingParameters {

pysplashsurf/src/reconstruction.rs

@@ -1,7 +1,7 @@
 use crate::mesh::PyTriMesh3d;
 use crate::neighborhood_search::PyNeighborhoodLists;
 use crate::uniform_grid::PyUniformGrid;
-use crate::utils;
+use crate::utils::{self, KernelType};
 use anyhow::anyhow;
 use ndarray::ArrayView1;
 use numpy as np;
@@ -124,6 +124,8 @@ impl PySurfaceReconstruction {
 ///     Flag to enable multi-threading for the reconstruction and post-processing steps.
 /// simd
 ///    Flag to enable SIMD vectorization for the reconstruction if supported by the CPU architecture.
+/// kernel_type
+///     Kernel function to use for the reconstruction.
 /// subdomain_grid
 ///     Flag to enable spatial decomposition by dividing the domain into subdomains with dense marching cube grids for efficient multi-threading.
 /// subdomain_grid_auto_disable
@@ -135,8 +137,8 @@ impl PySurfaceReconstruction {
 #[pyo3(name = "reconstruct_surface")]
 #[pyo3(signature = (particles, *,
     particle_radius, rest_density = 1000.0, smoothing_length, cube_size, iso_surface_threshold = 0.6,
-    aabb_min = None, aabb_max = None,
-    multi_threading = true, simd = true, global_neighborhood_list = false,
+    aabb_min = None, aabb_max = None, multi_threading = true, simd = true,
+    kernel_type = KernelType::CubicSpline, global_neighborhood_list = false,
     subdomain_grid = true, subdomain_grid_auto_disable = true, subdomain_num_cubes_per_dim = 64
 ))]
 pub fn reconstruct_surface<'py>(
@@ -150,6 +152,7 @@ pub fn reconstruct_surface<'py>(
     aabb_max: Option<[f64; 3]>,
     multi_threading: bool,
     simd: bool,
+    kernel_type: KernelType,
     global_neighborhood_list: bool,
     subdomain_grid: bool,
     subdomain_grid_auto_disable: bool,
@@ -181,7 +184,7 @@ pub fn reconstruct_surface<'py>(
         enable_simd: simd,
         spatial_decomposition,
         global_neighborhood_list,
-        kernel_type: splashsurf_lib::kernel::KernelType::CubicSpline,
+        kernel_type: kernel_type.into_lib_enum(),
     };
 
     let element_type = particles.dtype();

pysplashsurf/src/sph_interpolation.rs

@@ -1,12 +1,11 @@
-use crate::utils::*;
+use crate::utils::{KernelType, *};
 use numpy as np;
 use numpy::prelude::*;
 use numpy::{Element, PyArray1, PyArray2, PyUntypedArray};
 use pyo3::PyResult;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3_stub_gen::derive::*;
-use splashsurf_lib::kernel::KernelType;
 use splashsurf_lib::nalgebra::SVector;
 use splashsurf_lib::{
     Real,
@@ -19,7 +18,7 @@ enum PySphInterpolatorWrapper {
     F64(SphInterpolator<f64>),
 }
 
-/// Interpolator of per-particle quantities to arbitrary points using SPH interpolation (with cubic kernel)
+/// Interpolator of per-particle quantities to arbitrary points using SPH interpolation
 #[gen_stub_pyclass]
 #[pyclass]
 #[pyo3(name = "SphInterpolator")]
@@ -36,6 +35,7 @@ impl PySphInterpolator {
         particle_densities: &Bound<'py, PyUntypedArray>,
         particle_rest_mass: f64,
         compact_support_radius: f64,
+        kernel_type: KernelType,
     ) -> PyResult<PySphInterpolator>
     where
         PySphInterpolator: From<SphInterpolator<R>>,
@@ -55,7 +55,7 @@ impl PySphInterpolator {
                 densities,
                 R::from_float(particle_rest_mass),
                 R::from_float(compact_support_radius),
-                KernelType::CubicSpline,
+                kernel_type.into_lib_enum(),
             )))
         } else {
             Err(pyerr_scalar_type_mismatch())
@@ -160,7 +160,7 @@ impl PySphInterpolator {
 #[gen_stub_pymethods]
 #[pymethods]
 impl PySphInterpolator {
-    /// Constructs an SPH interpolator (with cubic kernels) for the given particles
+    /// Constructs an SPH interpolator for the given particles
     ///
     /// Parameters
     /// ----------
@@ -171,13 +171,16 @@ impl PySphInterpolator {
     /// particle_rest_mass
     ///     The rest mass of each particle (assumed to be the same for all particles).
     /// compact_support_radius
-    ///     The compact support radius of the cubic spline kernel used for interpolation.
+    ///     The compact support radius of the kernel used for interpolation.
+    /// kernel_type
+    ///     The kernel function used for interpolation
     #[new]
     fn py_new<'py>(
         particle_positions: &Bound<'py, PyUntypedArray>,
         particle_densities: &Bound<'py, PyUntypedArray>,
         particle_rest_mass: f64,
         compact_support_radius: f64,
+        kernel_type: KernelType,
     ) -> PyResult<Self> {
         let py = particle_positions.py();
         let element_type = particle_positions.dtype();
@@ -188,13 +191,15 @@ impl PySphInterpolator {
                 particle_densities,
                 particle_rest_mass,
                 compact_support_radius,
+                kernel_type,
             )
         } else if element_type.is_equiv_to(&np::dtype::<f64>(py)) {
             Self::new_generic::<f64>(
                 particle_positions,
                 particle_densities,
                 particle_rest_mass,
                 compact_support_radius,
+                kernel_type,
             )
         } else {
             Err(pyerr_unsupported_scalar())

pysplashsurf/src/utils.rs

@@ -3,9 +3,32 @@ use numpy::{Element, PyArray, PyUntypedArray};
 use pyo3::exceptions::PyTypeError;
 use pyo3::prelude::*;
 use pyo3::{Bound, PyAny, PyErr, PyResult};
+use pyo3_stub_gen::derive::gen_stub_pyclass_enum;
 use splashsurf_lib::Real;
 use splashsurf_lib::nalgebra::SVector;
 
+/// Enum for specifying the Kernel function used for the reconstruction
+#[gen_stub_pyclass_enum]
+#[pyclass]
+#[derive(Clone)]
+pub enum KernelType {
+    CubicSpline,
+    Poly6,
+    Spiky,
+    WendlandQuinticC2,
+}
+
+impl KernelType {
+    pub fn into_lib_enum(&self) -> splashsurf_lib::kernel::KernelType {
+        match self {
+            KernelType::CubicSpline => splashsurf_lib::kernel::KernelType::CubicSpline,
+            KernelType::Poly6 => splashsurf_lib::kernel::KernelType::Poly6,
+            KernelType::Spiky => splashsurf_lib::kernel::KernelType::Spiky,
+            KernelType::WendlandQuinticC2 => splashsurf_lib::kernel::KernelType::WendlandQuinticC2,
+        }
+    }
+}
+
 /// The index type used for all grids and reconstructions in this crate
 pub(crate) type IndexT = i64;
 

pysplashsurf/tests/test_basic.py

@@ -301,7 +301,7 @@ def interpolator_test(dtype):
     rest_mass = 1000.0 * 0.025**3
 
     interpolator = pysplashsurf.SphInterpolator(
-        particles, reconstruction.particle_densities, rest_mass, compact_support
+        particles, reconstruction.particle_densities, rest_mass, compact_support, pysplashsurf.KernelType.CubicSpline
     )
 
     assert type(interpolator) is pysplashsurf.SphInterpolator