Spill trees

src/mlpack/methods/neighbor_search/knn_main.cpp

@@ -70,6 +70,7 @@ PARAM_INT_IN("leaf_size", "Leaf size for tree building (used for kd-trees, R "
     "trees, R* trees, X trees, Hilbert R trees, R+ trees, R++ trees, and Spill "
     "trees).", "l", 20);
 PARAM_DOUBLE_IN("tau", "Overlapping size (for spill trees).", "u", 0);
+PARAM_DOUBLE_IN("rho", "Balance threshold (for spill trees).", "b", 0.7);
 
 PARAM_FLAG("random_basis", "Before tree-building, project the data onto a "
     "random orthogonal basis.", "R");
@@ -117,6 +118,9 @@ int main(int argc, char *argv[])
     if (CLI::HasParam("tau"))
       Log::Warn << "--tau (-u) will be ignored because --input_model_file"
           << " is specified." << endl;
+    if (CLI::HasParam("rho"))
+      Log::Warn << "--rho (-b) will be ignored because --input_model_file"
+          << " is specified." << endl;
     if (CLI::HasParam("random_basis"))
       Log::Warn << "--random_basis (-R) will be ignored because "
           << "--input_model_file is specified." << endl;
@@ -157,6 +161,14 @@ int main(int argc, char *argv[])
   if (CLI::HasParam("tau") && "spill" != CLI::GetParam<string>("tree_type"))
     Log::Fatal << "Tau parameter is only valid for spill trees." << endl;
 
+  // Sanity check on rho.
+  const double rho = CLI::GetParam<double>("rho");
+  if (rho < 0 || rho > 1)
+    Log::Fatal << "Invalid rho: " << rho << ".  Must be in the range [0,1]. "
+        << endl;
+  if (CLI::HasParam("rho") && "spill" != CLI::GetParam<string>("tree_type"))
+    Log::Fatal << "Rho parameter is only valid for spill trees." << endl;
+
   // Sanity check on epsilon.
   const double epsilon = CLI::GetParam<double>("epsilon");
   if (epsilon < 0)
@@ -204,6 +216,7 @@ int main(int argc, char *argv[])
     knn.RandomBasis() = randomBasis;
     knn.LeafSize() = size_t(lsInt);
     knn.Tau() = tau;
+    knn.Rho() = rho;
 
     arma::mat referenceSet;
     data::Load(referenceFile, referenceSet, true);
@@ -231,6 +244,7 @@ int main(int argc, char *argv[])
     knn.LeafSize() = size_t(lsInt);
     knn.Epsilon() = epsilon;
     knn.Tau() = tau;
+    knn.Rho() = rho;
   }
 
   // Perform search, if desired.

src/mlpack/methods/neighbor_search/ns_model.hpp

@@ -110,6 +110,8 @@ class BiSearchVisitor : public boost::static_visitor<void>
   const size_t leafSize;
   //! Overlapping size (for spill trees).
   const double tau;
+  //! Balance threshold (for spill trees).
+  const double rho;
 
   //! Bichromatic neighbor search on the given NSType considering the leafSize.
   template<typename NSType>
@@ -143,7 +145,8 @@ class BiSearchVisitor : public boost::static_visitor<void>
                   arma::Mat<size_t>& neighbors,
                   arma::mat& distances,
                   const size_t leafSize,
