FlannCache.h

Go to the documentation of this file.

#ifndef FLANNCACHE_H_
#define FLANNCACHE_H_
 
#include "MUQ/config.h"
 
#include <deque>
 
#include <nanoflann.hpp>
 
#include "MUQ/Modeling/LinearAlgebra/AnyAlgebra.h"
 
#include "MUQ/Modeling/ModPiece.h"
 
namespace muq {
  namespace Modeling {
 
    class FlannCache;
 
    template <class Distance = nanoflann::metric_L2, typename IndexType = size_t>
    struct DynamicKDTreeAdaptor {
 
      friend class FlannCache;
 
      typedef DynamicKDTreeAdaptor<Distance,IndexType> self_t;
      typedef typename Distance::template traits<double,self_t>::distance_t metric_t;
      typedef nanoflann::KDTreeSingleIndexDynamicAdaptor< metric_t,self_t,-1,IndexType>  index_t;
 
      std::shared_ptr<index_t> index; 
      std::deque<Eigen::VectorXd> m_data;
 
        inline DynamicKDTreeAdaptor(const int dim, const int leaf_max_size = 10) {
        index = std::make_shared<index_t>(dim, *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size ) );
      }
 
        inline void UpdateIndex(const int leaf_max_size = 10) {
          assert(m_data.size()>0);
        index = std::make_shared<index_t>(m_data.at(0).size(), *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size ) );
        index->addPoints(0, m_data.size()-1);
      }
 
      inline virtual ~DynamicKDTreeAdaptor() {}
 
      inline void add(Eigen::VectorXd const& newPt) {
        m_data.push_back(newPt);
          index->addPoints(m_data.size()-1, m_data.size()-1);
      }
 
      inline std::pair<std::vector<IndexType>, std::vector<double>> query(Eigen::VectorXd const& query_point, const size_t num_closest, const int nChecks_IGNORED = 10) const {
        std::vector<IndexType> out_indices(num_closest);
        std::vector<double> out_distances_sq(num_closest);
 
        nanoflann::KNNResultSet<double,IndexType> resultSet(num_closest);
        resultSet.init(&out_indices[0], &out_distances_sq[0]);
    #if MUQ_NANOFLAN_PARAMS_COMPILES
        index->findNeighbors(resultSet, query_point.data(), nanoflann::SearchParams());
    #else 
        index->findNeighbors(resultSet, query_point.data(), nanoflann::SearchParameters());
    #endif 
        return std::make_pair(out_indices, out_distances_sq);
      }
 
      inline const self_t & derived() const {
        return *this;
      }
 
      inline self_t& derived() {
        return *this;
      }
 
      // Must return the number of data points
      inline size_t kdtree_get_point_count() const {
        return m_data.size();
      }
 
      // Returns the dim'th component of the idx'th point in the class:
      inline double kdtree_get_pt(const size_t idx, int dim) const {
          assert(idx<m_data.size());
          assert(dim<m_data[idx].size());
 
          return m_data[idx][dim];
      }
 
      // Optional bounding-box computation: return false to default to a standard bbox computation loop.
      //   Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again.
      //   Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds)
      template <class BBOX>
      inline bool kdtree_get_bbox(BBOX & /*bb*/) const {
        return false;
      }
    }; // end of DynamicKDTreeAdaptor
 
 
 
    class FlannCache : public ModPiece {
    public:
 
      FlannCache(std::shared_ptr<ModPiece> function/* = nullptr*/);
 
      ~FlannCache();
 
 
      int InCache(Eigen::VectorXd const& input) const;
 
 
      std::vector<Eigen::VectorXd> Add(std::vector<Eigen::VectorXd> const& inputs);
 
 
      void Add(std::vector<Eigen::VectorXd> const& inputs, std::vector<Eigen::VectorXd> const& results);
 
 
      const Eigen::VectorXd at(unsigned int const index) const;
 
 
      Eigen::VectorXd at(unsigned int const index);
 
      Eigen::VectorXd const& OutputValue(unsigned int index) const;
 
 
      Eigen::VectorXd Add(Eigen::VectorXd const& input);
 
 
      unsigned int Add(Eigen::VectorXd const& input, Eigen::VectorXd const& result);
 
 
      void Remove(Eigen::VectorXd const& input);
 
 
      size_t NearestNeighborIndex(Eigen::VectorXd const& point) const;
 
 
      void NearestNeighbors(Eigen::VectorXd const& point,
                            unsigned int const k,
                            std::vector<Eigen::VectorXd>& neighbors,
                            std::vector<Eigen::VectorXd>& result) const;
 
 
      void NearestNeighbors(Eigen::VectorXd const& point,
                            unsigned int const k,
                            std::vector<Eigen::VectorXd>& neighbors) const;
 
 
      unsigned int Size() const;
 
 
      Eigen::VectorXd Centroid() const;
 
 
      std::shared_ptr<ModPiece> Function() const;
 
    private:
 
 
      void UpdateCentroid(Eigen::VectorXd const& point);
 
      // The vector of previous results
      std::vector<Eigen::VectorXd> outputCache;
 
      virtual void EvaluateImpl(ref_vector<Eigen::VectorXd> const& inputs) override;
 
      std::shared_ptr<ModPiece> function;
 
      std::shared_ptr<DynamicKDTreeAdaptor<>> kdTree;
 
 
      Eigen::VectorXd centroid;
    };
  } // namespace Utilities
} // namespace muq
 
#endif