nsearch.h

Go to the documentation of this file.

/*
 * -------------------------------------------
 * Copyright (c) 2021 - 2024 Prashant K. Jha
 * -------------------------------------------
 * PeriDEM https://github.com/prashjha/PeriDEM
 *
 * Distributed under the Boost Software License, Version 1.0. (See accompanying
 * file LICENSE)
 */
 
#ifndef NSEARCH_NSEARCH_H
#define NSEARCH_NSEARCH_H
 
#include "util/point.h" // definition of Point
#include "util/methods.h"
#include <cstdint> // uint8_t type
#include <string> // size_t type
#include <vector>
 
#include "nflannSetup.h"
 
namespace nsearch {
 
class BaseNSearch {
 
public:
  BaseNSearch(std::string name, size_t debug = 0)
      : d_debug(debug), d_treeType(name) {}
 
  virtual double updatePointCloud(const std::vector<util::Point> &x,
                        bool parallel = true) = 0;
 
  virtual double setInputCloud() = 0;
 
  virtual size_t radiusSearch(
      const util::Point &searchPoint, const double &search_r,
      std::vector<int> &neighs,
      std::vector<float> &sqr_dist) = 0;
 
  virtual size_t radiusSearch(
      const util::Point &searchPoint, const double &search_r,
      std::vector<size_t> &neighs,
      std::vector<double> &sqr_dist) = 0;
 
    virtual size_t radiusSearchExcludeTag(
            const util::Point &searchPoint,
            const double &search_r,
            std::vector<size_t> &neighs,
            std::vector<double> &sqr_dist,
            const size_t &searchPointTag,
            const std::vector<size_t> &dataTags) = 0;
 
    virtual size_t radiusSearchExcludeTag(
            const util::Point &searchPoint,
            const double &search_r,
            std::vector<int> &neighs,
            std::vector<float> &sqr_dist,
            const size_t &searchPointTag,
            const std::vector<size_t> &dataTags) = 0;
 
    virtual size_t radiusSearchIncludeTag(
            const util::Point &searchPoint,
            const double &search_r,
            std::vector<size_t> &neighs,
            std::vector<double> &sqr_dist,
            const size_t &searchPointTag,
            const std::vector<size_t> &dataTags) = 0;
 
    virtual size_t radiusSearchIncludeTag(
            const util::Point &searchPoint,
            const double &search_r,
            std::vector<int> &neighs,
            std::vector<float> &sqr_dist,
            const size_t &searchPointTag,
            const std::vector<size_t> &dataTags) = 0;
 
 
      virtual void closestPoint(
              const util::Point &searchPoint,
              size_t &neigh,
              double &sqrDistNeigh) = 0;
 
public:
  size_t d_debug;
 
  std::string d_treeType;
};
 
template <int dim = 3>
class NFlannSearchKd : public BaseNSearch {
 
public:
  explicit NFlannSearchKd(const PointCloud &x, size_t debug = 0, size_t max_leafs = 10)
          : BaseNSearch("nflann_kdtree", debug),
            d_cloud(x),
            d_tree(dim, d_cloud,
                   nanoflann::KDTreeSingleIndexAdaptorParams(max_leafs /* max leaf */)) {
    d_params.sorted = false;
  };
 
  double setInputCloud() override  {
    auto t1 = steady_clock::now();
    d_tree.buildIndex();
    auto t2 = steady_clock::now();
    return util::methods::timeDiff(t1, t2);
  };
 
  double updatePointCloud(const std::vector<util::Point> &x,
                          bool parallel = true) override {
    return 0;
  };
 
  size_t radiusSearch(
      const util::Point &searchPoint, const double &search_r,
      std::vector<size_t> &neighs,
      std::vector<double> &sqr_dist) override {
 
    double query_pt[3] = {searchPoint[0], searchPoint[1], searchPoint[2]};
 
    TreeSearchRes resultSet(search_r * search_r, neighs, sqr_dist);
    return d_tree.radiusSearchCustomCallback(&query_pt[0], resultSet, d_params);
  };
 
