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.sampleCount() - 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.sampleCount();
                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