You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
52 lines
1.9 KiB
52 lines
1.9 KiB
#ifndef MEDUSA_BITS_SPATIAL_SEARCH_KDTREE_HPP_
|
|
#define MEDUSA_BITS_SPATIAL_SEARCH_KDTREE_HPP_
|
|
|
|
/**
|
|
* @file
|
|
* Implementation of KDTree.
|
|
*/
|
|
|
|
#include "KDTree_fwd.hpp"
|
|
#include <medusa/bits/utils/assert.hpp>
|
|
#include <medusa/bits/utils/memutils.hpp>
|
|
|
|
namespace mm {
|
|
|
|
template <class vec_t>
|
|
std::pair<Range<int>, Range<typename vec_t::scalar_t>> KDTree<vec_t>::query(const vec_t& point,
|
|
int k) const {
|
|
assert_msg(point.array().isFinite().prod() == 1, "Invalid point.");
|
|
Range<int> ret_index(k);
|
|
Range<scalar_t> out_dist_sqr(k);
|
|
int actual_k = tree.knnSearch(point.data(), k, &ret_index[0], &out_dist_sqr[0]);
|
|
assert_msg(actual_k == k, "There were not enough points in the tree, you requested %d "
|
|
"points, the tree only contains %d points.", k, actual_k);
|
|
return {ret_index, out_dist_sqr};
|
|
}
|
|
|
|
template <class vec_t>
|
|
std::pair<Range<int>, Range<typename vec_t::scalar_t>> KDTree<vec_t>::query(
|
|
const vec_t& point, const scalar_t& radius_squared) const {
|
|
assert_msg(point.array().isFinite().prod() == 1, "Invalid point.");
|
|
std::vector<std::pair<int, scalar_t>> idx_dist;
|
|
int k = tree.radiusSearch(point.data(), radius_squared, idx_dist, nanoflann::SearchParams());
|
|
Range<int> idx(k); Range<scalar_t> dists(k);
|
|
for (int i = 0; i < k; ++i) {
|
|
std::tie(idx[i], dists[i]) = idx_dist[i];
|
|
}
|
|
return {idx, dists};
|
|
}
|
|
|
|
/// Output basic info about given tree.
|
|
template <class V>
|
|
std::ostream& operator<<(std::ostream& os, const KDTree<V>& tree) {
|
|
return os << "KDTree:\n"
|
|
<< " dimension: " << tree.dim << '\n'
|
|
<< " num of points: " << tree.size() << "\n"
|
|
<< " first point: " << tree.get(0) << "\n"
|
|
<< " last point: " << tree.get(tree.size() - 1) << "\n";
|
|
}
|
|
|
|
} // namespace mm
|
|
|
|
#endif // MEDUSA_BITS_SPATIAL_SEARCH_KDTREE_HPP_
|
|
|