From cb3af4bdfd0bfb9c42730e5f63a30d8c7df977c3 Mon Sep 17 00:00:00 2001 From: Ryan Curtin Date: Wed, 9 May 2012 20:26:02 +0000 Subject: [PATCH] Add implementation of rules for dual-tree search. --- .../neighbor_search/neighbor_search_rules.hpp | 13 ++ .../neighbor_search_rules_impl.hpp | 196 ++++++++++++++++++ 2 files changed, 209 insertions(+) diff --git a/src/mlpack/methods/neighbor_search/neighbor_search_rules.hpp b/src/mlpack/methods/neighbor_search/neighbor_search_rules.hpp index 03b3858d41..317cbaa75f 100644 --- a/src/mlpack/methods/neighbor_search/neighbor_search_rules.hpp +++ b/src/mlpack/methods/neighbor_search/neighbor_search_rules.hpp @@ -26,6 +26,19 @@ class NeighborSearchRules // For single-tree traversal. bool CanPrune(const size_t queryIndex, TreeType& referenceNode); + // For dual-tree traversal. + bool CanPrune(TreeType& queryNode, TreeType& referenceNode); + + // Get the order of points to recurse to. + void RecursionOrder(TreeType& queryNode, + TreeType& referenceNode, + arma::Mat& recursionOrder, + bool& queryRecurse, + bool& referenceRecurse); + + // Update bounds. Needs a better name. + void UpdateAfterRecursion(TreeType& queryNode, TreeType& referenceNode); + private: //! The reference set. const arma::mat& referenceSet; diff --git a/src/mlpack/methods/neighbor_search/neighbor_search_rules_impl.hpp b/src/mlpack/methods/neighbor_search/neighbor_search_rules_impl.hpp index ed14a602b3..87bffbdac8 100644 --- a/src/mlpack/methods/neighbor_search/neighbor_search_rules_impl.hpp +++ b/src/mlpack/methods/neighbor_search/neighbor_search_rules_impl.hpp @@ -69,6 +69,202 @@ inline bool NeighborSearchRules::CanPrune( return true; // There cannot be anything better in this node. So prune it. } +template +inline bool NeighborSearchRules::CanPrune( + TreeType& queryNode, + TreeType& referenceNode) +{ + const double distance = SortPolicy::BestNodeToNodeDistance( + &queryNode, &referenceNode); + const double bestDistance = queryNode.Stat().Bound(); + + if (SortPolicy::IsBetter(distance, bestDistance)) + return false; // Can't prune. + else + return true; +} + +// Return the order in which we should recurse. +template +inline void NeighborSearchRules< + SortPolicy, + MetricType, + TreeType>:: +RecursionOrder(TreeType& queryNode, + TreeType& referenceNode, + arma::Mat& recursionOrder, + bool& queryRecurse, + bool& referenceRecurse) +{ + queryRecurse = !(queryNode.IsLeaf()); + referenceRecurse = !(referenceNode.IsLeaf()); + + if (queryRecurse && !referenceRecurse) + { + // We only need to recurse into the query children. Therefore, the elements + // in row 1 can be ignored. + recursionOrder.set_size(2, queryNode.NumChildren()); + arma::vec recursionDistances(queryNode.NumChildren()); + recursionDistances.fill(SortPolicy::WorstDistance()); + size_t children = 0; // Number of children to recurse to. + + for (size_t i = 0; i < queryNode.NumChildren(); ++i) + { + double distance = SortPolicy::BestNodeToNodeDistance(&queryNode.Child(i), + &referenceNode); + + // Find where to insert. + size_t insertPosition; + for (insertPosition = 0; insertPosition < children; ++insertPosition) + if (SortPolicy::IsBetter(distance, recursionDistances[insertPosition])) + break; + + // Now perform the actual insertion. + if ((children - insertPosition) > 0) + { + memmove(recursionDistances.memptr() + insertPosition + 1, + recursionDistances.memptr() + insertPosition, + sizeof(double) * (children - insertPosition)); + memmove(recursionOrder.memptr() + (insertPosition + 1) * 2, + recursionOrder.memptr() + (insertPosition * 2), + sizeof(size_t) * (children - insertPosition) * 2); + } + + // Insert. + recursionDistances[insertPosition] = distance; + recursionOrder(0, insertPosition) = i; + ++children; + } + + // Strip extra columns. + if (children < queryNode.NumChildren()) + recursionOrder.shed_cols(children, queryNode.NumChildren() - 1); + } + else if (!queryRecurse && referenceRecurse) + { + // We only need to recurse into the reference children. Therefore, the + // elements in row 0 can be ignored. + recursionOrder.set_size(2, referenceNode.NumChildren()); + arma::vec recursionDistances(referenceNode.NumChildren()); + recursionDistances.fill(SortPolicy::WorstDistance()); + size_t children = 0; // Number of children to recurse into. + + for (size_t i = 0; i < referenceNode.NumChildren(); ++i) + { + double distance = SortPolicy::BestNodeToNodeDistance(&queryNode, + &referenceNode.Child(i)); + + // Find where to insert. + size_t insertPosition; + for (insertPosition = 0; insertPosition < children; ++insertPosition) + if (SortPolicy::IsBetter(distance, recursionDistances[insertPosition])) + break; + + // Now perform the actual insertion. + if ((children - insertPosition) > 0) + { + memmove(recursionDistances.memptr() + insertPosition + 1, + recursionDistances.memptr() + insertPosition, + sizeof(double) * (children - insertPosition)); + memmove(recursionOrder.memptr() + (insertPosition + 1) * 2, + recursionOrder.memptr() + (insertPosition * 2), + sizeof(size_t) * (children - insertPosition) * 2); + } + + // Insert. + recursionDistances[insertPosition] = distance; + recursionOrder(1, insertPosition) = i; + ++children; + } + + // Strip extra columns. + if (children < referenceNode.NumChildren()) + recursionOrder.shed_cols(children, referenceNode.NumChildren() - 1); + } + else if (queryRecurse && referenceRecurse) + { + // We need to recurse into both children. + const size_t maxChildren = referenceNode.NumChildren() * + queryNode.NumChildren(); + recursionOrder.set_size(2, maxChildren); + arma::vec recursionDistances(maxChildren); + recursionDistances.fill(SortPolicy::WorstDistance()); + size_t children = 0; // Number of children to recurse into. + + for (size_t i = 0; i < queryNode.NumChildren(); ++i) + { + // Check if we should even continue this direction. + if (CanPrune(queryNode.Child(i), referenceNode)) + continue; // Don't go this way. + + for (size_t j = 0; j < referenceNode.NumChildren(); ++j) + { + double distance = SortPolicy::BestNodeToNodeDistance( + &queryNode.Child(i), &referenceNode.Child(j)); + + // Find where to insert. + size_t insertPosition; + for (insertPosition = 0; insertPosition < children; ++insertPosition) + if (SortPolicy::IsBetter(distance, + recursionDistances[insertPosition])) + break; + + // Move things to prepare for insertion. + if ((children - insertPosition) > 0) + { + memmove(recursionDistances.memptr() + insertPosition + 1, + recursionDistances.memptr() + insertPosition, + sizeof(double) * (children - insertPosition)); + memmove(recursionOrder.memptr() + (insertPosition + 1) * 2, + recursionOrder.memptr() + (insertPosition * 2), + sizeof(size_t) * (children - insertPosition) * 2); + } + + // Insert. + recursionDistances[insertPosition] = distance; + recursionOrder(0, insertPosition) = i; + recursionOrder(1, insertPosition) = j; + ++children; + } + } + + // Strip extra columns. + if (children < maxChildren) + recursionOrder.shed_cols(children, maxChildren - 1); + } +} + +template +void NeighborSearchRules< + SortPolicy, + MetricType, + TreeType>:: +UpdateAfterRecursion(TreeType& queryNode, TreeType& /* referenceNode */) +{ + // Find the worst distance that the children found (including any points), and + // update the bound accordingly. + double worstDistance = SortPolicy::BestDistance(); + + // First look through children nodes. + for (size_t i = 0; i < queryNode.NumChildren(); ++i) + { + if (SortPolicy::IsBetter(worstDistance, queryNode.Child(i).Stat().Bound())) + worstDistance = queryNode.Child(i).Stat().Bound(); + } + + // Now look through children points. + for (size_t i = 0; i < queryNode.NumPoints(); ++i) + { + if (SortPolicy::IsBetter(worstDistance, + distances(distances.n_rows - 1, queryNode.Point(i)))) + worstDistance = distances(distances.n_rows - 1, queryNode.Point(i)); + } + + // Take the worst distance from all of these, and update our bound to reflect + // that. + queryNode.Stat().Bound() = worstDistance; +} + /** * Helper function to insert a point into the neighbors and distances matrices. *