TreeNearestNeighbors.h

Go to the documentation of this file.

//===========================================================================
/*!
 * 
 *
 * \brief       Efficient Nearest neighbor queries.
 * 
 * 
 *
 * \author      T. Glasmachers
 * \date        2011
 *
 *
 * \par Copyright 1995-2017 Shark Development Team
 * 
 * <BR><HR>
 * This file is part of Shark.
 * <http://shark-ml.org/>
 * 
 * Shark is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published 
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * Shark is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 * 
 * You should have received a copy of the GNU Lesser General Public License
 * along with Shark.  If not, see <http://www.gnu.org/licenses/>.
 *
 */
//===========================================================================

#ifndef SHARK_ALGORITHMS_NEARESTNEIGHBORS_TREENEARESTNEIGHBORS_H
#define SHARK_ALGORITHMS_NEARESTNEIGHBORS_TREENEARESTNEIGHBORS_H


#include <boost/intrusive/rbtree.hpp>
#include <shark/Models/Trees/BinaryTree.h>
#include <shark/Algorithms/NearestNeighbors/AbstractNearestNeighbors.h>
#include <shark/Data/DataView.h>
namespace shark {


///
/// \brief Iterative nearest neighbors query.
///
/// \par
/// The IterativeNNQuery class (Iterative Nearest Neighbor
/// Query) allows the nearest neighbors of a reference point
/// to be queried iteratively. Given the reference point, a
/// query is set up that returns the nearest neighbor first,
/// then the second nearest neighbor, and so on.
/// Thus, nearest neighbor queries are treated in an "online"
/// fashion. The algorithm follows the paper (generalized to
/// arbitrary space-partitioning trees):
///
/// \par
/// Strategies for efficient incremental nearest neighbor search.
/// A. J. Broder. Pattern Recognition 23(1/2), pp 171-178, 1990.
///
/// \par
/// The algorithm is based on traversing a BinaryTree that
/// partitions the space into nested cells. The triangle
/// inequality is applied to exclude cells from the search.
/// Furthermore, candidate points are cached in a queue,
/// such that subsequent queries profit from points that
/// could not be excluded this way, but that did not turn
/// out the be the (current) nearest neighbor.
///
/// \par
/// The tree must have a bucket size of one, but leaf nodes
/// with multiple copies of the same point are allowed.
/// This means that the space partitioning must be carried
/// out to the finest possible scale.
///
/// The Data must be sotred in a random access container. This means that elements
/// have O(1) access time. This is crucial for the performance of the tree lookup.
/// When data is stored in a Data<T>, a View should be chosen as template parameter.
template <class DataContainer>
class IterativeNNQuery
{
public:
    typedef typename DataContainer::value_type value_type;
    typedef BinaryTree<value_type> tree_type;
    typedef AbstractKernelFunction<value_type> kernel_type;
    typedef std::pair<double, std::size_t> result_type;

    /// create a new query
    /// \param  tree    Underlying space-partitioning tree (this is assumed to persist for the lifetime of the query object).
    /// \param  data    Container holding the stored data which is referenced by the tree
    /// \param  point   Point whose nearest neighbors are to be found.
    IterativeNNQuery(tree_type const* tree, DataContainer const& data, value_type const& point)
    : m_data(data)
    , m_reference(point)
    , m_nextIndex(0)
    , mp_trace(NULL)
    , mep_head(NULL)
    , m_squaredRadius(0.0)
    , m_neighbors(0)
    {
        // Initialize the recursion trace: descend to the
        // leaf covering the reference point and queue it.
        // The parent of this leaf becomes the "head".
        mp_trace = new TraceNode(tree, NULL, m_reference);
        TraceNode* tn = mp_trace;
        while (tree->hasChildren())
        {
            tn->createLeftNode(tree, m_data, m_reference);
            tn->createRightNode(tree, m_data, m_reference);
            bool left = tree->isLeft(m_reference);
            tn = (left ? tn->mep_left : tn->mep_right);
            tree = (left ? tree->left() : tree->right());
        }
        mep_head = tn->mep_parent;
        insertIntoQueue((TraceLeaf*)tn);
        m_squaredRadius = mp_trace->squaredRadius(m_reference);
    }

    /// destroy the query object and its internal data structures
    ~IterativeNNQuery() {
        m_queue.clear();
        delete mp_trace;
    }


    /// return the number of neighbors already found
    std::size_t neighbors() const {
        return m_neighbors;
    }