-                  const double tau);
+                  const double tau,
+                  const double rho);
 };
 
 /**
@@ -162,6 +165,8 @@ class TrainVisitor : public boost::static_visitor<void>
   size_t leafSize;
   //! Overlapping size (for spill trees).
   const double tau;
+  //! Balance threshold (for spill trees).
+  const double rho;
 
   //! Train on the given NSType considering the leafSize.
   template<typename NSType>
@@ -190,10 +195,11 @@ class TrainVisitor : public boost::static_visitor<void>
   void operator()(NSSpillType* ns) const;
 
   //! Construct the TrainVisitor object with the given reference set, leafSize
-  //! for BinarySpaceTrees, and tau for spill trees.
+  //! for BinarySpaceTrees, and tau and rho for spill trees.
   TrainVisitor(arma::mat&& referenceSet,
                const size_t leafSize,
-               const double tau);
+               const double tau,
+               const double rho);
 };
 
 /**
@@ -289,6 +295,8 @@ class NSModel
 
   //! Overlapping size (for spill trees).
   double tau;
+  //! Balance threshold (for spill trees).
+  double rho;
 
   //! If true, random projections are used.
   bool randomBasis;
@@ -348,6 +356,10 @@ class NSModel
   double Tau() const { return tau; }
   double& Tau() { return tau; }
 
+  //! Expose rho.
+  double Rho() const { return rho; }
+  double& Rho() { return rho; }
+
   //! Expose treeType.
   TreeTypes TreeType() const { return treeType; }
   TreeTypes& TreeType() { return treeType; }

src/mlpack/methods/neighbor_search/ns_model_impl.hpp

@@ -34,13 +34,15 @@ BiSearchVisitor<SortPolicy>::BiSearchVisitor(const arma::mat& querySet,
                                              arma::Mat<size_t>& neighbors,
                                              arma::mat& distances,
                                              const size_t leafSize,
-                                             const double tau) :
+                                             const double tau,
+                                             const double rho) :
     querySet(querySet),
     k(k),
     neighbors(neighbors),
     distances(distances),
     leafSize(leafSize),
-    tau(tau)
+    tau(tau),
+    rho(rho)
 {}
 
 //! Default Bichromatic neighbor search on the given NSType instance.
@@ -84,7 +86,7 @@ void BiSearchVisitor<SortPolicy>::operator()(NSSpillType* ns) const
       // For Dual Tree Search on SpillTrees, the queryTree must be built with
       // non overlapping (tau = 0).
       typename NSSpillType::Tree queryTree(std::move(querySet), 0 /* tau*/,
-          leafSize);
+          leafSize, rho);
       ns->Search(&queryTree, k, neighbors, distances);
     }
     else
@@ -126,10 +128,12 @@ void BiSearchVisitor<SortPolicy>::SearchLeaf(NSType* ns) const
 template<typename SortPolicy>
 TrainVisitor<SortPolicy>::TrainVisitor(arma::mat&& referenceSet,
                                        const size_t leafSize,
-                                       const double tau) :
+                                       const double tau,
+                                       const double rho) :
     referenceSet(std::move(referenceSet)),
     leafSize(leafSize),
-    tau(tau)
+    tau(tau),
+    rho(rho)
 {}
 
 //! Default Train on the given NSType instance.
@@ -173,7 +177,7 @@ void TrainVisitor<SortPolicy>::operator ()(NSSpillType* ns) const
     else
     {
       typename NSSpillType::Tree* tree = new typename NSSpillType::Tree(
-          std::move(referenceSet), tau, leafSize);
+          std::move(referenceSet), tau, leafSize, rho);
       ns->Train(tree);
       // Give the model ownership of the tree.
       ns->neighborSearch.treeOwner = true;
@@ -257,6 +261,7 @@ NSModel<SortPolicy>::NSModel(TreeTypes treeType, bool randomBasis) :
     treeType(treeType),
     leafSize(20),
     tau(0),
+    rho(0.7),
     randomBasis(randomBasis)
 {
   // Nothing to do.
@@ -307,6 +312,7 @@ void NSModel<SortPolicy>::Serialize(Archive& ar,
   ar & data::CreateNVP(treeType, "treeType");
   ar & data::CreateNVP(leafSize, "leafSize");
   ar & data::CreateNVP(tau, "tau");
+  ar & data::CreateNVP(rho, "rho");
   ar & data::CreateNVP(randomBasis, "randomBasis");
   ar & data::CreateNVP(q, "q");
 
@@ -454,11 +460,11 @@ void NSModel<SortPolicy>::BuildModel(arma::mat&& referenceSet,
           epsilon);
       break;
     case SPILL_TREE:
-      nSearch = new NSSpillType(naive, singleMode, tau, epsilon);
+      nSearch = new NSSpillType(naive, singleMode, tau, leafSize, rho, epsilon);
       break;
   }
 
-  TrainVisitor<SortPolicy> tn(std::move(referenceSet), leafSize, tau);
+  TrainVisitor<SortPolicy> tn(std::move(referenceSet), leafSize, tau, rho);
   boost::apply_visitor(tn, nSearch);
 
   if (!naive)
@@ -491,7 +497,7 @@ void NSModel<SortPolicy>::Search(arma::mat&& querySet,
         << std::endl;
 
   BiSearchVisitor<SortPolicy> search(querySet, k, neighbors, distances,
-      leafSize, tau);
+      leafSize, tau, rho);
   boost::apply_visitor(search, nSearch);
 }
 

