neighborGraph.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"
 
namespace Ponca
{
    template <typename _Traits>
    struct NeighborGraphBuffer : public NeighborGraphBufferBase<_Traits>
    {
        WRITE_NEIGHBOR_GRAPH_ALIASES
        using Base = NeighborGraphBufferBase<_Traits>;
 
        IndexContainer ranges;
 
        PONCA_MULTIARCH inline NeighborGraphBuffer() = default;
        PONCA_MULTIARCH inline NeighborGraphBuffer(PointContainerConstRef _points) : Base(_points) {}
        PONCA_MULTIARCH inline NeighborGraphBuffer(PointContainerConstRef _points,
                                                   typename Traits::IndexContainerRef _indices,
                                                   const size_t _points_size, const size_t _indices_size,
                                                   typename Traits::IndexContainerRef _ranges)
            : Base(_points, _indices, _points_size, _indices_size), ranges(_ranges)
        {
        }
    };
 
    template <typename _Traits>
    class StaticNeighborGraphBase
        : public AbstractNeighborGraph<_Traits, NeighborGraphBuffer,
                                       NeighborGraphOneConnectedQuery<StaticNeighborGraphBase<_Traits>>,
                                       NeighborGraphRangeQuery<StaticNeighborGraphBase<_Traits>>>
    {
    public:
        WRITE_NEIGHBOR_GRAPH_ALIASES
 
        friend class NeighborGraphOneConnectedQuery<StaticNeighborGraphBase<Traits>>;
        friend class NeighborGraphRangeQuery<StaticNeighborGraphBase<Traits>>;
 
        using Base    = AbstractNeighborGraph<Traits, NeighborGraphBuffer,
                                              NeighborGraphOneConnectedQuery<StaticNeighborGraphBase<Traits>>,
                                              NeighborGraphRangeQuery<StaticNeighborGraphBase<Traits>>>;
        using Buffers = typename Base::Buffers;
 
        PONCA_MULTIARCH inline StaticNeighborGraphBase(const Buffers& _bufs) : Base(_bufs) {}
 
        PONCA_MULTIARCH [[nodiscard]] inline int k(int vId = 0) const { return endId(vId) * beginId(vId); }
        PONCA_MULTIARCH [[nodiscard]] inline int beginId(int vId) const { return Base::buffers().ranges[vId]; }
        PONCA_MULTIARCH [[nodiscard]] inline int endId(int vId) const { return Base::buffers().ranges[vId + 1]; }
    };
 
    template <typename _Traits>
    class NeighborGraphBase : public StaticNeighborGraphBase<_Traits>
    {
    public:
        WRITE_NEIGHBOR_GRAPH_ALIASES
 
    private:
        using Base    = StaticNeighborGraphBase<Traits>;
        using Buffers = typename Base::Buffers;
 
    public:
        template <typename KdTreeTraits>
        PONCA_MULTIARCH_HOST inline NeighborGraphBase(const KdTreeBase<KdTreeTraits>& _kdtree, const Scalar range)
            : Base(_kdtree.points())
        {
            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::PointContainerConstRef, 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.ranges.resize(cloudSize + 1); // we need one more index (see StaticNeighborGraphBase::endId)
            Base::m_bufs.ranges[0] = 0;                // first element starts at 0;
 
            for (int i = 0; i < cloudSize; ++i)
            {
                for (auto n : _kdtree.rangeNeighbors(typename KdTreeTraits::IndexType(i), range))
                {
                    Base::m_bufs.indices.push_back(n);
                }
                Base::m_bufs.ranges[i + 1] = IndexType(Base::m_bufs.indices.size());
            }
            Base::m_bufs.indices_size = Base::m_bufs.indices.size();
        }
    };
 
    template <typename DataPoint>
    using NeighborGraph = NeighborGraphBase<NeighborGraphDefaultTraits<DataPoint>>;
 
} // namespace Ponca
 
#undef WRITE_TRAITS