    /// find and return the next nearest neighbor
    result_type next() {
        SHARK_RUNTIME_CHECK(m_neighbors < mp_trace->m_tree->size(), "No more neighbors available");

        assert(! m_queue.empty());

        // Check whether the current node has points
        // left, or whether it should be discarded.
        if (m_neighbors > 0){
            TraceLeaf& q = *m_queue.begin();
            if (m_nextIndex < q.m_tree->size()){
                return getNextPoint(q);
            }
            else
                m_queue.erase(q);
        }
        if (m_queue.empty() || (*m_queue.begin()).m_squaredPtDistance > m_squaredRadius){
            // enqueue more points
            TraceNode* tn = mep_head;
            while (tn != NULL){
                enqueue(tn);
                if (tn->m_status == COMPLETE) mep_head = tn->mep_parent;
                tn = tn->mep_parent;
            }

            // re-compute the radius
            m_squaredRadius = mp_trace->squaredRadius(m_reference);
        }
        m_nextIndex = 0;
        ++m_neighbors;
        return getNextPoint(*m_queue.begin());
    }

    /// return the size of the queue,
    /// which is a measure of the
    /// overhead of the search
    std::size_t queuesize() const{ 
        return m_queue.size();
    }

private:

    /// status of a TraceNode object during the search
    enum Status
    {
        NONE,            // no points of this node have been queued yet
        PARTIAL,         // some of the points of this node have been queued
        COMPLETE,        // all points of this node have been queued
    };

    /// The TraceNode class builds up a tree during
    /// the search. This tree covers only those parts
    /// of the space partirioning tree that need to be
    /// traversed in order to find the next nearest
    /// neighbor.
    class TraceNode
    {
    public:
        /// Constructor
        TraceNode(tree_type const* tree, TraceNode* parent, value_type const& reference)
        : m_tree(tree)
        , m_status(NONE)
        , mep_parent(parent)
        , mep_left(NULL)
        , mep_right(NULL)
        , m_squaredDistance(tree->squaredDistanceLowerBound(reference))
        { }

        /// Destructor
        virtual ~TraceNode()
        {
            if (mep_left != NULL) delete mep_left;
            if (mep_right != NULL) delete mep_right;
        }
        
        void createLeftNode(tree_type const* tree, DataContainer const& data, value_type const& reference){
            if (tree->left()->hasChildren())
                mep_left = new TraceNode(tree->left(), this, reference);
            else
                mep_left = new TraceLeaf(tree->left(), this, data, reference);
        }
        void createRightNode(tree_type const* tree, DataContainer const& data, value_type const& reference){
            if (tree->right()->hasChildren())
                mep_right = new TraceNode(tree->right(), this, reference);
            else
                mep_right = new TraceLeaf(tree->right(), this, data, reference);
        }

        /// Compute the squared distance of the area not
        /// yet covered by the queue to the reference point.
        /// This is also referred to as the squared "radius"
        /// of the area covered by the queue (in fact, it is
        /// the radius of the largest sphere around the
        /// reference point that fits into the covered area).
        double squaredRadius(value_type const& ref) const{
            if (m_status == NONE) return m_squaredDistance;
            else if (m_status == PARTIAL)
            {
                double l = mep_left->squaredRadius(ref);
                double r = mep_right->squaredRadius(ref);
                return std::min(l, r);
            }
            else return 1e100;
        }

        /// node of the tree
        tree_type const* m_tree;

        /// status of the search
        Status m_status;

        /// parent node
        TraceNode* mep_parent;

        /// "left" child
        TraceNode* mep_left;

        /// "right" child
        TraceNode* mep_right;

        /// squared distance of the box to the reference point
        double m_squaredDistance;
    };

    /// hook type for intrusive container
    typedef boost::intrusive::set_base_hook<> HookType;

    /// Leaves of the three have three roles:
    /// (1) they are tree nodes holding exactly one point
    ///     (possibly multiple copies of this point),
    /// (2) they know the distance of their point to the
    ///     reference point,
    /// (3) they can be added to the candidates queue.
    class TraceLeaf : public TraceNode, public HookType
    {
    public:
        /// Constructor
        TraceLeaf(tree_type const* tree, TraceNode* parent, DataContainer const& data, value_type const& ref)
        : TraceNode(tree, parent, ref){
            //check whether the tree uses a differen metric than a linear one.
            if(tree->kernel() != NULL)
                m_squaredPtDistance = tree->kernel()->featureDistanceSqr(data[tree->index(0)], ref);
            else
                m_squaredPtDistance = distanceSqr(data[tree->index(0)], ref);
        }

        /// Destructor
        ~TraceLeaf() { }


        /// Comparison by distance, ties are broken arbitrarily,
        /// but deterministically, by tree node pointer.
        inline bool operator < (TraceLeaf const& rhs) const{
            if (m_squaredPtDistance == rhs.m_squaredPtDistance) 
                return (this->m_tree < rhs.m_tree);
            else
                return (m_squaredPtDistance < rhs.m_squaredPtDistance);
        }

        /// Squared distance of the single point in the leaf to the reference point.
        double m_squaredPtDistance;
    };

