itkKdTree.h

Go to the documentation of this file.

/*=========================================================================
 *
 *  Copyright NumFOCUS
 *
 *  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
 *
 *         https://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::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;

  [[nodiscard]] virtual bool
  IsTerminal() const = 0;

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

  virtual Self *
  Left() = 0;

  [[nodiscard]] virtual const Self *
  Left() const = 0;

  virtual Self *
  Right() = 0;

  [[nodiscard]] virtual const Self *
  Right() const = 0;

  [[nodiscard]] virtual unsigned int
  Size() const = 0;

  virtual void
  GetWeightedCentroid(CentroidType &) = 0;

  virtual void
  GetCentroid(CentroidType &) = 0;

  [[nodiscard]] 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 typename Superclass::MeasurementType;
  using typename Superclass::CentroidType;
  using typename Superclass::InstanceIdentifier;
 
  KdTreeNonterminalNode(unsigned int, MeasurementType, Superclass *, Superclass *);
 
  ~KdTreeNonterminalNode() override = default;
 
  [[nodiscard]] 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;
  }

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

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

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

  void
  GetWeightedCentroid(CentroidType &) override
  {}

  void
  GetCentroid(CentroidType &) override
  {}

  [[nodiscard]] 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 typename Superclass::MeasurementType;
  using typename Superclass::CentroidType;
  using typename Superclass::InstanceIdentifier;
  using MeasurementVectorSizeType = typename TSample::MeasurementVectorSizeType;
 
  KdTreeWeightedCentroidNonterminalNode(unsigned int,
                                        MeasurementType,
                                        Superclass *,
                                        Superclass *,
                                        CentroidType &,
                                        unsigned int);
 
  ~KdTreeWeightedCentroidNonterminalNode() override = default;

  [[nodiscard]] bool
  IsTerminal() const override
  {
    return false;
  }

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

  [[nodiscard]] MeasurementVectorSizeType
  GetMeasurementVectorSize() const
  {
    return m_MeasurementVectorSize;
  }

  Superclass *
  Left() override
  {
    return m_Left;
  }

  Superclass *
  Right() override
  {
    return m_Right;
  }

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

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

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

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

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

  [[nodiscard]] 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 typename Superclass::MeasurementType;
  using typename Superclass::CentroidType;
  using typename Superclass::InstanceIdentifier;
 
  KdTreeTerminalNode() = default;
 
  ~KdTreeTerminalNode() override { this->m_InstanceIdentifiers.clear(); }

  [[nodiscard]] bool
  IsTerminal() const override
  {
    return true;
  }

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

  Superclass *
  Left() override
  {
    return nullptr;
  }

  Superclass *
  Right() override
  {
    return nullptr;
  }

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

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

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

  void
  GetWeightedCentroid(CentroidType &) override
  {}

  void
  GetCentroid(CentroidType &) override
  {}

  [[nodiscard]] 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_MOVE(KdTree);

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

  itkOverrideGetNameOfClassMacro(KdTree);

  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() = delete;

    ~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;
        }
      }
    }

    [[nodiscard]] const InstanceIdentifierVectorType &
    GetNeighbors() const
    {
      return m_Identifiers;
    }

    [[nodiscard]] InstanceIdentifier
    GetNeighbor(unsigned int index) const
    {
      return m_Identifiers[index];
    }
 
  private:
    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
} // namespace itk::Statistics
 
#ifndef ITK_MANUAL_INSTANTIATION
#  include "itkKdTree.hxx"
#endif
 
#endif