Merge pull request #3567 from shrit/float_knn

Change the template type from arma::mat to MatType in NS

src/mlpack/core/tree/binary_space_tree/binary_space_tree.hpp

@@ -504,7 +504,7 @@ class BinarySpaceTree
   size_t& Count() { return count; }
 
   //! Store the center of the bounding region in the given vector.
-  void Center(arma::vec& center) const { bound.Center(center); }
+  void Center(arma::Col<ElemType>& center) const { bound.Center(center); }
 
  private:
   /**

src/mlpack/core/tree/binary_space_tree/binary_space_tree_impl.hpp

@@ -909,7 +909,7 @@ void BinarySpaceTree<MetricType, StatisticType, MatType, BoundType, SplitType>::
       splitter, maxLeafSize);
 
   // Calculate parent distances for those two nodes.
-  arma::vec center, leftCenter, rightCenter;
+  arma::Col<ElemType> center, leftCenter, rightCenter;
   Center(center);
   left->Center(leftCenter);
   right->Center(rightCenter);
@@ -977,7 +977,7 @@ SplitNode(std::vector<size_t>& oldFromNew,
       oldFromNew, splitter, maxLeafSize);
 
   // Calculate parent distances for those two nodes.
-  arma::vec center, leftCenter, rightCenter;
+  arma::Col<ElemType> center, leftCenter, rightCenter;
   Center(center);
   left->Center(leftCenter);
   right->Center(rightCenter);

src/mlpack/methods/neighbor_search/neighbor_search.hpp

@@ -82,6 +82,8 @@ class NeighborSearch
  public:
   //! Convenience typedef.
   typedef TreeType<MetricType, NeighborSearchStat<SortPolicy>, MatType> Tree;
+  //! The type of element held in MatType.
+  typedef typename MatType::elem_type ElemType;
 
   /**
    * Initialize the NeighborSearch object, passing a reference dataset (this is
@@ -223,7 +225,7 @@ class NeighborSearch
   void Search(const MatType& querySet,
               const size_t k,
               arma::Mat<size_t>& neighbors,
-              arma::mat& distances);
+              arma::Mat<ElemType>& distances);
 
   /**
    * Given a pre-built query tree, search for the nearest neighbors of each
@@ -248,7 +250,7 @@ class NeighborSearch
   void Search(Tree& queryTree,
               const size_t k,
               arma::Mat<size_t>& neighbors,
-              arma::mat& distances,
+              arma::Mat<ElemType>& distances,
               bool sameSet = false);
 
   /**
@@ -267,7 +269,7 @@ class NeighborSearch
    */
   void Search(const size_t k,
               arma::Mat<size_t>& neighbors,
-              arma::mat& distances);
+              arma::Mat<ElemType>& distances);
 
   /**
    * Calculate the average relative error (effective error) between the
@@ -284,8 +286,8 @@ class NeighborSearch
    *     query point.
    * @return Average relative error.
    */
-  static double EffectiveError(arma::mat& foundDistances,
-                               arma::mat& realDistances);
+  static double EffectiveError(arma::Mat<ElemType>& foundDistances,
+                               arma::Mat<ElemType>& realDistances);
 
   /**
    * Calculate the recall (% of neighbors found) given the list of found

src/mlpack/methods/neighbor_search/neighbor_search_impl.hpp

@@ -104,7 +104,7 @@ SingleTreeTraversalType>::NeighborSearch(const NeighborSearchMode mode,
   // Build the tree on the empty dataset, if necessary.
   if (mode != NAIVE_MODE)
   {
-    referenceTree = BuildTree<Tree>(std::move(arma::mat()),
+    referenceTree = BuildTree<Tree>(std::move(MatType()),
         oldFromNewReferences);
     referenceSet = &referenceTree->Dataset();
   }
@@ -255,7 +255,7 @@ NeighborSearch<SortPolicy,
   if (!other.referenceTree)
     delete other.referenceSet;
 
-  other.referenceTree = BuildTree<Tree>(std::move(arma::mat()),
+  other.referenceTree = BuildTree<Tree>(std::move(MatType()),
       other.oldFromNewReferences);
   other.referenceSet = &other.referenceTree->Dataset();
   other.searchMode = DUAL_TREE_MODE,
@@ -365,7 +365,7 @@ DualTreeTraversalType, SingleTreeTraversalType>::Search(
     const MatType& querySet,
     const size_t k,
     arma::Mat<size_t>& neighbors,
-    arma::mat& distances)
+    arma::Mat<ElemType>& distances)
 {
   if (k > referenceSet->n_cols)
   {
@@ -386,14 +386,14 @@ DualTreeTraversalType, SingleTreeTraversalType>::Search(
   // To avoid an extra copy, we will store the neighbors and distances in a
   // separate matrix.
   arma::Mat<size_t>* neighborPtr = &neighbors;
-  arma::mat* distancePtr = &distances;
+  arma::Mat<ElemType>* distancePtr = &distances;
 
   // Mapping is only necessary if the tree rearranges points.
   if (TreeTraits<Tree>::RearrangesDataset)
   {
     if (searchMode == DUAL_TREE_MODE)
     {
-      distancePtr = new arma::mat; // Query indices need to be mapped.
+      distancePtr = new arma::Mat<ElemType>; // Query indices need to be mapped.
       neighborPtr = new arma::Mat<size_t>;
     }
     else if (!oldFromNewReferences.empty())
@@ -570,7 +570,7 @@ DualTreeTraversalType, SingleTreeTraversalType>::Search(
     Tree& queryTree,
     const size_t k,
     arma::Mat<size_t>& neighbors,
-    arma::mat& distances,
+    arma::Mat<ElemType>& distances,
     bool sameSet)
 {
   if (k > referenceSet->n_cols)
@@ -648,7 +648,7 @@ void NeighborSearch<SortPolicy, MetricType, MatType, TreeType,
 DualTreeTraversalType, SingleTreeTraversalType>::Search(
     const size_t k,
     arma::Mat<size_t>& neighbors,
-    arma::mat& distances)
+    arma::Mat<ElemType>& distances)
 {
   if (k > referenceSet->n_cols)
   {
@@ -670,12 +670,12 @@ DualTreeTraversalType, SingleTreeTraversalType>::Search(
   scores = 0;
 
   arma::Mat<size_t>* neighborPtr = &neighbors;
-  arma::mat* distancePtr = &distances;
+  arma::Mat<ElemType>* distancePtr = &distances;
 
   if (!oldFromNewReferences.empty() && TreeTraits<Tree>::RearrangesDataset)
   {
     // We will always need to rearrange in this case.
-    distancePtr = new arma::mat;
+    distancePtr = new MatType;
     neighborPtr = new arma::Mat<size_t>;
   }
 
@@ -825,8 +825,8 @@ template<typename SortPolicy,
          template<typename> class SingleTreeTraversalType>
 double NeighborSearch<SortPolicy, MetricType, MatType, TreeType,
 DualTreeTraversalType, SingleTreeTraversalType>::EffectiveError(
-    arma::mat& foundDistances,
-    arma::mat& realDistances)
+    arma::Mat<ElemType>& foundDistances,
+    arma::Mat<ElemType>& realDistances)
 {
   if (foundDistances.n_rows != realDistances.n_rows ||
       foundDistances.n_cols != realDistances.n_cols)

src/mlpack/methods/neighbor_search/neighbor_search_rules.hpp

@@ -34,6 +34,9 @@ template<typename SortPolicy, typename MetricType, typename TreeType>
 class NeighborSearchRules
 {
  public:
+  //! The type of element held in MatType.
+  typedef typename TreeType::Mat::elem_type ElemType;
+
   /**
    * Construct the NeighborSearchRules object.  This is usually done from within
    * the NeighborSearch class at search time.
@@ -60,7 +63,7 @@ class NeighborSearchRules
    * @param distances Matrix storing distances of neighbors for each query
    *     point.
    */
