kdtree.hpp

Go to the documentation of this file.

#ifndef KDTREE_HPP
#define KDTREE_HPP
 
/**
 * \file kdtree.hpp
 * \brief A k-d tree wrapper for broad-phase collision detection using
 *        nanoflann.
 * \author Le Bars, Yoann
 * \ingroup PhysicsAcceleration
 *
 * This file is part of the pure C++ benchmark.
 */
 
#include <nanoflann.hpp>
 
#include "body.hpp"
 
namespace Model {
 
    /**
     * \brief An adapter to allow nanoflann to work directly with our `Bodies`
     *        SoA container.
     *
     * Note: The method names in this class follow the `snake_case` convention
     * required by the nanoflann library's adapter interface, rather than the
     * project's standard `camelCase`.
     * \ingroup PhysicsAcceleration
     */
    class BodiesAdaptor {
       public:
        /**
         * \brief Construct a new Bodies Adaptor object.
         * \param in_bodies The `Bodies` container.
         */
        explicit BodiesAdaptor(const Bodies& in_bodies) : bodies_(in_bodies) {}
 
        /**
         * \brief Returns the number of bodies (points) in the dataset.
         * \return The number of bodies.
         */
        [[nodiscard]] inline std::size_t kdtree_get_point_count() const {
            return bodies_.size();
        }
 
        /**
         * \brief Returns the value of a single dimension of a point.
         * \param idx The index of the point.
         * \param dim The dimension (0 for x, 1 for y, 2 for z).
         * \return The coordinate value.
         */
        [[nodiscard]] inline double kdtree_get_pt(const std::size_t idx,
                                                  const std::size_t dim) const {
            return bodies_.x_[idx][dim];
        }
 
        /**
         * \brief Optional bounding box calculation (not used here).
         * \tparam BBOX The type of the bounding box.
         * \return Always false.
         */
        template <class BBOX>
        bool kdtree_get_bbox(BBOX& /*bb*/) const {
            return false;
        }
 
       private:
        /// \brief A const reference to the bodies container.
        const Bodies& bodies_;
    };
 
    /**
     * \brief A k-d tree for fast nearest neighbour searches, specialized for
     *        our `Bodies` container.
     */
    class KDTree {
       public:
        /// \brief The specific KD-Tree type using our adapter.
        using TreeType = nanoflann::KDTreeSingleIndexAdaptor<
            nanoflann::L2_Simple_Adaptor<double, BodiesAdaptor>,
            BodiesAdaptor,  // NOLINT
            3 /* dimensionality */, std::size_t /* index type */>;
 
        /**
         * \brief Construct a new KDTree object and build the index.
         * \param in_bodies The `Bodies` container to build the tree from.
         */
        explicit KDTree(const Bodies& in_bodies)
            : adaptor_(in_bodies),
              tree_(3 /* dimensionality */, adaptor_,
                    // Set leaf_max_size to 10. This is a good balance for
                    // performance.
                    nanoflann::KDTreeSingleIndexAdaptorParams(10)) {
            tree_.buildIndex();
        }
 
        /**
         * \brief Finds all bodies within a given radius of a query point.
         *
         * \param query_point The centre point of the search.
         * \param search_radius The radius of the search sphere.
         * \param out_results A vector to store the results (pairs of index and
         *                    squared distance).
         * \return The number of neighbours found.
         */
        template <typename VectorOfPairs>
        std::size_t radiusSearch(const double* query_point,
                                 const double search_radius,
                                 VectorOfPairs& out_results) {
/* The SearchParameters class was renamed from SearchParams in v1.5.0.
   We check for the legacy version macro to decide which API to use. */
#if defined(NANOFLANN_VERSION) && NANOFLANN_VERSION < 0x150
            // Use the legacy API for nanoflann < 1.5.0
            return tree_.radiusSearch(query_point,
                                      search_radius * search_radius,
                                      out_results, nanoflann::SearchParams());
#else
            // Use the modern API for nanoflann >= 1.5.0
            return tree_.radiusSearch(
                query_point, search_radius * search_radius, out_results,
                nanoflann::SearchParameters());
#endif
        }
 
       private:
        /// \brief The adapter that lets nanoflann access our data.
        BodiesAdaptor adaptor_;
        /// \brief The underlying nanoflann k-d tree index.
        TreeType tree_;
    };
 
}  // namespace Model
 
#endif  // KDTREE_HPP