  size_t radiusSearch(
      const util::Point &searchPoint, const double &search_r,
      std::vector<int> &neighs,
      std::vector<float> &sqr_dist) override {
 
    // ugly but quick fix
    // first, get results using int and float and then convert
    std::vector<size_t> neighs_temp;
    std::vector<double> sqr_dist_temp;
    auto N =
            this->radiusSearch(searchPoint, search_r, neighs_temp, sqr_dist_temp);
 
    if (N > 0) {
      neighs.resize(N);
      sqr_dist.resize(N);
      for (size_t i=0; i<N; i++) {
        neighs.push_back(int(neighs_temp[i]));
        sqr_dist.push_back(float(sqr_dist_temp[i]));
      }
    }
 
    return N;
  };
 
 
    size_t radiusSearchExcludeTag(
            const util::Point &searchPoint,
            const double &search_r,
            std::vector<size_t> &neighs,
            std::vector<double> &sqr_dist,
            const size_t &searchPointTag,
            const std::vector<size_t> &dataTags) override {
 
      double query_pt[3] = {searchPoint[0], searchPoint[1], searchPoint[2]};
 
      TreeSearchCheckIDExcludeRes resultSet(search_r * search_r,
                            neighs, sqr_dist, searchPointTag, dataTags);
 
      return d_tree.radiusSearchCustomCallback(&query_pt[0], resultSet, d_params);
    };
 
    size_t radiusSearchExcludeTag(
            const util::Point &searchPoint,
            const double &search_r,
            std::vector<int> &neighs,
            std::vector<float> &sqr_dist,
            const size_t &searchPointTag,
            const std::vector<size_t> &dataTags) override {
 
      // first, get results using int and float and then convert
      std::vector<size_t> neighs_temp;
      std::vector<double> sqr_dist_temp;
      auto N =
              this->radiusSearchExcludeTag(searchPoint, search_r,
                                           neighs_temp, sqr_dist_temp,
                                           searchPointTag, dataTags);
 
      if (N > 0) {
        neighs.resize(N);
        sqr_dist.resize(N);
        for (size_t i=0; i<N; i++) {
          neighs.push_back(int(neighs_temp[i]));
          sqr_dist.push_back(float(sqr_dist_temp[i]));
        }
      }
 
      return N;
    };
 
    size_t radiusSearchIncludeTag(
            const util::Point &searchPoint,
            const double &search_r,
            std::vector<size_t> &neighs,
            std::vector<double> &sqr_dist,
            const size_t &searchPointTag,
            const std::vector<size_t> &dataTags) override {
 
      double query_pt[3] = {searchPoint[0], searchPoint[1], searchPoint[2]};
 
      TreeSearchCheckIDIncludeRes resultSet(search_r * search_r,
                                            neighs, sqr_dist,
                                            searchPointTag, dataTags);
 
      return d_tree.radiusSearchCustomCallback(&query_pt[0], resultSet, d_params);
    };
 
    size_t radiusSearchIncludeTag(
            const util::Point &searchPoint,
            const double &search_r,
            std::vector<int> &neighs,
            std::vector<float> &sqr_dist,
            const size_t &searchPointTag,
            const std::vector<size_t> &dataTags) override {
 
      // first, get results using int and float and then convert
      std::vector<size_t> neighs_temp;
      std::vector<double> sqr_dist_temp;
      auto N =
              this->radiusSearchIncludeTag(searchPoint, search_r,
                                           neighs_temp, sqr_dist_temp,
                                           searchPointTag, dataTags);
 
      if (N > 0) {
        neighs.resize(N);
        sqr_dist.resize(N);
        for (size_t i=0; i<N; i++) {
          neighs.push_back(int(neighs_temp[i]));
          sqr_dist.push_back(float(sqr_dist_temp[i]));
        }
      }
 
      return N;
    };
 
    void closestPoint(
            const util::Point &searchPoint,
            size_t &neigh,
            double &sqrDistNeigh) override {
 
      double query_pt[3] = {searchPoint[0], searchPoint[1], searchPoint[2]};
      unsigned int neigh_temp = 0;
      d_tree.knnSearch(&query_pt[0], 1, &neigh_temp, &sqrDistNeigh);
      neigh = neigh_temp;
    };
 
public:
  PointCloudAdaptor d_cloud;
 
  nanoflann::KDTreeSingleIndexAdaptor<
          nanoflann::L2_Simple_Adaptor<double, PointCloudAdaptor>, PointCloudAdaptor,
          dim
  > d_tree;
 
  nanoflann::SearchParameters d_params;
};
 
} // namespace nsearch
 
#endif // NSEARCH_NSEARCH_H