kdtree_ndarray_adapter.h

/*!
 *  Copyright (c) 2021 by Contributors
 * \file graph/transform/cpu/kdtree_ndarray_adapter.h
 * \brief NDArray adapter for nanoflann, without
 *        duplicating the storage
 */
#ifndef DGL_GRAPH_TRANSFORM_CPU_KDTREE_NDARRAY_ADAPTER_H_
#define DGL_GRAPH_TRANSFORM_CPU_KDTREE_NDARRAY_ADAPTER_H_

#include <dgl/array.h>
#include <dmlc/logging.h>
#include <nanoflann.hpp>
#include "../../../c_api_common.h"

namespace dgl {
namespace transform {
namespace knn_utils {

/*!
 * \brief A simple 2D NDArray adapter for nanoflann, without duplicating the storage.
 *
 * \tparam FloatType: The type of the point coordinates (typically, double or float).
 * \tparam IdType: The type for indices in the KD-tree index (typically, size_t of int)
 * \tparam FeatureDim: If set to > 0, it specifies a compile-time fixed dimensionality
 *         for the points in the data set, allowing more compiler optimizations.
 * \tparam Dist: The distance metric to use: nanoflann::metric_L1, nanoflann::metric_L2,
 *         nanoflann::metric_L2_Simple, etc.
 * \note The spelling of dgl's adapter ("adapter") is different from naneflann ("adaptor")
 */
template <typename FloatType,
          typename IdType,
          int FeatureDim = -1,
          typename Dist = nanoflann::metric_L2>
class KDTreeNDArrayAdapter {
 public:
  using self_type = KDTreeNDArrayAdapter<FloatType, IdType, FeatureDim, Dist>;
  using metric_type = typename Dist::template traits<FloatType, self_type>::distance_t;
  using index_type = nanoflann::KDTreeSingleIndexAdaptor<
    metric_type, self_type, FeatureDim, IdType>;

  KDTreeNDArrayAdapter(const size_t /* dims */,
                       const NDArray data_points,
                       const int leaf_max_size = 10)
      : data_(data_points) {
    CHECK(data_points->shape[0] != 0 && data_points->shape[1] != 0)
      << "Tensor containing input data point set must be 2D.";
    const size_t dims = data_points->shape[1];
    CHECK(!(FeatureDim > 0 && static_cast<int>(dims) != FeatureDim))
      << "Data set feature dimension does not match the 'FeatureDim' "
      << "template argument.";
    index_ = new index_type(
      static_cast<int>(dims), *this, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size));
    index_->buildIndex();
  }

  ~KDTreeNDArrayAdapter() {
    delete index_;
  }

  index_type* GetIndex() {
    return index_;
  }

  /*!
   * \brief Query for the \a num_closest points to a given point
   *  Note that this is a short-cut method for GetIndex()->findNeighbors().
   */
  void query(const FloatType* query_pt, const size_t num_closest,
             IdType* out_idxs, FloatType* out_dists) const {
    nanoflann::KNNResultSet<FloatType, IdType> resultSet(num_closest);
    resultSet.init(out_idxs, out_dists);
    index_->findNeighbors(resultSet, query_pt, nanoflann::SearchParams());
  }

  /*! \brief Interface expected by KDTreeSingleIndexAdaptor */
  const self_type& derived() const {
    return *this;
  }

  /*! \brief Interface expected by KDTreeSingleIndexAdaptor */
  self_type& derived() {
    return *this;
  }

  /*!
   * \brief Interface expected by KDTreeSingleIndexAdaptor,
   *  return the number of data points
   */
  size_t kdtree_get_point_count() const {
    return data_->shape[0];
  }

  /*!
   * \brief Interface expected by KDTreeSingleIndexAdaptor,
   *  return the dim'th component of the idx'th point
   */
  FloatType kdtree_get_pt(const size_t idx, const size_t dim) const {
    return data_.Ptr<FloatType>()[idx * data_->shape[1] + dim];
  }

  /*!
   * \brief Interface expected by KDTreeSingleIndexAdaptor.
   *  Optional bounding-box computation: return false to
   *  default to a standard bbox computation loop.
   *
   */
  template <typename BBOX>
  bool kdtree_get_bbox(BBOX& /* bb */) const {
    return false;
  }

 private:
  index_type* index_;    // The kd tree index
  const NDArray data_;  // data points
};

}  // namespace knn_utils
}  // namespace transform
}  // namespace dgl

#endif  // DGL_GRAPH_TRANSFORM_CPU_KDTREE_NDARRAY_ADAPTER_H_