knnGraph.h

/*
 This Source Code Form is subject to the terms of the Mozilla Public
 License, v. 2.0. If a copy of the MPL was not distributed with this
 file, You can obtain one at http://mozilla.org/MPL/2.0/.
*/
 
#pragma once
 
#include "./knnGraphTraits.h"
 
#include "Query/knnGraphKNearestQuery.h"
#include "Query/knnGraphRangeQuery.h"
 
#include "../KdTree/kdTree.h"
 
#include <memory>
 
namespace Ponca {
 
template <typename Traits> class KnnGraphBase;
 
template <typename DataPoint>
using KnnGraph = KnnGraphBase<KnnGraphDefaultTraits<DataPoint>>;
 
template <typename Traits> class KnnGraphBase
{
public:
    using DataPoint  = typename Traits::DataPoint; 
    using Scalar     = typename DataPoint::Scalar; 
    using VectorType = typename DataPoint::VectorType; 
 
    using IndexType      = typename Traits::IndexType;
    using PointContainer = typename Traits::PointContainer; 
    using IndexContainer = typename Traits::IndexContainer; 
 
    using KNearestIndexQuery = KnnGraphKNearestQuery<Traits>;
    using RangeIndexQuery    = KnnGraphRangeQuery<Traits>;
 
    friend class KnnGraphKNearestQuery<Traits>; // This type must be equal to KnnGraphBase::KNearestIndexQuery
    friend class KnnGraphRangeQuery<Traits>;    // This type must be equal to KnnGraphBase::RangeIndexQuery
 
    // knnGraph ----------------------------------------------------------------
public:
    template<typename KdTreeTraits>
    inline KnnGraphBase(const KdTreeBase<KdTreeTraits>& kdtree, int k = 6)
            : m_k(std::min(k,kdtree.sample_count()-1)),
              m_kdTreePoints(kdtree.points())
    {
        static_assert( std::is_same<typename Traits::DataPoint, typename KdTreeTraits::DataPoint>::value,
                       "KdTreeTraits::DataPoint is not equal to Traits::DataPoint" );
        static_assert( std::is_same<typename Traits::PointContainer, typename KdTreeTraits::PointContainer>::value,
                       "KdTreeTraits::PointContainer is not equal to Traits::PointContainer" );
        static_assert( std::is_same<typename Traits::IndexContainer, typename KdTreeTraits::IndexContainer>::value,
                       "KdTreeTraits::IndexContainer is not equal to Traits::IndexContainer" );
 
        // We need to account for the entire point set, irrespectively of the sampling. This is because the kdtree
        // (kNearestNeighbors) return ids of the entire point set, not it sub-sampled list of ids.
        // \fixme Update API to properly handle kdtree subsampling
        const int cloudSize   = kdtree.pointCount();
        {
            const int samplesSize = kdtree.sample_count();
            eigen_assert(cloudSize == samplesSize);
        }
 
        m_indices.resize(cloudSize * m_k, -1);
 
#pragma omp parallel for shared(kdtree, cloudSize) default(none)
        for(int i=0; i<cloudSize; ++i)
        {
            int j = 0;
            for(auto n : kdtree.kNearestNeighbors(typename KdTreeTraits::IndexType(i),
                                                   typename KdTreeTraits::IndexType(m_k)))
            {
                m_indices[i * m_k + j] = n;
                ++j;
            }
        }
    }
 
    // Query -------------------------------------------------------------------
public:
    inline KNearestIndexQuery kNearestNeighbors(int index) const{
        return KNearestIndexQuery(this, index);
    }
 
    inline RangeIndexQuery    rangeNeighbors(int index, Scalar r) const{
        return RangeIndexQuery(this, r, index);
    }
 
    inline KNearestIndexQuery kNearestNeighborsIndexQuery() const
    {
        return KNearestIndexQuery(this, 0);
    }
 
    inline RangeIndexQuery rangeNeighborsIndexQuery() const
    {
        return RangeIndexQuery(this, 0, 0);
    }
 
    // Accessors ---------------------------------------------------------------
public:
    [[nodiscard]] inline int k() const { return m_k; }
    [[nodiscard]] inline size_t size() const { return m_indices.size()/static_cast<size_t>(m_k); }
 
    // Data --------------------------------------------------------------------
private:
    const int m_k;
    IndexContainer m_indices; 
 
protected: // for friends relations
    const PointContainer& m_kdTreePoints;
    inline const IndexContainer& index_data() const { return m_indices; };
};
 
} // namespace Ponca