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 "abstractNeighborGraph.h"
#include "Query/neighborGraphOneConnectedQuery.h"
#include "Query/neighborGraphRangeQuery.h"
 
#include "../KdTree/kdTree.h"
#include "../../Common/Assert.h"
 
namespace Ponca
{

    template <typename _Traits>
    struct KnnGraphBuffers : public NeighborGraphBufferBase<_Traits>
    {
        WRITE_NEIGHBOR_GRAPH_ALIASES
        using Base = NeighborGraphBufferBase<_Traits>;

        static constexpr int DefaultKInKnnGraph() { return 6; }

        const int k{DefaultKInKnnGraph()};
 
        PONCA_MULTIARCH inline KnnGraphBuffers() = default;
        PONCA_MULTIARCH inline KnnGraphBuffers(PointContainerConstRef _points, const int _k = DefaultKInKnnGraph())
            : Base(_points), k(_k)
        {
        }
        PONCA_MULTIARCH inline KnnGraphBuffers(PointContainerConstRef _points, IndexContainerRef _indices,
                                               const size_t _points_size, const size_t _indices_size, const int _k)
            : Base(_points, _indices, _points_size, _indices_size), k(_k)
        {
        }
    };

    template <typename _Traits>
    class StaticKnnGraphBase : public AbstractNeighborGraph<_Traits, KnnGraphBuffers,
                                                            NeighborGraphOneConnectedQuery<StaticKnnGraphBase<_Traits>>,
                                                            NeighborGraphRangeQuery<StaticKnnGraphBase<_Traits>>>
    {
    public:
        WRITE_NEIGHBOR_GRAPH_ALIASES
 
        friend class NeighborGraphOneConnectedQuery<StaticKnnGraphBase<Traits>>; 
        friend class NeighborGraphRangeQuery<StaticKnnGraphBase<_Traits>>;       
        using Base =
            AbstractNeighborGraph<Traits, KnnGraphBuffers, NeighborGraphOneConnectedQuery<StaticKnnGraphBase<Traits>>,
                                  NeighborGraphRangeQuery<StaticKnnGraphBase<_Traits>>>;
        using Buffers = typename Base::Buffers;
 
        PONCA_MULTIARCH inline StaticKnnGraphBase<Traits>(Buffers& _bufs) : Base(_bufs) {}
 
    protected:
        PONCA_MULTIARCH inline StaticKnnGraphBase(PointContainerConstRef _points, const int _k)
            : Base(Buffers(_points, _k))
        {
        }
 
        // Accessors ---------------------------------------------------------------
    public:
        PONCA_MULTIARCH [[nodiscard]] inline int k(int /*index*/ = 0) const { return Base::buffers().k; }
        PONCA_MULTIARCH [[nodiscard]] inline int beginId(int vId) const { return vId * k(); }
        PONCA_MULTIARCH [[nodiscard]] inline int endId(int vId) const { return (vId + 1) * k(); }

        PONCA_MULTIARCH [[nodiscard]] inline typename Base::OneConnectedIndexQuery kNearestNeighbors(int index) const
        {
            return Base::oneConnectedNeighbors(index);
        }

        PONCA_MULTIARCH [[nodiscard]] inline typename Base::OneConnectedIndexQuery kNearestNeighborsIndexQuery() const
        {
            return Base::oneConnectedNeighbors();
        }
    };
 
    template <typename _Traits>
    class KnnGraphBase : public StaticKnnGraphBase<_Traits>
    {
    public:
        WRITE_NEIGHBOR_GRAPH_ALIASES
 
    private:
        using Base    = StaticKnnGraphBase<Traits>;
        using Buffers = typename Base::Buffers;
        // knnGraph ----------------------------------------------------------------
    public:
        template <typename KdTreeTraits>
        PONCA_MULTIARCH_HOST inline KnnGraphBase(const KdTreeBase<KdTreeTraits>& _kdtree,
                                                 const int _k = Buffers::DefaultKInKnnGraph())
            : 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;
                }
            }
        }
    };

    template <typename DataPoint>
    using KnnGraph = KnnGraphBase<NeighborGraphDefaultTraits<DataPoint>>;
 
} // namespace Ponca
 
#undef WRITE_TRAITS