-  void GetResults(arma::Mat<size_t>& neighbors, arma::mat& distances);
+  void GetResults(arma::Mat<size_t>& neighbors, arma::Mat<ElemType>& distances);
 
   /**
    * Get the distance from the query point to the reference point.

src/mlpack/methods/neighbor_search/neighbor_search_rules_impl.hpp

@@ -61,7 +61,7 @@ NeighborSearchRules<SortPolicy, MetricType, TreeType>::NeighborSearchRules(
 template<typename SortPolicy, typename MetricType, typename TreeType>
 void NeighborSearchRules<SortPolicy, MetricType, TreeType>::GetResults(
     arma::Mat<size_t>& neighbors,
-    arma::mat& distances)
+    arma::Mat<ElemType>& distances)
 {
   neighbors.set_size(k, querySet.n_cols);
   distances.set_size(k, querySet.n_cols);

src/mlpack/tests/knn_test.cpp

@@ -16,6 +16,18 @@
 
 using namespace mlpack;
 
+/**
+ * A couple of handful declarations for float32 testing.
+ * These will be removed when we refactor the Bounds to accept MatType.
+ * For now, we will keep the following declarations.
+ */
+template<typename MetricType>
+using FloatHRectBound = HRectBound<MetricType, float>;
+
+template<typename MetricType, typename StatisticType, typename MatType>
+using FloatKDTree = BinarySpaceTree<MetricType, StatisticType, MatType,
+                                    FloatHRectBound, MidpointSplit>;
+
 /**
  * Test that Unmap() works in the dual-tree case (see unmap.hpp).
  */
@@ -746,6 +758,48 @@ TEST_CASE("KNNSingleTreeVsNaive", "[KNNTest]")
   }
 }
 
+/**
+ * Test the single-tree nearest-neighbors method with the naive method.
+ *
+ * The main difference with the above test is that this one loads the reference
+ * dataset as a float32, and the distances as a float32 as well.
+ *
+ * Errors are produced if the results are not identical.
+ */
+TEST_CASE("KNNSingleTreeVsNaiveF32", "[KNNTest]")
+{
+  arma::fmat dataset;
+
+  // Hard-coded filename: bad?
+  // Code duplication: also bad!
+  if (!data::Load("test_data_3_1000.csv", dataset))
+    FAIL("Cannot load test dataset test_data_3_1000.csv!");
+
+  NeighborSearch<NearestNeighborSort,
+                 EuclideanDistance,
+                 arma::fmat,
+                 FloatKDTree> knn(dataset, SINGLE_TREE_MODE);
+
+  // Set up computation for naive mode.
+  NeighborSearch<NearestNeighborSort,
+                 EuclideanDistance,
+                 arma::fmat,
+                 FloatKDTree> naive(dataset, NAIVE_MODE);
+
+  arma::Mat<size_t> neighborsTree;
+  arma::fmat distancesTree;
+  knn.Search(15, neighborsTree, distancesTree);
+
+  arma::Mat<size_t> neighborsNaive;
+  arma::fmat distancesNaive;
+  naive.Search(15, neighborsNaive, distancesNaive);
+
+  for (size_t i = 0; i < neighborsTree.n_elem; ++i)
+  {
+    REQUIRE(neighborsTree[i] ==neighborsNaive[i]);
+    REQUIRE(distancesTree[i] == Approx(distancesNaive[i]).epsilon(1e-7));
+  }
+}
 /**
  * Test the cover tree single-tree nearest-neighbors method against the naive
  * method.  This uses only a random reference dataset.