Review/Statistics/itkKdTree.h

Go to the documentation of this file.

/*=========================================================================

  Program:   Insight Segmentation & Registration Toolkit
  Module:    $RCSfile: itkKdTree.h,v $
  Language:  C++
  Date:      $Date: 2009-05-02 05:43:56 $
  Version:   $Revision: 1.1 $

  Copyright (c) Insight Software Consortium. All rights reserved.
  See ITKCopyright.txt or http://www.itk.org/HTML/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notices for more information.

=========================================================================*/
#ifndef __itkKdTree_h
#define __itkKdTree_h

#include <queue>
#include <vector>

#include "itkMacro.h"
#include "itkPoint.h"
#include "itkSize.h"
#include "itkObject.h"
#include "itkNumericTraits.h"
#include "itkArray.h"

#include "itkSample.h"
#include "itkSubsample.h"

#include "itkEuclideanDistanceMetric.h"

namespace itk {
namespace Statistics  {

template< class TSample >
struct KdTreeNode
{
  typedef KdTreeNode< TSample> Self;

  typedef typename TSample::MeasurementType MeasurementType;

  typedef Array< double > CentroidType;

  typedef typename TSample::InstanceIdentifier InstanceIdentifier;

  virtual bool IsTerminal() const = 0;

  virtual void GetParameters(unsigned int &partitionDimension,
                             MeasurementType &partitionValue) const = 0;

  virtual       Self* Left()       = 0;
  virtual const Self* Left() const = 0;

  virtual       Self* Right()       = 0;
  virtual const Self* Right() const = 0;

  virtual unsigned int Size() const = 0;

  virtual void GetWeightedCentroid(CentroidType &centroid) = 0;

  virtual void GetCentroid(CentroidType &centroid) = 0;

  virtual InstanceIdentifier GetInstanceIdentifier(size_t index) const = 0;

  virtual void AddInstanceIdentifier(InstanceIdentifier id) = 0;

  virtual ~KdTreeNode() {}; // needed to subclasses will actually be deleted
}; // end of class

template< class TSample >
struct KdTreeNonterminalNode: public KdTreeNode< TSample >
{
  typedef KdTreeNode< TSample >                       Superclass;
  typedef typename Superclass::MeasurementType        MeasurementType;
  typedef typename Superclass::CentroidType           CentroidType;
  typedef typename Superclass::InstanceIdentifier     InstanceIdentifier;

  KdTreeNonterminalNode(unsigned int partitionDimension,
                        MeasurementType partitionValue,
                        Superclass* left,
                        Superclass* right);

  virtual ~KdTreeNonterminalNode() {}

  virtual bool IsTerminal() const
    { 
    return false;
    }

  void GetParameters(unsigned int &partitionDimension,
                     MeasurementType &partitionValue) const;

  Superclass* Left()
    { 
    return m_Left;
    }

  Superclass* Right()
    {
    return m_Right;
    }

  const Superclass* Left() const
    {
    return m_Left;
    }

  const Superclass* Right() const
    {
    return m_Right; 
    }

  unsigned int Size() const
    { 
    return 0;
    }

  void GetWeightedCentroid( CentroidType & ) 
    { 
    /* do nothing */
    }

  void GetCentroid( CentroidType & ) 
    { 
    /* do nothing */
    }

  // Returns the identifier of the only MeasurementVector associated with 
  // this node in the tree. This MeasurementVector will be used later during
  // the distance computation when querying the tree.
  InstanceIdentifier GetInstanceIdentifier(size_t) const
  { return this->m_InstanceIdentifier; }

  void AddInstanceIdentifier(InstanceIdentifier valueId) 
  { this->m_InstanceIdentifier = valueId; }

private:
  unsigned int          m_PartitionDimension;
  MeasurementType       m_PartitionValue;
  InstanceIdentifier    m_InstanceIdentifier;
  Superclass* m_Left;
  Superclass* m_Right;
}; // end of class

template< class TSample >
struct KdTreeWeightedCentroidNonterminalNode: public KdTreeNode< TSample >
{
  typedef KdTreeNode< TSample >                           Superclass;
  typedef typename Superclass::MeasurementType            MeasurementType;
  typedef typename Superclass::CentroidType               CentroidType;
  typedef typename Superclass::InstanceIdentifier         InstanceIdentifier;
  typedef typename TSample::MeasurementVectorSizeType     MeasurementVectorSizeType;

  KdTreeWeightedCentroidNonterminalNode(unsigned int partitionDimension,
                                         MeasurementType partitionValue,
                                         Superclass* left,
                                         Superclass* right,
                                         CentroidType &centroid,
                                         unsigned int size);
  virtual ~KdTreeWeightedCentroidNonterminalNode() {}

  virtual bool IsTerminal() const
    {
    return false;
    }

  void GetParameters(unsigned int &partitionDimension,
                     MeasurementType &partitionValue) const;

