neighbor_search.hpp

Go to the documentation of this file.

#ifndef MLPACK_METHODS_NEIGHBOR_SEARCH_NEIGHBOR_SEARCH_HPP
#define MLPACK_METHODS_NEIGHBOR_SEARCH_NEIGHBOR_SEARCH_HPP

#include <mlpack/prereqs.hpp>
#include <vector>
#include <string>

#include <mlpack/core/tree/binary_space_tree.hpp>
#include <mlpack/core/tree/rectangle_tree.hpp>
#include <mlpack/core/tree/binary_space_tree/binary_space_tree.hpp>

#include "neighbor_search_stat.hpp"
#include "sort_policies/nearest_neighbor_sort.hpp"
#include "neighbor_search_rules.hpp"

namespace mlpack {
namespace neighbor  {

// Forward declaration.
template<typename SortPolicy>
class TrainVisitor;

enum NeighborSearchMode
{
  NAIVE_MODE,
  SINGLE_TREE_MODE,
  DUAL_TREE_MODE,
  GREEDY_SINGLE_TREE_MODE
};

template<typename SortPolicy = NearestNeighborSort,
         typename MetricType = mlpack::metric::EuclideanDistance,
         typename MatType = arma::mat,
         template<typename TreeMetricType,
                  typename TreeStatType,
                  typename TreeMatType> class TreeType = tree::KDTree,
         template<typename RuleType> class DualTreeTraversalType =
             TreeType<MetricType,
                      NeighborSearchStat<SortPolicy>,
                      MatType>::template DualTreeTraverser,
         template<typename RuleType> class SingleTreeTraversalType =
             TreeType<MetricType,
                      NeighborSearchStat<SortPolicy>,
                      MatType>::template SingleTreeTraverser>
class NeighborSearch
{
 public:
  typedef TreeType<MetricType, NeighborSearchStat<SortPolicy>, MatType> Tree;

  NeighborSearch(const MatType& referenceSet,
                 const NeighborSearchMode mode = DUAL_TREE_MODE,
                 const double epsilon = 0,
                 const MetricType metric = MetricType());

  NeighborSearch(MatType&& referenceSet,
                 const NeighborSearchMode mode = DUAL_TREE_MODE,
                 const double epsilon = 0,
                 const MetricType metric = MetricType());

  NeighborSearch(
      const Tree& referenceTree,
      const NeighborSearchMode mode = DUAL_TREE_MODE,
      const double epsilon = 0,
      const MetricType metric = MetricType());

  NeighborSearch(
      Tree&& referenceTree,
      const NeighborSearchMode mode = DUAL_TREE_MODE,
      const double epsilon = 0,
      const MetricType metric = MetricType());

  NeighborSearch(const NeighborSearchMode mode = DUAL_TREE_MODE,
                 const double epsilon = 0,
                 const MetricType metric = MetricType());

  NeighborSearch(const NeighborSearch& other);

  NeighborSearch(NeighborSearch&& other);

  NeighborSearch& operator=(const NeighborSearch& other);

  NeighborSearch& operator=(NeighborSearch&& other);

  ~NeighborSearch();

  void Train(const MatType& referenceSet);

  void Train(MatType&& referenceSet);

  void Train(const Tree& referenceTree);

  void Train(Tree&& referenceTree);

  void Search(const MatType& querySet,
              const size_t k,
              arma::Mat<size_t>& neighbors,
              arma::mat& distances);

  void Search(Tree& queryTree,
              const size_t k,
              arma::Mat<size_t>& neighbors,
              arma::mat& distances,
              bool sameSet = false);

  void Search(const size_t k,
              arma::Mat<size_t>& neighbors,
              arma::mat& distances);

  static double EffectiveError(arma::mat& foundDistances,
                               arma::mat& realDistances);

  static double Recall(arma::Mat<size_t>& foundNeighbors,
                       arma::Mat<size_t>& realNeighbors);

  size_t BaseCases() const { return baseCases; }

  size_t Scores() const { return scores; }

  NeighborSearchMode SearchMode() const { return searchMode; }
  NeighborSearchMode& SearchMode() { return searchMode; }

  double Epsilon() const { return epsilon; }
  double& Epsilon() { return epsilon; }

  const MatType& ReferenceSet() const { return *referenceSet; }

  const Tree& ReferenceTree() const { return *referenceTree; }
  Tree& ReferenceTree() { return *referenceTree; }

  template<typename Archive>
  void Serialize(Archive& ar, const unsigned int /* version */);

 private:
  std::vector<size_t> oldFromNewReferences;
  Tree* referenceTree;
  const MatType* referenceSet;

  bool treeOwner;
  bool setOwner;

  NeighborSearchMode searchMode;
  double epsilon;

  MetricType metric;

  size_t baseCases;
  size_t scores;

  bool treeNeedsReset;

  template<typename SortPol>
  friend class TrainVisitor;
}; // class NeighborSearch

} // namespace neighbor
} // namespace mlpack

// Include implementation.
#include "neighbor_search_impl.hpp"

// Include convenience typedefs.
#include "typedef.hpp"

#endif