    /// insert a point into the queue
    void insertIntoQueue(TraceLeaf* leaf){
        m_queue.insert_unique(*leaf);

        // traverse up the tree, updating the state
        TraceNode* tn = leaf;
        tn->m_status = COMPLETE;
        while (true){
            TraceNode* par = tn->mep_parent;
            if (par == NULL) break;
            if (par->m_status == NONE){
                par->m_status = PARTIAL;
                break;
            }
            else if (par->m_status == PARTIAL){
                if (par->mep_left == tn){
                    if (par->mep_right->m_status == COMPLETE) par->m_status = COMPLETE;
                    else break;
                }
                else{
                    if (par->mep_left->m_status == COMPLETE) par->m_status = COMPLETE;
                    else break;
                }
            }
            tn = par;
        }
    }

    result_type getNextPoint(TraceLeaf const& leaf){
        double dist = std::sqrt(leaf.m_squaredPtDistance);
        std::size_t index = leaf.m_tree->index(m_nextIndex);
        ++m_nextIndex;
        return std::make_pair(dist,index);
    }

    /// Recursively descend the node and enqueue
    /// all points in cells intersecting the
    /// current bounding sphere.
    void enqueue(TraceNode* tn){
        // check whether this node needs to be enqueued
        if (tn->m_status == COMPLETE) return;
        if (! m_queue.empty() && tn->m_squaredDistance >= (*m_queue.begin()).m_squaredPtDistance) return;

        const tree_type* tree = tn->m_tree;
        if (tree->hasChildren()){
            // extend the tree at need
            if (tn->mep_left == NULL){
                tn->createLeftNode(tree,m_data,m_reference);
            }
            if (tn->mep_right == NULL){
                tn->createRightNode(tree,m_data,m_reference);
            }

            // first descend into the closer sub-tree
            if (tree->isLeft(m_reference))
            {
                // left first
                enqueue(tn->mep_left);
                enqueue(tn->mep_right);
            }
            else
            {
                // right first
                enqueue(tn->mep_right);
                enqueue(tn->mep_left);
            }
        }
        else
        {
            TraceLeaf* leaf = (TraceLeaf*)tn;
            insertIntoQueue(leaf);
        }
    }

    /// the queue is a self-balancing tree of sorted entries
    typedef boost::intrusive::rbtree<TraceLeaf> QueueType;


    ///\brief Datastorage to lookup the points referenced by the space partitioning tree.
    DataContainer const& m_data;

    /// reference point for this query
    value_type m_reference;

    /// queue of candidates
    QueueType m_queue;

    /// index of the next not yet returned element
    /// of the current leaf.
    std::size_t m_nextIndex;

    /// recursion trace tree
    TraceNode* mp_trace;

    /// "head" of the trace tree. This is the
    /// node containing the reference point,
    /// but so high up in the tree that it is
    /// not fully queued yet.
    TraceNode* mep_head;

    /// squared radius of the "covered" area
    double m_squaredRadius;

    /// number of neighbors already returned
    std::size_t m_neighbors;
};


///\brief Nearest Neighbors implementation using binary trees
///
/// Returns the labels and distances of the k nearest neighbors of a point.
template<class InputType, class LabelType>
class TreeNearestNeighbors:public AbstractNearestNeighbors<InputType,LabelType>
{
private:
    typedef AbstractNearestNeighbors<InputType,LabelType> base_type;

public:
    typedef LabeledData<InputType, LabelType> Dataset;
    typedef BinaryTree<InputType> Tree;
    typedef typename base_type::DistancePair DistancePair;
    typedef typename Batch<InputType>::type BatchInputType;

    TreeNearestNeighbors(Dataset const& dataset, Tree const* tree)
    : m_dataset(dataset)
    , m_inputs(dataset.inputs())
    , m_labels(dataset.labels())
    , mep_tree(tree)
    {
        this->m_inputShape = dataset.inputShape();
    }

    ///\brief returns the k nearest neighbors of the point
    std::vector<DistancePair> getNeighbors(BatchInputType const& patterns, std::size_t k)const{
        std::size_t numPoints = batchSize(patterns);
        std::vector<DistancePair> results(k*numPoints);
        for(std::size_t p = 0; p != numPoints; ++p){
            IterativeNNQuery<DataView<Data<InputType> const> > query(mep_tree, m_inputs, row(patterns, p));
            //find the neighbors using the queries
            for(std::size_t i = 0; i != k; ++i){
                typename IterativeNNQuery<DataView<Data<InputType> const> >::result_type result = query.next();
                results[i+p*k].key=result.first;
                results[i+p*k].value= m_labels[result.second]; 
            }
        }
        return results;
    }

    LabeledData<InputType,LabelType>const& dataset()const {
        return m_dataset;
    }

private:
    Dataset const& m_dataset;
    DataView<Data<InputType> const> m_inputs;
    DataView<Data<LabelType> const> m_labels;
    Tree const* mep_tree;
    
};


}
#endif