  MeasurementVectorSizeType GetMeasurementVectorSize() const
    {
    return m_MeasurementVectorSize;
    }

  Superclass* Left()
    {
    return m_Left;
    }

  Superclass* Right()
    {
    return m_Right;
    }


  const Superclass* Left() const
    {
    return m_Left;
    }

  const Superclass* Right() const
    {
    return m_Right;
    }

  unsigned int Size() const
    {
    return m_Size;
    }

  void GetWeightedCentroid(CentroidType &centroid)
    {
    centroid = m_WeightedCentroid;
    }

  void GetCentroid(CentroidType &centroid)
    {
    centroid = m_Centroid;
    }

  InstanceIdentifier GetInstanceIdentifier(size_t) const
    {
    return this->m_InstanceIdentifier;
    }

  void AddInstanceIdentifier(InstanceIdentifier valueId) 
    {
    this->m_InstanceIdentifier = valueId;
    }

private:
  MeasurementVectorSizeType   m_MeasurementVectorSize;
  unsigned int                m_PartitionDimension;
  MeasurementType             m_PartitionValue;
  CentroidType                m_WeightedCentroid;
  CentroidType                m_Centroid;
  InstanceIdentifier          m_InstanceIdentifier;
  unsigned int                m_Size;
  Superclass*                 m_Left;
  Superclass*                 m_Right;
}; // end of class


template< class TSample >
struct KdTreeTerminalNode: public KdTreeNode< TSample >
{
  typedef KdTreeNode< TSample >                   Superclass;
  typedef typename Superclass::MeasurementType    MeasurementType;
  typedef typename Superclass::CentroidType       CentroidType;
  typedef typename Superclass::InstanceIdentifier InstanceIdentifier;

  KdTreeTerminalNode() {}

  virtual ~KdTreeTerminalNode() {}

  bool IsTerminal() const
    {
    return true;
    }

  void GetParameters(unsigned int &,
                     MeasurementType &) const {}

  Superclass* Left()
    {
    return 0;
    }

  Superclass* Right()
    {
    return 0;
    }

  const Superclass* Left() const
    {
    return 0;
    }

  const Superclass* Right() const
    {
    return 0;
    }

  unsigned int Size() const
    {
    return static_cast<unsigned int>( m_InstanceIdentifiers.size() );
    }

  void GetWeightedCentroid(CentroidType &)
    { 
    /* do nothing */ 
    }

  void GetCentroid(CentroidType &)
    {
    /* do nothing */ 
    }

  InstanceIdentifier GetInstanceIdentifier(size_t index) const
    { 
    return m_InstanceIdentifiers[index];
    }

  void AddInstanceIdentifier(InstanceIdentifier id)
    {
    m_InstanceIdentifiers.push_back(id);
    }

private:
  std::vector< InstanceIdentifier > m_InstanceIdentifiers;
}; // end of class

template < class TSample >
class ITK_EXPORT KdTree : public Object
{
public:
  typedef KdTree                                    Self;
  typedef Object                                    Superclass;
  typedef SmartPointer<Self>                        Pointer;
  typedef SmartPointer<const Self>                  ConstPointer;

  itkTypeMacro(KdTree, Object);

  itkNewMacro(Self);

  typedef TSample                                     SampleType;
  typedef typename TSample::MeasurementVectorType     MeasurementVectorType;
  typedef typename TSample::MeasurementType           MeasurementType;
  typedef typename TSample::InstanceIdentifier        InstanceIdentifier;
  typedef typename TSample::AbsoluteFrequencyType     AbsoluteFrequencyType;

  typedef unsigned int                    MeasurementVectorSizeType;

  itkGetConstMacro( MeasurementVectorSize, MeasurementVectorSizeType );

  typedef EuclideanDistanceMetric< MeasurementVectorType > DistanceMetricType;

  typedef KdTreeNode< TSample > KdTreeNodeType;

  typedef std::pair< InstanceIdentifier, double > NeighborType;

  typedef std::vector< InstanceIdentifier > InstanceIdentifierVectorType;

  class NearestNeighbors {
  public:
    NearestNeighbors() {}

    ~NearestNeighbors() {}

    void resize(unsigned int k)
      {
      m_Identifiers.clear();
      m_Identifiers.resize(k, NumericTraits< unsigned long >::max());
      m_Distances.clear();
      m_Distances.resize(k, NumericTraits< double >::max());
      m_FarthestNeighborIndex = 0;
      }

    double GetLargestDistance()
      {
      return m_Distances[m_FarthestNeighborIndex];
      }

    void ReplaceFarthestNeighbor(InstanceIdentifier id, double distance)
      {
      m_Identifiers[m_FarthestNeighborIndex] = id;
      m_Distances[m_FarthestNeighborIndex] = distance;
      double farthestDistance = NumericTraits< double >::min();
      const unsigned int size = static_cast<unsigned int>( m_Distances.size() );
      for ( unsigned int i = 0; i < size; i++ )
        {
        if ( m_Distances[i] > farthestDistance )
          {
          farthestDistance = m_Distances[i];
          m_FarthestNeighborIndex = i;
          }
        }
      }

