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 at
// the top level of the deal.II distribution.
//
// ---------------------------------------------------------------------

#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 <memory>

#include <nanoflann.hpp>


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
  {
    typedef double coord_t;


    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;
  };


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


  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