itkKdTree.h

Go to the documentation of this file.

/*=========================================================================
 *
 *  Copyright Insight Software Consortium
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/
#ifndef itkKdTree_h
#define itkKdTree_h

#include <queue>
#include <vector>

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

#include "itkSubsample.h"

#include "itkEuclideanDistanceMetric.h"

namespace itk
{
namespace Statistics
{
template<typename TSample>

struct ITK_TEMPLATE_EXPORT KdTreeNode
  {
  using Self = KdTreeNode<TSample>;

  using MeasurementType = typename TSample::MeasurementType;

  using CentroidType = Array<double>;

  using InstanceIdentifier = typename TSample::InstanceIdentifier;

  virtual bool IsTerminal() const = 0;

  virtual void GetParameters( unsigned int &, MeasurementType & ) 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 & ) = 0;

  virtual void GetCentroid( CentroidType & ) = 0;

  virtual InstanceIdentifier GetInstanceIdentifier( InstanceIdentifier ) const = 0;

  virtual void AddInstanceIdentifier( InstanceIdentifier ) = 0;

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

template<typename TSample>

struct ITK_TEMPLATE_EXPORT KdTreeNonterminalNode:public KdTreeNode<TSample>
  {
  using Superclass = KdTreeNode<TSample>;
  using MeasurementType = typename Superclass::MeasurementType;
  using CentroidType = typename Superclass::CentroidType;
  using InstanceIdentifier = typename Superclass::InstanceIdentifier;

  KdTreeNonterminalNode( unsigned int, MeasurementType, Superclass *,
    Superclass * );

  ~KdTreeNonterminalNode() override = default;

  bool IsTerminal() const override
  {
    return false;
  }

  void GetParameters( unsigned int &, MeasurementType & ) const override;

  Superclass * Left() override
  {
    return m_Left;
  }

  Superclass * Right() override
  {
    return m_Right;
  }

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

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

  unsigned int Size() const override
  {
    return 0;
  }

  void GetWeightedCentroid( CentroidType & ) override {}

  void GetCentroid( CentroidType & ) override {}

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

  void AddInstanceIdentifier( InstanceIdentifier valueId ) override
  {
    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<typename TSample>
struct ITK_TEMPLATE_EXPORT KdTreeWeightedCentroidNonterminalNode:public KdTreeNode<TSample>
  {
  using Superclass = KdTreeNode<TSample>;
  using MeasurementType = typename Superclass::MeasurementType;
  using CentroidType = typename Superclass::CentroidType;
  using InstanceIdentifier = typename Superclass::InstanceIdentifier;
  using MeasurementVectorSizeType = typename TSample::MeasurementVectorSizeType;

  KdTreeWeightedCentroidNonterminalNode( unsigned int, MeasurementType,
    Superclass *, Superclass *, CentroidType &, unsigned int );

  ~KdTreeWeightedCentroidNonterminalNode() override = default;

  bool IsTerminal() const override
  {
    return false;
  }

  void GetParameters( unsigned int &, MeasurementType & ) const override;

  MeasurementVectorSizeType GetMeasurementVectorSize() const
  {
    return m_MeasurementVectorSize;
  }

  Superclass * Left() override
  {
    return m_Left;
  }

  Superclass * Right() override
  {
    return m_Right;
  }

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

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

  unsigned int Size() const override
  {
    return m_Size;
  }

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

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

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

  void AddInstanceIdentifier(InstanceIdentifier valueId) override
  {
    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<typename TSample>
struct ITK_TEMPLATE_EXPORT KdTreeTerminalNode:public KdTreeNode<TSample>
  {
  using Superclass = KdTreeNode<TSample>;
  using MeasurementType = typename Superclass::MeasurementType;
  using CentroidType = typename Superclass::CentroidType;
  using InstanceIdentifier = typename Superclass::InstanceIdentifier;

  KdTreeTerminalNode() = default;

  ~KdTreeTerminalNode() override
  {
    this->m_InstanceIdentifiers.clear();
  }

  bool IsTerminal() const override
  {
    return true;
  }

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

  Superclass * Left() override
  {
    return nullptr;
  }

  Superclass * Right() override
  {
    return nullptr;
  }

  const Superclass * Left() const override
  {
    return nullptr;
  }

  const Superclass * Right() const override
  {
    return nullptr;
  }

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

  void GetWeightedCentroid( CentroidType & ) override {}

  void GetCentroid( CentroidType & ) override {}

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

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

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

template<typename TSample>
class ITK_TEMPLATE_EXPORT KdTree:public Object
{
public:
  ITK_DISALLOW_COPY_AND_ASSIGN(KdTree);

  using Self = KdTree;
  using Superclass = Object;
  using Pointer = SmartPointer< Self >;
  using ConstPointer = SmartPointer< const Self >;

  itkTypeMacro(KdTree, Object);

  itkNewMacro(Self);

  using SampleType = TSample;
  using MeasurementVectorType = typename TSample::MeasurementVectorType;
  using MeasurementType = typename TSample::MeasurementType;
  using InstanceIdentifier = typename TSample::InstanceIdentifier;
  using AbsoluteFrequencyType = typename TSample::AbsoluteFrequencyType;

  using MeasurementVectorSizeType = unsigned int;

  itkGetConstMacro( MeasurementVectorSize, MeasurementVectorSizeType );

  using DistanceMetricType = EuclideanDistanceMetric< MeasurementVectorType >;

  using KdTreeNodeType = KdTreeNode<TSample>;

  using NeighborType = std::pair< InstanceIdentifier, double >;

  using InstanceIdentifierVectorType = std::vector< InstanceIdentifier >;

  class NearestNeighbors
  {
  public:
    NearestNeighbors( std::vector<double> & cache_vector) : m_FarthestNeighborIndex(0), m_Distances(cache_vector)  {}

    ~NearestNeighbors() = default;

    void resize( unsigned int k )
    {
      m_Identifiers.clear();
      m_Identifiers.resize( k, NumericTraits< IdentifierType >::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 auto 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];
    }

  private:
    NearestNeighbors() = delete;
    unsigned int m_FarthestNeighborIndex;

    InstanceIdentifierVectorType m_Identifiers;

    std::vector<double> & m_Distances;
  };

  void SetBucketSize( unsigned int );

  void SetSample( const TSample * );

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

  SizeValueType Size() const
  {
    return m_Sample->Size();
  }

  KdTreeNodeType * GetEmptyTerminalNode()
  {
    return m_EmptyTerminalNode;
  }

  void SetRoot(KdTreeNodeType *root)
  {
    if ( this->m_Root )
      {
      this->DeleteNode( this->m_Root );
      }
    this->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 &, unsigned int,
              InstanceIdentifierVectorType & ) const;

  void Search( const MeasurementVectorType &, unsigned int,
    InstanceIdentifierVectorType &, std::vector<double> & ) const;

  void Search( const MeasurementVectorType &, double,
    InstanceIdentifierVectorType & ) const;

  bool BallWithinBounds( const MeasurementVectorType &,
    MeasurementVectorType &, MeasurementVectorType &, double ) const;

  bool BoundsOverlapBall( const MeasurementVectorType &,
    MeasurementVectorType &, MeasurementVectorType &, double) const;

  void DeleteNode( KdTreeNodeType * );

  void PrintTree( std::ostream & ) const;

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

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

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

  using Iterator = typename TSample::Iterator;
  using ConstIterator = typename TSample::ConstIterator;

protected:
  KdTree();

  ~KdTree() override;

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

  int NearestNeighborSearchLoop( const KdTreeNodeType *,
    const MeasurementVectorType &, MeasurementVectorType &,
    MeasurementVectorType &, NearestNeighbors & ) const;

  int SearchLoop( const KdTreeNodeType *, const MeasurementVectorType &,
    double, MeasurementVectorType &, MeasurementVectorType &,
    InstanceIdentifierVectorType & ) const;

private:
  const TSample *m_Sample;

  int m_BucketSize;

  KdTreeNodeType *m_Root;

  KdTreeNodeType *m_EmptyTerminalNode;

  typename DistanceMetricType::Pointer m_DistanceMetric;

  MeasurementVectorSizeType m_MeasurementVectorSize;
};  // end of class
} // end of namespace Statistics
} // end of namespace itk

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

#endif