    const InstanceIdentifierVectorType & GetNeighbors() const
      {
      return m_Identifiers;
      }

    InstanceIdentifier GetNeighbor(unsigned int index) const
      {
      return m_Identifiers[index];
      }

    const std::vector< double >& GetDistances() const
      { 
      return m_Distances;
      }

  private:
    unsigned int m_FarthestNeighborIndex;

    InstanceIdentifierVectorType m_Identifiers;

    std::vector< double > m_Distances;
  };

  void SetBucketSize(unsigned int size);

  void SetSample(const TSample* sample);

  const TSample* GetSample() const
    {
    return m_Sample;
    }

  unsigned long Size() const
    {
    return m_Sample->Size();
    }

  KdTreeNodeType* GetEmptyTerminalNode()
    {
    return m_EmptyTerminalNode;
    }

  void SetRoot(KdTreeNodeType* root)
    {
    m_Root = root;
    }

  KdTreeNodeType* GetRoot()
    {
    return m_Root;
    }

  const MeasurementVectorType & GetMeasurementVector(InstanceIdentifier id) const
    {
    return m_Sample->GetMeasurementVector(id); 
    }

  AbsoluteFrequencyType GetFrequency(InstanceIdentifier id) const
    {
    return m_Sample->GetFrequency( id ); 
    }

  DistanceMetricType* GetDistanceMetric()
    {
    return m_DistanceMetric.GetPointer();
    }

  void Search(const MeasurementVectorType &query,
              unsigned int numberOfNeighborsRequested,
              InstanceIdentifierVectorType& result) const;

  void Search(const MeasurementVectorType &query,
              double radius,
              InstanceIdentifierVectorType& result) const;

  int GetNumberOfVisits() const
    {
    return m_NumberOfVisits; 
    }

  bool BallWithinBounds(const MeasurementVectorType &query,
                        MeasurementVectorType &lowerBound,
                        MeasurementVectorType &upperBound,
                        double radius) const;

  bool BoundsOverlapBall(const MeasurementVectorType &query,
                         MeasurementVectorType &lowerBound,
                         MeasurementVectorType &upperBound,
                         double radius) const;

  void DeleteNode(KdTreeNodeType *node);

  void PrintTree( std::ostream & os ) const;

  void PrintTree(KdTreeNodeType *node, unsigned int level,
                 unsigned int activeDimension,
                 std::ostream & os = std::cout ) const;

  void PlotTree( std::ostream & os ) const;

  void PlotTree(KdTreeNodeType *node, std::ostream & os = std::cout ) const;


  typedef typename TSample::Iterator      Iterator;
  typedef typename TSample::ConstIterator ConstIterator;

  Iterator Begin()
    {
    typename TSample::ConstIterator iter = m_Sample->Begin();
    return iter;
    }

  Iterator End()
    {
    Iterator iter = m_Sample->End();
    return iter;
    }

  ConstIterator Begin() const
    {
    typename TSample::ConstIterator iter = m_Sample->Begin();
    return iter;
    }

  ConstIterator End() const
    {
    ConstIterator iter = m_Sample->End();
    return iter;
    } 

protected:
  KdTree();

  virtual ~KdTree();

  void PrintSelf(std::ostream& os, Indent indent) const;

  int NearestNeighborSearchLoop(const KdTreeNodeType* node,
                                const MeasurementVectorType &query,
                                MeasurementVectorType &lowerBound,
                                MeasurementVectorType &upperBound) const;

  int SearchLoop(const KdTreeNodeType* node, const MeasurementVectorType &query,
                 MeasurementVectorType &lowerBound,
                 MeasurementVectorType &upperBound) const;
private:
  KdTree(const Self&); //purposely not implemented
  void operator=(const Self&); //purposely not implemented

  const TSample* m_Sample;

  int m_BucketSize;

  KdTreeNodeType* m_Root;

  KdTreeNodeType* m_EmptyTerminalNode;

  typename DistanceMetricType::Pointer m_DistanceMetric;

  mutable bool m_IsNearestNeighborSearch;

  mutable double m_SearchRadius;

  mutable InstanceIdentifierVectorType m_Neighbors;

  mutable NearestNeighbors m_NearestNeighbors;

  mutable MeasurementVectorType m_LowerBound;

  mutable MeasurementVectorType m_UpperBound;

  mutable int m_NumberOfVisits;

  mutable bool m_StopSearch;

  mutable NeighborType m_TempNeighbor;

  MeasurementVectorSizeType m_MeasurementVectorSize;
}; // end of class

} // end of namespace Statistics
} // end of namespace itk

#ifndef ITK_MANUAL_INSTANTIATION
#include "itkKdTree.txx"
#endif

#endif