Skip to content

Commit

Permalink
Update to nanoflann 1.6.1
Browse files Browse the repository at this point in the history
- Add threading in kd three construction. It should improve runtime most
of the time.
  • Loading branch information
rjanvier committed Sep 14, 2024
1 parent 0d7b74f commit 564cfcc
Show file tree
Hide file tree
Showing 4 changed files with 9 additions and 10 deletions.
4 changes: 2 additions & 2 deletions include/nn_search.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ static std::pair<nb::ndarray<nb::numpy, uint32_t, nb::ndim<2>>, nb::ndarray<nb::

if (knn > data.rows()) { throw std::invalid_argument("knn size is greater than the data point cloud size"); }

kd_tree_t kd_tree(3, data, 10);
kd_tree_t kd_tree(3, data, 10, 0);
const Eigen::Index n_points = query.rows();
uint32_t* indices = new uint32_t[knn * n_points];
nb::capsule owner_indices(indices, [](void* p) noexcept { delete[] (uint32_t*)p; });
Expand Down Expand Up @@ -94,7 +94,7 @@ static std::pair<nb::ndarray<nb::numpy, int32_t, nb::ndim<2>>, nb::ndarray<nb::n
throw std::invalid_argument("max knn size is greater than the data point cloud size");
}

kd_tree_t kd_tree(3, data, 10);
kd_tree_t kd_tree(3, data, 10, 0);
const real_t sq_search_radius = search_radius * search_radius;

const Eigen::Index n_points = query.rows();
Expand Down
10 changes: 5 additions & 5 deletions include/pgeof.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>


#include <algorithm>
#include <cmath>
#include <cstdio>
Expand Down Expand Up @@ -74,9 +73,10 @@ static inline void flush() { std::cout << std::endl; };
* @return the geometric features associated with each point's neighborhood in a (num_points, features_count) nd::array.
*/
template <typename real_t = float, const size_t feature_count = 11>
static nb::ndarray<nb::numpy, real_t, nb::shape<-1, static_cast<nb::ssize_t>(feature_count)>> compute_geometric_features(
RefCloud<real_t> xyz, nb::ndarray<const uint32_t, nb::ndim<1>> nn, nb::ndarray<const uint32_t, nb::ndim<1>> nn_ptr,
const size_t k_min, const bool verbose)
static nb::ndarray<nb::numpy, real_t, nb::shape<-1, static_cast<nb::ssize_t>(feature_count)>>
compute_geometric_features(
RefCloud<real_t> xyz, nb::ndarray<const uint32_t, nb::ndim<1>> nn,
nb::ndarray<const uint32_t, nb::ndim<1>> nn_ptr, const size_t k_min, const bool verbose)
{
if (k_min < 1) { throw std::invalid_argument("k_min should be > 1"); }
// Each point can be treated in parallel
Expand Down Expand Up @@ -330,7 +330,7 @@ static nb::ndarray<nb::numpy, real_t, nb::shape<-1, -1>> compute_geometric_featu
using kd_tree_t = nanoflann::KDTreeEigenMatrixAdaptor<RefCloud<real_t>, 3, nanoflann::metric_L2_Simple>;
// TODO: where knn < num of points

kd_tree_t kd_tree(3, xyz, 10);
kd_tree_t kd_tree(3, xyz, 10, 0);
const size_t feature_count = selected_features.size();
const Eigen::Index n_points = xyz.rows();
real_t sq_search_radius = search_radius * search_radius;
Expand Down
3 changes: 1 addition & 2 deletions src/pgeof_ext.cpp
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@

#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>
#include <nanobind/stl/pair.h>
#include <nanobind/stl/vector.h>

#include "nn_search.hpp"
Expand Down Expand Up @@ -176,4 +175,4 @@ NB_MODULE(pgeof_ext, m)
:param selected_features: List of selected features. See EFeatureID
:return: Geometric features associated with each point's neighborhood in a (num_points, features_count) numpy array.
)");
}
}
2 changes: 1 addition & 1 deletion third_party/nanoflann

0 comments on commit 564cfcc

Please sign in to comment.