KDTreeMutable.hpp

Go to the documentation of this file.

#ifndef MEDUSA_BITS_SPATIAL_SEARCH_KDTREEMUTABLE_HPP_
#define MEDUSA_BITS_SPATIAL_SEARCH_KDTREEMUTABLE_HPP_
 
#include "KDTreeMutable_fwd.hpp"
#include <medusa/bits/utils/assert.hpp>
#include <cmath>
 
namespace mm {
 
template <class V>
std::ostream& operator<<(std::ostream& os, const KDTreeMutable<V>& tree)  {
    return os << "KDTreeMutable:\n"
              << "    dimension: " << tree.dim << '\n'
              << "    num of points: " << tree.size() << std::endl;
}
 
template <class vec_t>
void KDTreeMutable<vec_t>::insert(const vec_t& point) {
    assert_msg(point.array().isFinite().prod() == 1, "Invalid point.");
    auto n = points_.kdtree_get_point_count();
    points_.add(point);
    tree.addPoints(n, n);
    ++size_;
}
 
template <class vec_t>
void KDTreeMutable<vec_t>::insert(const Range<vec_t>& points) {
    auto n = points_.kdtree_get_point_count();
    for (const auto& p : points) {
        assert_msg(p.array().isFinite().prod() == 1, "One of the points is invalid.");
        points_.add(p);
    }
    size_ += points.size();
    tree.addPoints(n, n + points.size() - 1);
}
 
template <class vec_t>
std::pair<mm::Range<int>, mm::Range<double>>
KDTreeMutable<vec_t>::query(const vec_t& point, int k) {
    assert_msg(point.array().isFinite().prod() == 1, "Invalid query point.");
    nanoflann::KNNResultSet<scalar_t, int> resultSet(k);
    Range<int> ret_index(k);
    Range<scalar_t> out_dist_sqr(k);
    resultSet.init(&ret_index[0], &out_dist_sqr[0]);
    tree.findNeighbors(resultSet, point.data(), nanoflann::SearchParams(k));
    assert_msg(resultSet.full(), "Not enough points in the tree, you requested %d points, "
                                 "but the tree contains only %d points.", k, size());
    return {ret_index, out_dist_sqr};
}
 
}  // namespace mm
 
#endif  // MEDUSA_BITS_SPATIAL_SEARCH_KDTREEMUTABLE_HPP_