kdtree.h

// ---------------------------------------------------------------------
//
// Copyright (C) 2017 - 2018 by the deal.II authors
//
// This file is part of the deal.II library.
//
// The deal.II library is free software; you can use it, redistribute
// it, and/or modify it under the terms of the GNU Lesser General
// Public License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
// The full text of the license can be found in the file LICENSE.md at
// the top level directory of deal.II.
//
// ---------------------------------------------------------------------

#ifndef dealii_numerics_kdtree_h
#define dealii_numerics_kdtree_h

#include <deal.II/base/config.h>

#ifdef DEAL_II_WITH_NANOFLANN

#  include <deal.II/base/point.h>

#  include <nanoflann.hpp>

#  include <memory>


DEAL_II_NAMESPACE_OPEN

template <int dim>
class KDTree
{
public:
  KDTree(const unsigned int             max_leaf_size = 10,
         const std::vector<Point<dim>> &pts = std::vector<Point<dim>>());


  struct PointCloudAdaptor
  {
    using coord_t = double;


    const std::vector<Point<dim>> &points;


    PointCloudAdaptor(const std::vector<Point<dim>> &_points);


    size_t
    kdtree_get_point_count() const;


    coord_t
    kdtree_distance(const coord_t *p1,
                    const size_t   idx_p2,
                    const size_t   size) const;


    coord_t
    kdtree_get_pt(const size_t idx, const int d) const;


    template <class BBOX>
    bool
    kdtree_get_bbox(BBOX &) const;
  };


  using NanoFlannKDTree = typename nanoflann::KDTreeSingleIndexAdaptor<
    nanoflann::L2_Simple_Adaptor<double, PointCloudAdaptor>,
    PointCloudAdaptor,
    dim,
    unsigned int>;


  void
  set_points(const std::vector<Point<dim>> &pts);


  const Point<dim> &operator[](const unsigned int i) const;


  unsigned int
  size() const;


  std::vector<std::pair<unsigned int, double>>
  get_points_within_ball(const Point<dim> &target,
                         const double &    radius,
                         const bool        sorted = false) const;

  std::vector<std::pair<unsigned int, double>>
  get_closest_points(const Point<dim> & target,
                     const unsigned int n_points) const;

private:
  const unsigned int max_leaf_size;


  std::unique_ptr<PointCloudAdaptor> adaptor;


  std::unique_ptr<NanoFlannKDTree> kdtree;
};


//------------ inline functions -------------
#  ifndef DOXYGEN

template <int dim>
inline unsigned int
KDTree<dim>::size() const
{
  if (adaptor)
    return adaptor->points.size();
  else
    return 0;
}



template <int dim>
inline const Point<dim> &KDTree<dim>::operator[](const unsigned int i) const
{
  AssertIndexRange(i, size());
  return adaptor->points[i];
}



template <int dim>
KDTree<dim>::PointCloudAdaptor::PointCloudAdaptor(
  const std::vector<Point<dim>> &_points)
  : points(_points)
{}



template <int dim>
inline size_t
KDTree<dim>::PointCloudAdaptor::kdtree_get_point_count() const
{
  return points.size();
}



template <int dim>
inline double
KDTree<dim>::PointCloudAdaptor::kdtree_get_pt(const size_t idx, int d) const
{
  AssertIndexRange(d, dim);
  return points[idx][d];
}



template <int dim>
template <class BBOX>
inline bool
KDTree<dim>::PointCloudAdaptor::kdtree_get_bbox(BBOX &) const
{
  return false;
}



template <int dim>
inline double
KDTree<dim>::PointCloudAdaptor::kdtree_distance(const double *p1,
                                                const size_t  idx_p2,
                                                const size_t  size) const
{
  AssertDimension(size, dim);
  double res = 0.0;
  for (size_t d = 0; d < size; ++d)
    res += (p1[d] - points[idx_p2][d]) * (p1[d] - points[idx_p2][d]);
  return std::sqrt(res);
}
#  endif

DEAL_II_NAMESPACE_CLOSE

#endif // DEAL_II_WITH_NANO_FLANN
#endif