feat: npdu+kdtree as v0.2.0

Cargo.toml

@@ -12,3 +12,4 @@ crate-type = ["cdylib"]
 [dependencies]
 pyo3 = { version = "0.19", features = ["extension-module"] }
 numpy = "0.19"
+kdtree = "0.7"

README.md

@@ -4,12 +4,18 @@
 [![license badge](https://img.shields.io/github/license/leonardodalinky/fpsample)](https://github.com/leonardodalinky/fpsample/blob/main/LICENSE)
 [![star badge](https://img.shields.io/github/stars/leonardodalinky/fpsample?style=social)](https://github.com/leonardodalinky/fpsample)
 
-Python farthest point sampling library built on Rust pyo3 bindings. Compatible with numpy.
+Python efficient farthest point sampling (FPS) library, 100x faster than pure Python implementation.
+
+`fpsample` is coupled with `numpy` and built upon Rust pyo3 bindings.
+
+*Library for GPU is under construction*. Any issues are welcome.
 
 ## Installation
 
 ### Install from PyPI
 
+`numpy>=1.16.0` is required. Install `fpsample` using pip:
+
 ```shell
 pip install -U fpsample
 ```
@@ -32,26 +38,37 @@ import numpy as np
 # Generate random point cloud
 pc = np.random.rand(4096, 3)
 # sample 1024 points
-fps_samples = fpsample.fps_sampling(pc, 1024)
+fps_samples_idx = fpsample.fps_sampling(pc, 1024)
 
-fps_npdu_samples = fpsample.fps_npdu_sampling(pc, 1024)
+fps_npdu_samples_idx = fpsample.fps_npdu_sampling(pc, 1024)
 # or specify the windows size
-fps_npdu_samples = fpsample.fps_npdu_sampling(pc, 1024, k=64)
+fps_npdu_samples_idx = fpsample.fps_npdu_sampling(pc, 1024, k=64)
+
+fps_npdu_kdtree_samples_idx = fpsample.fps_npdu_kdtree_sampling(pc, 1024)
+# or specify the windows size
+fps_npdu_kdtree_samples_idx = fpsample.fps_npdu_kdtree_sampling(pc, 1024, k=64)
 ```
 
-> NOTE: NPDU method only gives sub-optimal answers. And it assumes that the points are approximately sorted or have **dimensional locality**. Otherwise, the result may be **worse**. Check the paper for details.
+**NOTE**: NPDU method only gives sub-optimal answers. And it assumes that the points are approximately sorted or have **dimensional locality**. Otherwise, the result may be **worse**. Check the paper for details.
+
+NPDU+KDTree method is more robust than NPDU method. It does not require the dimensional locality. But it is slightly slower than vanilla NPDU method. **It is recommended to use NPDU+KDTree method in general cases**.
 
 ## Performance
 Setup:
   - CPU: Intel(R) Core(TM) i9-10940X CPU @ 3.30GHz
   - RAM: 128 GiB
 
-|  Method  | #samples | #points |        Time       |
-|:--------:|:--------:|:-------:|:-----------------:|
-|    FPS   |   1024   |   4096  |  18.6 ms ± 1.3 ms |
-| FPS+NPDU |   1024   |   4096  | 5.97 ms ± 0.61 ms |
-|    FPS   |  50,000  | 100,000 |  27.4 s ± 551 ms  |
-| FPS+NPDU |  50,000  | 100,000 |  5.36 s ± 152 ms  |
+|     Method      | #samples | #points |        Time        |
+| :-------------: | :------: | :-----: | :----------------: |
+|       FPS       |   1024   |  4096   | 18.6 ms ± 0.17 ms  |
+|    FPS+NPDU     |   1024   |  4096   | 3.68 ms ± 0.10 ms  |
+| FPS+NPDU+KDTree |   1024   |  4096   | 13.10 ms ± 0.16 ms |
+|       FPS       |   4000   | 50,000  |  832 ms ± 9.01 ms  |
+|    FPS+NPDU     |   4000   | 50,000  |  143 ms ± 1.98 ms  |
+| FPS+NPDU+KDTree |   4000   | 50,000  |  294 ms ± 9.06 ms  |
+|       FPS       |  50,000  | 100,000 |  22.1 s ± 207 ms   |
+|    FPS+NPDU     |  50,000  | 100,000 |  3.35 s ± 66.8 ms  |
+| FPS+NPDU+KDTree |  50,000  | 100,000 |  4.08 s ± 60.4 ms  |
 
 ## Reference
 The nearest-point-distance-updating (NPDU) heuristic strategy is proposed in the following paper:

pyproject.toml

@@ -4,15 +4,15 @@ build-backend = "maturin"
 
 [project]
 name = "fpsample"
-version = "0.1.0"
+version = "0.2.0"
 authors = [{ name = "Leonard Lin", email = "leonard.keilin@gmail.com" }]
 keyword = [
     "sampling",
     "farthest point sampling",
     "furthest point sampling",
     "point cloud",
 ]
-description = "A efficient CPU implementation of farthest point sampling (FPS) for point clouds."
+description = "An efficient CPU implementation of farthest point sampling (FPS) for point clouds."
 readme = "README.md"
 license = { file = "LICENSE" }
 requires-python = ">=3.7"
@@ -24,8 +24,8 @@ classifiers = [
 dependencies = ["numpy>=1.16.0"]
 
 [project.urls]
-Homepage = "https://github.com/leonardodalinky/fpsample"
 Repository = "https://github.com/leonardodalinky/fpsample"
+Tracker = "https://github.com/leonardodalinky/fpsample/issues"
 
 [tool.maturin]
 python-source = "python"

python/fpsample/wrapper.py

@@ -1,8 +1,9 @@
+import warnings
 from typing import Optional
 
 import numpy as np
 
-from .fpsample import _fps_npdu_sampling, _fps_sampling
+from .fpsample import _fps_npdu_kdtree_sampling, _fps_npdu_sampling, _fps_sampling
 
 
 def fps_sampling(pc: np.ndarray, n_samples: int) -> np.ndarray:
@@ -40,6 +41,32 @@ def fps_npdu_sampling(pc: np.ndarray, n_samples: int, k: Optional[int] = None) -
     assert n_pts >= n_samples, "n_pts should be >= n_samples"
     pc = pc.astype(np.float32)
     k = k or int(n_pts / n_samples * 16)
+    if k >= n_pts - 1:
+        warnings.warn(f"k is too large, set to {n_pts - 1}")
+        k = n_pts - 1
     # Random pick a start
     start_idx = np.random.randint(low=0, high=n_pts)
     return _fps_npdu_sampling(pc, n_samples, k, start_idx)
+
+
+def fps_npdu_kdtree_sampling(pc: np.ndarray, n_samples: int, k: Optional[int] = None) -> np.ndarray:
+    """
+    Args:
+        pc (np.ndarray): The input point cloud of shape (n_pts, D).
+        n_samples (int): Number of samples.
+        k (int, default=None): Windows size of local heuristic search. If set to None, it will be set to `n_pts / n_samples * 16`.
+    Returns:
+        np.ndarray: The selected indices of shape (n_samples,).
+    """
+    assert n_samples >= 1, "n_samples should be >= 1"
+    assert pc.ndim == 2
+    n_pts, _ = pc.shape
+    assert n_pts >= n_samples, "n_pts should be >= n_samples"
+    pc = pc.astype(np.float32)
+    k = k or int(n_pts / n_samples * 16)
+    if k >= n_pts:
+        warnings.warn(f"k is too large, set to {n_pts}")
+        k = n_pts
+    # Random pick a start
+    start_idx = np.random.randint(low=0, high=n_pts)
+    return _fps_npdu_kdtree_sampling(pc, n_samples, k, start_idx)

src/lib.rs

@@ -1,3 +1,5 @@
+use kdtree::distance::squared_euclidean;
+use kdtree::KdTree;
 use numpy::{
     ndarray::s,
     ndarray::{Array1, ArrayView2, Axis, Zip},
@@ -174,10 +176,92 @@ fn fps_npdu_sampling_py<'py>(
     Ok(ret)
 }
 
+fn fps_npdu_kdtree_sampling(
+    points: ArrayView2<f32>,
+    n_samples: usize,
+    k: usize,
+    start_idx: usize,
+) -> Array1<usize> {
+    let [p, c] = points.shape() else {
+        panic!("points must be a 2D array")
+    };
+    // contruct kd-tree
+    let std_points = points.as_standard_layout();
+    let mut kdtree: KdTree<f32, usize, _> = KdTree::new(*c);
+    let std_points_vec = std_points.outer_iter().collect::<Vec<_>>();
+    let std_points_slice_vec = std_points_vec
+        .iter()
+        .map(|x| x.as_slice().unwrap())
+        .collect::<Vec<_>>();
+    for (i, point) in std_points_slice_vec.iter().enumerate() {
+        kdtree.add(*point, i).unwrap();
+    }
+
+    // previous round selected point index
+    let mut res_selected_idx: Option<usize> = None;
+    // distance from each point to the selected point set
+    let mut dist_pts_to_selected_min = Array1::<f32>::from_elem((*p,), f32::INFINITY);
+    // selected points index
+    let mut selected_pts_idx = Vec::<usize>::with_capacity(n_samples);
+
+    while selected_pts_idx.len() < n_samples {
+        if let Some(prev_idx) = res_selected_idx {
+            // find nearest point
+            let nearest_idx = kdtree
+                .nearest(std_points_slice_vec[prev_idx], k, &squared_euclidean)
+                .unwrap();
+            for (dist, idx) in nearest_idx {
+                // dist_pts_to_selected_min[idx] = f32::min(dist_pts_to_selected_min[idx], dist);
+                let d = dist_pts_to_selected_min.get_mut(*idx).unwrap();
+                *d = f32::min(*d, dist);
+            }
+
+            let max_idx = dist_pts_to_selected_min
+                .indexed_iter()
+                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
+                .unwrap()
+                .0;
+            selected_pts_idx.push(max_idx);
+            res_selected_idx = Some(max_idx);
+        } else {
+            // update distance at the first round
+            let dist = &points - &points.slice(s![start_idx, ..]);
+            let dist = dist.mapv(|x| x.powi(2)).sum_axis(Axis(1));
+            // update min distance
+            Zip::from(&mut dist_pts_to_selected_min)
+                .and(&dist)
+                .for_each(|x, &y| {
+                    *x = f32::min(*x, y);
+                });
+            // first point
+            selected_pts_idx.push(start_idx);
+            res_selected_idx = Some(start_idx);
+        }
+    }
+    selected_pts_idx.into()
+}
+
+#[pyfunction]
+#[pyo3(name = "_fps_npdu_kdtree_sampling")]
+fn fps_npdu_kdtree_sampling_py<'py>(
+    py: Python<'py>,
+    points: PyReadonlyArray2<f32>,
+    n_samples: usize,
+    k: usize,
+    start_idx: usize,
+) -> PyResult<&'py PyArray1<usize>> {
+    check_py_input(&points, n_samples, start_idx)?;
+    let points = points.as_array();
+    let idxs = py.allow_threads(|| fps_npdu_kdtree_sampling(points, n_samples, k, start_idx));
+    let ret = idxs.to_pyarray(py);
+    Ok(ret)
+}
+
 #[pymodule]
 fn fpsample(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(fps_sampling_py, m)?)?;
     m.add_function(wrap_pyfunction!(fps_npdu_sampling_py, m)?)?;
+    m.add_function(wrap_pyfunction!(fps_npdu_kdtree_sampling_py, m)?)?;
 
     Ok(())
 }