1 #ifndef FLANNCACHE_H_
2 #define FLANNCACHE_H_
3 
4 #include <deque>
5 
6 #include <nanoflann.hpp>
7 
8 #include "MUQ/Modeling/LinearAlgebra/AnyAlgebra.h"
9 
10 #include "MUQ/Modeling/ModPiece.h"
11 
12 namespace muq {
13   namespace Modeling {
14 
15     class FlannCache;
16 
17     template <class Distance = nanoflann::metric_L2, typename IndexType = size_t>
18     struct DynamicKDTreeAdaptor {
19 
20       friend class FlannCache;
21 
22       typedef DynamicKDTreeAdaptor<Distance,IndexType> self_t;
23       typedef typename Distance::template traits<double,self_t>::distance_t metric_t;
24       typedef nanoflann::KDTreeSingleIndexDynamicAdaptor< metric_t,self_t,-1,IndexType>  index_t;
25 
26       std::shared_ptr<index_t> index; //! The kd-tree index for the user to call its methods as usual with any other FLANN index.
27       std::deque<Eigen::VectorXd> m_data;
28 
29       /// Constructor: takes a const ref to the vector of vectors object with the data points
30 	    inline DynamicKDTreeAdaptor(const int dim, const int leaf_max_size = 10) {
31         index = std::make_shared<index_t>(dim, *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size ) );
32       }
33 
34 	    inline void UpdateIndex(const int leaf_max_size = 10) {
35 	      assert(m_data.size()>0);
36         index = std::make_shared<index_t>(m_data.at(0).size(), *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size ) );
37         index->addPoints(0, m_data.size()-1);
38       }
39 
~DynamicKDTreeAdaptorDynamicKDTreeAdaptor40       inline virtual ~DynamicKDTreeAdaptor() {}
41 
addDynamicKDTreeAdaptor42       inline void add(Eigen::VectorXd const& newPt) {
43         m_data.push_back(newPt);
44 	      index->addPoints(m_data.size()-1, m_data.size()-1);
45       }
46 
47       /** Query for the \a num_closest closest points to a given point (entered as query_point[0:dim-1]).
48     	 *  Note that this is a short-cut method for index->findNeighbors().
49     	 *  The user can also call index->... methods as desired.
50     	 * \note nChecks_IGNORED is ignored but kept for compatibility with the original FLANN interface.
51     	 */
52       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 {
53         std::vector<IndexType> out_indices(num_closest);
54         std::vector<double> out_distances_sq(num_closest);
55 
56         nanoflann::KNNResultSet<double,IndexType> resultSet(num_closest);
57         resultSet.init(&out_indices[0], &out_distances_sq[0]);
58         index->findNeighbors(resultSet, query_point.data(), nanoflann::SearchParams());
59 
60         return std::make_pair(out_indices, out_distances_sq);
61       }
62 
derivedDynamicKDTreeAdaptor63       inline const self_t & derived() const {
64         return *this;
65       }
66 
derivedDynamicKDTreeAdaptor67       inline self_t& derived() {
68         return *this;
69       }
70 
71       // Must return the number of data points
kdtree_get_point_countDynamicKDTreeAdaptor72       inline size_t kdtree_get_point_count() const {
73         return m_data.size();
74       }
75 
76       // Returns the dim'th component of the idx'th point in the class:
kdtree_get_ptDynamicKDTreeAdaptor77       inline double kdtree_get_pt(const size_t idx, int dim) const {
78 	      assert(idx<m_data.size());
79 	      assert(dim<m_data[idx].size());
80 
81 	      return m_data[idx][dim];
82       }
83 
84       // Optional bounding-box computation: return false to default to a standard bbox computation loop.
85       //   Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again.
86       //   Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds)
87       template <class BBOX>
kdtree_get_bboxDynamicKDTreeAdaptor88       inline bool kdtree_get_bbox(BBOX & /*bb*/) const {
89         return false;
90       }
91     }; // end of DynamicKDTreeAdaptor
92 
93 
94     /// Create a cache of model evaluations (input/output pairs)
95     /**
96        Caches the input/output pairs for a muq::Modeling::WorkPiece that has one input and one output.
97      */
98     class FlannCache : public ModPiece {
99     public:
100 
101       /**
102 	 If we pass a nullptr (default), then only construct a cache of points (no outputs)
103 	 @param[in] function The function whose input/output pairs we want to cache
104        */
105       FlannCache(std::shared_ptr<ModPiece> function/* = nullptr*/);
106 
107       ~FlannCache();
108 
109       /// Determine if an entry is in the cache
110       /**
111 	     @param[in] input Check if this input is in the cache
112 	      \return -1: the input vector does not have an entry in the cache, otherise: return the index of the input in the cache
113        */
114       int InCache(Eigen::VectorXd const& input) const;
115 
116       /// Add new points to the cache
117       /**
118 	 @param[in] inputs The points to add
119 	 \return The function evaluation at each input point
120        */
121       std::vector<Eigen::VectorXd> Add(std::vector<Eigen::VectorXd> const& inputs);
122 
123       /// Add new points to the cache given the result
124       /**
125 	 @param[in] inputs The points to add
126 	 @param[in] results The function evaluation at each input point
127        */
128       void Add(std::vector<Eigen::VectorXd> const& inputs, std::vector<Eigen::VectorXd> const& results);
129 
130       /// Get an input point from the cache
131       /**
132 	 @param[in] index The point to returnn
133 	 \return The input point associated with the index
134        */
135       const Eigen::VectorXd at(unsigned int const index) const;
136 
137       /// Get an input point from the cache
138       /**
139 	 @param[in] index The point to returnn
140 	 \return The input point associated with the index
141        */
142       Eigen::VectorXd at(unsigned int const index);
143 
144       /// Returns the model for a specific cache index
145       Eigen::VectorXd const& OutputValue(unsigned int index) const;
146 
147       /// Add a new point to the cache
148       /**
149 	 @param[in] input The entry we would like to add to the cache (if it is not there already)
150 	 \return The function result at that point
151        */
152       Eigen::VectorXd Add(Eigen::VectorXd const& input);
153 
154       /// Add a new point to the cache given the result
155       /**
156 	 @param[in] input The entry we would like to add to the cache (if it is not there already)
157 	 @param[in] result The result of the function
158    @return Returns the index of the added point
159       */
160       unsigned int Add(Eigen::VectorXd const& input, Eigen::VectorXd const& result);
161 
162       /// Remove point from the cache
163       /**
164 	 @param[in] input The entry we would like to remove from the cache
165        */
166       void Remove(Eigen::VectorXd const& input);
167 
168       /// The index of the nearest neighbor
169       /**
170 	 @param[in] point We want the index of the nearest neighbor that is closest to this point.
171 	 \return The index of the nearest neighbor
172        */
173       size_t NearestNeighborIndex(Eigen::VectorXd const& point) const;
174 
175       /// Find the \f$k\f$ nearest neighbors
176       /**
177 	 @param[in] point The point whose nearest neighbors we want to find
178 	 @param[in] k We want to find this many nearest neighbors
179 	 @param[out] neighbors A vector of the \fk\f$ nearest neighbors
180 	 @param[out] result The output corresponding to the \f$k\f$ nearest neighbors
181        */
182       void NearestNeighbors(Eigen::VectorXd const& point,
183                             unsigned int const k,
184                             std::vector<Eigen::VectorXd>& neighbors,
185                             std::vector<Eigen::VectorXd>& result) const;
186 
187       /// Find the \f$k\f$ nearest neighbors (don't bother getting the result too)
188       /**
189 	 @param[in] point The point whose nearest neighbors we want to find
190 	 @param[in] k We want to find this many nearest neighbors
191 	 @param[out] neighbors A vector of the \fk\f$ nearest neighbors
192        */
193       void NearestNeighbors(Eigen::VectorXd const& point,
194                             unsigned int const k,
195                             std::vector<Eigen::VectorXd>& neighbors) const;
196 
197       /// Get the size of the cache
198       /**
199 	     \return The size of the cache
200        */
201       unsigned int Size() const;
202 
203       /// Get the centroid of the cache
204       /**
205       \return The cache centroid
206       */
207       Eigen::VectorXd Centroid() const;
208 
209       /// Get the underlying function
210       /**
211       \return The function whose output we are storing
212       */
213       std::shared_ptr<ModPiece> Function() const;
214 
215     private:
216 
217       /// Update the centroid when a new point is added to the cache
218       /**
219         @param[in] point The new point
220       */
221       void UpdateCentroid(Eigen::VectorXd const& point);
222 
223       // The vector of previous results
224       std::vector<Eigen::VectorXd> outputCache;
225 
226       virtual void EvaluateImpl(ref_vector<Eigen::VectorXd> const& inputs) override;
227 
228       /// The function whose input/outputs we are caching
229       std::shared_ptr<ModPiece> function;
230 
231       /// The nearest neighbor index, used to perform searches
232       std::shared_ptr<DynamicKDTreeAdaptor<>> kdTree;
233 
234       /// The centroid of the input addPoints
235       /**
236         The centroid is
237         \f{equation}{
238           \bar{x} = \frac{1}{N} \sum_{i=1}^N x_i,
239         \f}
240         where \f$N\f$ is the cache size.
241 
242         Note: the centroid is not necessarily in the cache.
243       */
244       Eigen::VectorXd centroid;
245     };
246   } // namespace Utilities
247 } // namespace muq
248 
249 #endif
250