Feature/use kd tree for landmark map search

MODULE.bazel

@@ -19,6 +19,7 @@ module(
 
 bazel_dep(name = "bazel_skylib", version = "1.8.1")
 bazel_dep(name = "eigen", version = "3.4.0.bcr.3")
+bazel_dep(name = "nanoflann", version = "1.5.5")
 bazel_dep(name = "onetbb", version = "2022.1.0")
 bazel_dep(name = "package_metadata", version = "0.0.5")
 bazel_dep(name = "platforms", version = "1.0.0")

beluga/BUILD.bazel

@@ -28,6 +28,7 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         "@eigen",
+        "@nanoflann",
         "@onetbb//:tbb",
         "@range-v3",
         "@sophus",

beluga/CMakeLists.txt

@@ -48,12 +48,14 @@ endif()
 
 find_package(Eigen3 REQUIRED NO_MODULE)
 find_package(HDF5 COMPONENTS CXX)
+find_package(nanoflann REQUIRED)
 find_package(range-v3 REQUIRED)
 find_package(Sophus REQUIRED)
 find_package(TBB REQUIRED)
 
 set(_deps
     Eigen3::Eigen
+    nanoflann::nanoflann
     range-v3::range-v3
     Sophus::Sophus
     TBB::tbb)

beluga/README.md

@@ -34,6 +34,7 @@ Auto-generated Doxygen documentation can be found in https://ekumen-os.github.io
 Beluga is built on top of the following open source libraries:
 
 - [Eigen](https://gitlab.com/libeigen/eigen): A well-known C++ template library for linear algebra: matrices, vectors, numerical solvers, and related algorithms.
+- [Nanoflann](https://github.com/jlblancoc/nanoflann): A C++11 header-only library for building KD-Trees.
 - [Sophus](https://github.com/strasdat/Sophus): A C++ implementation of Lie groups using Eigen.
 - [Range](https://github.com/ericniebler/range-v3): The basis library for C++20's `std::ranges`.
 

beluga/cmake/Config.cmake.in

@@ -2,6 +2,7 @@
 
 include(CMakeFindDependencyMacro)
 find_dependency(Eigen3 REQUIRED NO_MODULE)
+find_dependency(nanoflann REQUIRED)
 find_dependency(range-v3 REQUIRED)
 find_dependency(HDF5 COMPONENTS CXX REQUIRED)
 find_dependency(Sophus REQUIRED)

beluga/docs/index.md

@@ -38,6 +38,7 @@ The current set of features includes:
 Beluga is built on top of the following open source libraries:
 
 - [Eigen](https://gitlab.com/libeigen/eigen): A well-known C++ template library for linear algebra: matrices, vectors, numerical solvers, and related algorithms.
+- [nanoflann](https://github.com/jlblancoc/nanoflann): A C++11 header-only library for building KD-Trees.
 - [Sophus](https://github.com/strasdat/Sophus): A C++ implementation of Lie groups using Eigen.
 - [Range](https://github.com/ericniebler/range-v3): The basis library for C++20's `std::ranges`.
 

beluga/include/beluga/sensor/data/landmark_map.hpp

@@ -16,13 +16,21 @@
 #define BELUGA_SENSOR_DATA_LANDMARK_MAP_HPP
 
 // external
+// Bazel exposes nanoflann with a prefix, CMake does not
+#if __has_include(<nanoflann/nanoflann.hpp>)
+#include <nanoflann/nanoflann.hpp>
+#else
+#include <nanoflann.hpp>
+#endif
 #include <range/v3/view/filter.hpp>
 #include <range/v3/view/tail.hpp>
 #include <sophus/se3.hpp>
 
 // standard library
 #include <algorithm>
 #include <cstdint>
+#include <memory>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -48,7 +56,9 @@ class LandmarkMap {
   /// @param boundaries Limits of the map.
   /// @param landmarks List of landmarks that can be expected to be detected.
   explicit LandmarkMap(const LandmarkMapBoundaries& boundaries, landmarks_set_position_data landmarks)
-      : landmarks_(std::move(landmarks)), map_boundaries_(std::move(boundaries)) {}
+      : landmarks_(std::move(landmarks)), map_boundaries_(std::move(boundaries)) {
+    build_category_indices();
+  }
 
   /// @brief Constructor with implicit map boundaries (computed from landmarks).
   /// @details Note that computing map boundaries from landmarks will effectively
@@ -67,9 +77,46 @@ class LandmarkMap {
         map_boundaries_.min() = map_boundaries_.min().cwiseMin(position);
         map_boundaries_.max() = map_boundaries_.max().cwiseMax(position);
       }
+      build_category_indices();
+    }
+  }
+
+  /// @brief Copy constructor.
+  /// @details Copying requires expensive reconstruction of the cached kd-tree indices.
+  /// @deprecated Use move semantics (std::move) instead to avoid performance overhead.
+  /// @param other Landmark map to copy from.
+  [[deprecated("LandmarkMap copying is expensive. Use std::move() instead.")]] LandmarkMap(const LandmarkMap& other)
+      : landmarks_(other.landmarks_), map_boundaries_(other.map_boundaries_) {
+    build_category_indices();
+  }
+
+  /// @brief Copy assignment operator.
+  /// @details Copying requires expensive reconstruction of the cached kd-tree indices.
+  /// @deprecated Use move semantics (std::move) instead to avoid performance overhead.
+  /// @param other Landmark map to copy from.
+  /// @return Reference to this landmark map.
+  [[deprecated("LandmarkMap copying is expensive. Use std::move() instead.")]] LandmarkMap& operator=(
+      const LandmarkMap& other) {
+    if (this != &other) {
+      landmarks_ = other.landmarks_;
+      map_boundaries_ = other.map_boundaries_;
+      category_indices_.clear();
+      build_category_indices();
     }
+    return *this;
   }
 
+  /// @brief Move constructor.
+  /// @details Explicitly defaulted so landmark data and cached kd-tree indices
+  /// can be transferred efficiently while preserving unique ownership semantics.
+  LandmarkMap(LandmarkMap&&) = default;
+
+  /// @brief Move assignment operator.
+  /// @details Explicitly defaulted so ownership of the cached kd-tree indices can
+  /// be transferred efficiently without rebuilding them.
+  /// @return Reference to this landmark map.
+  LandmarkMap& operator=(LandmarkMap&&) = default;
+
   /// @brief Returns the map boundaries.
   /// @return The map boundaries.
   [[nodiscard]] LandmarkMapBoundaries map_limits() const { return map_boundaries_; }
@@ -81,32 +128,19 @@ class LandmarkMap {
   [[nodiscard]] std::optional<LandmarkPosition3> find_nearest_landmark(
       const LandmarkPosition3& detection_position_in_world,
       const LandmarkCategory& detection_category) const {
-    // only consider those that have the same id
-    auto same_category_landmarks_view =
-        landmarks_ | ranges::views::filter([detection_category = detection_category](const auto& l) {
-          return detection_category == l.category;
-        });
-
-    // find the landmark that minimizes the distance to the detection position
-    // This is O(n). A spatial data structure should be used instead.
-    auto min = std::min_element(
-        same_category_landmarks_view.begin(), same_category_landmarks_view.end(),
-        [&detection_position_in_world](const auto& a, const auto& b) {
-          const auto& landmark_a_position_in_world = a.detection_position_in_robot;
-          const auto& landmark_b_position_in_world = b.detection_position_in_robot;
-
-          const auto landmark_b_squared_in_world_squared =
-              (landmark_a_position_in_world - detection_position_in_world).squaredNorm();
-          const auto landmark_b_distance_in_world_squared =
-              (landmark_b_position_in_world - detection_position_in_world).squaredNorm();
-          return landmark_b_squared_in_world_squared < landmark_b_distance_in_world_squared;
-        });
-
-    if (min == same_category_landmarks_view.end()) {
+    const auto it = category_indices_.find(detection_category);
+    if (it == category_indices_.end()) {
       return std::nullopt;
     }
-
-    return min->detection_position_in_robot;
+    const auto& index = *it->second;
+    const std::array<double, 3> query = {
+        detection_position_in_world.x(), detection_position_in_world.y(), detection_position_in_world.z()};
+    CategoryIndexType result_idx;
+    double result_dist_sq;
+    if (index.tree->knnSearch(query.data(), 1, &result_idx, &result_dist_sq) == 0) {
+      return std::nullopt;
+    }
+    return index.cloud.pts[result_idx];
   }
 
   /// @brief Finds the landmark that minimizes the bearing error to a given detection and returns its data.
@@ -161,8 +195,52 @@ class LandmarkMap {
   }
 
  private:
+  /// Point cloud adapter used by the category kd-trees.
+  struct PositionCloud {
+    std::vector<LandmarkPosition3> pts;
+    [[nodiscard]] std::size_t kdtree_get_point_count() const { return pts.size(); }
+    [[nodiscard]] double kdtree_get_pt(std::size_t i, std::size_t dim) const { return pts[i](static_cast<int>(dim)); }
+    template <class BBox>
+    bool kdtree_get_bbox(BBox&) const {
+      return false;
+    }
+  };
+
+  /// Index type used by the category kd-trees.
+  using CategoryIndexType = std::uint32_t;
+
+  /// kd-tree type used to query landmarks within each category.
+  using CategoryKDTree = nanoflann::KDTreeSingleIndexAdaptor<
+      nanoflann::L2_Simple_Adaptor<double, PositionCloud>,
+      PositionCloud,
+      3,
+      CategoryIndexType>;
+
+  /// Cached landmark positions and search tree for a single category.
+  struct CategoryIndex {
+    PositionCloud cloud;
+    std::unique_ptr<CategoryKDTree> tree;
+  };
+
   landmarks_set_position_data landmarks_;
   LandmarkMapBoundaries map_boundaries_;
+  std::unordered_map<LandmarkCategory, std::unique_ptr<CategoryIndex>> category_indices_;
+
+  /// @brief Builds per-category kd-tree indices for nearest-neighbor search.
+  void build_category_indices() {
+    for (const auto& l : landmarks_) {
+      auto& entry = category_indices_[l.category];
+      if (!entry) {
+        entry = std::make_unique<CategoryIndex>();
+      }
+      entry->cloud.pts.push_back(l.detection_position_in_robot);
+    }
+    for (auto& [_, entry] : category_indices_) {
+      entry->tree = std::make_unique<CategoryKDTree>(
+          3, entry->cloud, nanoflann::KDTreeSingleIndexAdaptorParams(/*leaf_max_size=*/10));
+      entry->tree->buildIndex();
+    }
+  }
 };
 
 }  // namespace beluga

beluga/package.xml

@@ -17,6 +17,7 @@
 
   <depend>eigen</depend>
   <depend>libhdf5-dev</depend>
+  <depend>libnanoflann-dev</depend>
   <depend>range-v3</depend>
   <depend>sophus</depend>
   <depend>tbb</depend>