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 "../../Common/Assert.h"
 
#include <memory>
 
namespace Ponca
{
 
    template <typename Traits>
    class KnnGraphBase;
 
    template <typename DataPoint>
    using KnnGraph = KnnGraphBase<KnnGraphDefaultTraits<DataPoint>>;
 
    template <typename Traits>
    class StaticKnnGraphBase
    {
    public:
#define WRITE_TRAITS                                                                                                 \
    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; 
        WRITE_TRAITS
 
        using KNearestIndexQuery = KnnGraphKNearestQuery<Traits>;
        using RangeIndexQuery    = KnnGraphRangeQuery<Traits>;
        friend class KnnGraphKNearestQuery<Traits>; 
        friend class KnnGraphRangeQuery<Traits>;    
    protected:
        PONCA_MULTIARCH inline const IndexType* getIndexPtr() const { return Traits::getIndexRawPtr(m_bufs.indices); }
        PONCA_MULTIARCH inline IndexType* getIndexPtr() { return Traits::getIndexRawPtr(m_bufs.indices); }
 
    public:
        struct Buffers
        {
            PointContainer points;  
            IndexContainer indices; 
 
            size_t points_size{0};
            size_t indices_size{0};
            int k{0};
 
            PONCA_MULTIARCH inline Buffers() = default;
            PONCA_MULTIARCH inline Buffers(PointContainer _points, const int _k) : points(_points), k(_k) {}
 
            PONCA_MULTIARCH inline Buffers(PointContainer _points, typename Traits::IndexContainerRef _indices,
                                           const size_t _points_size, const size_t _indices_size, const int _k)
                : points(_points), indices(_indices), points_size(_points_size), indices_size(_indices_size), k(_k)
            {
            }
        };
 
    protected:
        PONCA_MULTIARCH inline StaticKnnGraphBase(PointContainer _points, const int _k) : m_bufs(_points, _k) {}
 
    public:
        PONCA_MULTIARCH inline StaticKnnGraphBase(Buffers& _bufs) : m_bufs(_bufs) {}
 
        // Query -------------------------------------------------------------------
    public:
        PONCA_MULTIARCH inline KNearestIndexQuery kNearestNeighbors(int index) const
        {
            return KNearestIndexQuery(this, index);
        }
 
        PONCA_MULTIARCH inline RangeIndexQuery rangeNeighbors(int index, Scalar r) const
        {
            return RangeIndexQuery(this, r, index);
        }
 
        PONCA_MULTIARCH inline KNearestIndexQuery kNearestNeighborsIndexQuery() const
        {
            return KNearestIndexQuery(this, 0);
        }
 
        PONCA_MULTIARCH inline RangeIndexQuery rangeNeighborsIndexQuery() const { return RangeIndexQuery(this, 0, 0); }
 
        // Accessors ---------------------------------------------------------------
    public:
        PONCA_MULTIARCH [[nodiscard]] inline int k() const { return m_bufs.k; }
        PONCA_MULTIARCH [[nodiscard]] inline size_t size() const
        {
            return m_bufs.indices_size / static_cast<size_t>(m_bufs.k);
        }
        PONCA_MULTIARCH [[nodiscard]] inline IndexType sampleCount() const { return (IndexType)m_bufs.indices_size; }
        PONCA_MULTIARCH [[nodiscard]] inline IndexType pointCount() const { return (IndexType)m_bufs.points_size; }
        PONCA_MULTIARCH [[nodiscard]] inline PointContainer points() const { return m_bufs.points; };
        PONCA_MULTIARCH [[nodiscard]] inline IndexContainer samples() const { return m_bufs.indices; };
        PONCA_MULTIARCH [[nodiscard]] inline const Buffers& buffers() const { return m_bufs; }
 
        // Data --------------------------------------------------------------------
    protected:          // for friends relations
        Buffers m_bufs; 
    };
 
    template <typename Traits>
    class KnnGraphBase : public StaticKnnGraphBase<Traits>
    {
    public:
        WRITE_TRAITS
    private:
        using Base    = StaticKnnGraphBase<Traits>;
        using Buffers = typename StaticKnnGraphBase<Traits>::Buffers;
        // knnGraph ----------------------------------------------------------------
    public:
        template <typename KdTreeTraits>
        PONCA_MULTIARCH_HOST inline KnnGraphBase(const KdTreeBase<KdTreeTraits>& _kdtree, const int _k = 6)
            : Base(_kdtree.points(), std::min(_k, _kdtree.sampleCount() - 1))
        {
            Base::m_bufs.points_size = _kdtree.pointCount();
#define CHECK_TRAITS_TYPENAME_COMPAT(A, B)                                                            \
    static_assert(std::is_same_v<A, B> || std::is_convertible_v<A, B> || std::is_convertible_v<B, A>, \
                  "KdTreeTraits::DataPoint is not equal to Traits::DataPoint");
 
            static_assert(std::is_same_v<typename Traits::DataPoint, typename KdTreeTraits::DataPoint>,
                          "KdTreeTraits::DataPoint is not equal to Traits::DataPoint");
 
            CHECK_TRAITS_TYPENAME_COMPAT(typename Traits::PointContainer, typename KdTreeTraits::PointContainer)
            CHECK_TRAITS_TYPENAME_COMPAT(typename Traits::IndexContainer, typename KdTreeTraits::IndexContainer)
 
#undef CHECK_TRAITS_TYPENAME_COMPAT
 
            // 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();
                PONCA_ASSERT(cloudSize == samplesSize);
            }
 
            Base::m_bufs.indices_size = cloudSize * Base::m_bufs.k;
            Base::m_bufs.indices.resize(Base::m_bufs.indices_size, -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(Base::m_bufs.k)))
                {
                    Base::m_bufs.indices[i * Base::m_bufs.k + j] = n;
                    ++j;
                }
            }
        }
    };
 
} // namespace Ponca
 
#undef WRITE_TRAITS