src/mlpack/methods/neighbor_search/spill_search.hpp

@@ -68,13 +68,17 @@ class SpillSearch
    * @param singleMode If true, single-tree search will be used (as opposed to
    *      dual-tree search).
    * @param tau Overlapping size (non-negative).
+   * @param leafSize Max size of each leaf in the tree.
+   * @param rho Balance threshold (non-negative).
    * @param epsilon Relative approximate error (non-negative).
    * @param metric An optional instance of the MetricType class.
    */
   SpillSearch(const MatType& referenceSet,
               const bool naive = false,
               const bool singleMode = false,
               const double tau = 0,
+              const double leafSize = 20,
+              const double rho = 0.7,
               const double epsilon = 0,
               const MetricType metric = MetricType());
 
@@ -91,13 +95,17 @@ class SpillSearch
    * @param singleMode If true, single-tree search will be used (as opposed to
    *      dual-tree search).
    * @param tau Overlapping size (non-negative).
+   * @param leafSize Max size of each leaf in the tree.
+   * @param rho Balance threshold (non-negative).
    * @param epsilon Relative approximate error (non-negative).
    * @param metric An optional instance of the MetricType class.
    */
   SpillSearch(MatType&& referenceSet,
               const bool naive = false,
               const bool singleMode = false,
               const double tau = 0,
+              const double leafSize = 20,
+              const double rho = 0.7,
               const double epsilon = 0,
               const MetricType metric = MetricType());
 
@@ -113,12 +121,16 @@ class SpillSearch
    * @param singleMode Whether single-tree computation should be used (as
    *      opposed to dual-tree computation).
    * @param tau Overlapping size (non-negative).
+   * @param leafSize Max size of each leaf in the tree.
+   * @param rho Balance threshold (non-negative).
    * @param epsilon Relative approximate error (non-negative).
    * @param metric Instantiated distance metric.
    */
   SpillSearch(Tree* referenceTree,
               const bool singleMode = false,
               const double tau = 0,
+              const double leafSize = 20,
+              const double rho = 0.7,
               const double epsilon = 0,
               const MetricType metric = MetricType());
 
@@ -131,12 +143,16 @@ class SpillSearch
    * @param singleMode Whether single-tree computation should be used (as
    *      opposed to dual-tree computation).
    * @param tau Overlapping size (non-negative).
+   * @param leafSize Max size of each leaf in the tree.
+   * @param rho Balance threshold (non-negative).
    * @param epsilon Relative approximate error (non-negative).
    * @param metric Instantiated metric.
    */
   SpillSearch(const bool naive = false,
               const bool singleMode = false,
               const double tau = 0,
+              const double leafSize = 20,
+              const double rho = 0.7,
               const double epsilon = 0,
               const MetricType metric = MetricType());
 
@@ -262,6 +278,12 @@ class SpillSearch
   //! Access the overlapping size.
   double Tau() const { return tau; }
 
+  //! Access the balance threshold.
+  double Rho() const { return rho; }
+
+  //! Access the leaf size.
+  double LeafSize() const { return leafSize; }
+
   //! Access the reference dataset.
   const MatType& ReferenceSet() const { return neighborSearch.ReferenceSet(); }
 
@@ -277,6 +299,12 @@ class SpillSearch
   //! Overlapping size.
   double tau;
 
+  //! Balance threshold.
+  double rho;
+
+  //! Max leaf size.
+  double leafSize;
+
   //! The NSModel class should have access to internal members.
   template<typename SortPolicy>
   friend class TrainVisitor;