Nanoflann KD-tree spatial indexing for SVA/SRA

CMakeLists.txt

@@ -50,6 +50,7 @@ find_external_dependency("Eigen3" "Eigen3::Eigen" "${CMAKE_CURRENT_LIST_DIR}/cma
 find_package(OpenMP QUIET REQUIRED)
 robin_download_pmc()
 robin_download_xenium()
+robin_download_nanoflann()
 add_subdirectory(${pmc_SOURCE_DIR} ${pmc_BINARY_DIR})
 
 #
@@ -66,6 +67,7 @@ add_library(robin::robin ALIAS robin)
 target_link_libraries(robin
         PUBLIC Eigen3::Eigen
                xenium
+               nanoflann
         PRIVATE OpenMP::OpenMP_CXX
         PRIVATE pmc)
 target_include_directories(
@@ -112,7 +114,7 @@ export(
 )
 
 install(
-    TARGETS xenium pmc
+    TARGETS xenium pmc nanoflann
     EXPORT robinTargets
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}

cmake/DownloadExternal.cmake

@@ -90,6 +90,39 @@ function(robin_download_xenium)
 	install(DIRECTORY ${xenium_SOURCE_DIR}/xenium DESTINATION ${CMAKE_INSTALL_INCLUDEDIR})
 endfunction()
 
+# nanoflann
+function(robin_download_nanoflann)
+	download_project(PROJ nanoflann
+			GIT_REPOSITORY https://github.com/jlblancoc/nanoflann.git
+			GIT_TAG        v1.6.2
+			QUIET
+			)
+	add_library(nanoflann INTERFACE)
+	target_include_directories(nanoflann
+			INTERFACE
+			$<BUILD_INTERFACE:${nanoflann_SOURCE_DIR}/include>
+			$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>)
+
+	install(
+			TARGETS nanoflann
+			EXPORT nanoflannTargets
+			ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+	)
+
+	export(
+			EXPORT nanoflannTargets
+			FILE ${CMAKE_CURRENT_BINARY_DIR}/nanoflannTargets.cmake
+			NAMESPACE nanoflann::
+	)
+
+	install(
+			EXPORT nanoflannTargets
+			FILE nanoflannTargets.cmake
+			DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/nanoflann/
+			NAMESPACE nanoflann::
+	)
+endfunction()
+
 # For handling Eigen3 library
 function(find_external_dependency PACKAGE_NAME TARGET_NAME INCLUDED_CMAKE_PATH)
   string(TOUPPER ${PACKAGE_NAME} PACKAGE_NAME_UP)

src/robin.cpp

@@ -6,6 +6,51 @@
 
 #include <robin/robin.hpp>
 
+#include <nanoflann.hpp>
+
+namespace {
+
+/**
+ * @brief Convert a directed half-edge adjacency list (where only edges i->j with j>i are stored)
+ *        into a full undirected graph of the requested storage type.
+ */
+robin::IGraph* BuildGraphFromHalfEdges(std::vector<std::vector<size_t>>& half_adj_list,
+                                       robin::GraphsStorageType graph_storage_type) {
+  const size_t N = half_adj_list.size();
+
+  size_t total_edges = 0;
+  for (size_t i = 0; i < N; ++i) {
+    total_edges += half_adj_list[i].size();
+  }
+
+  if (graph_storage_type == robin::GraphsStorageType::ADJ_LIST) {
+    std::vector<std::vector<size_t>> full_adj_list(N);
+    for (size_t i = 0; i < N; ++i) {
+      for (size_t j : half_adj_list[i]) {
+        full_adj_list[i].push_back(j);
+        full_adj_list[j].push_back(i);
+      }
+    }
+    return new robin::AdjListGraph(std::move(full_adj_list), total_edges);
+  }
+
+  // CSR and ATOMIC_CSR use edge-list constructors
+  std::vector<std::pair<size_t, size_t>> edge_list;
+  edge_list.reserve(total_edges);
+  for (size_t i = 0; i < N; ++i) {
+    for (size_t j : half_adj_list[i]) {
+      edge_list.emplace_back(i, j);
+    }
+  }
+
+  if (graph_storage_type == robin::GraphsStorageType::CSR) {
+    return new robin::CSRGraph(edge_list);
+  }
+  return new robin::AtomicCSRGraph(edge_list);
+}
+
+} // anonymous namespace
+
 std::vector<size_t> robin::FindInlierStructure(const IGraph* g,
                                                InlierGraphStructure graph_structure) {
   // identify inlier structures
@@ -27,65 +72,100 @@ std::vector<size_t> robin::FindInlierStructure(const IGraph* g,
 
 robin::IGraph* robin::MakeVecAvgInvGraph(const Eigen::MatrixXd& measurements, double noise_bound,
                                          GraphsStorageType graph_storage_type) {
-  // Measurements set
-  VectorY y_set(measurements);
+  const size_t N = measurements.cols();
+  const size_t dim = measurements.rows();
 
-  // Compatibility check function
-  robin::SvaCompCheck sva_comp_check(noise_bound);
+  // Transpose to row-major (N x dim) for nanoflann
+  Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> data_rowmajor =
+      measurements.transpose();
 
-  // Compatibility graph constructor
-  robin::CompGraphConstructor<VectorY, robin::SvaCompCheck, 2> graph_constructor;
-  graph_constructor.SetCompCheckFunction(&sva_comp_check);
-  graph_constructor.SetMeasurements(&y_set);
+  // Build KD-tree using nanoflann's Eigen matrix adaptor
+  using kd_tree_t = nanoflann::KDTreeEigenMatrixAdaptor<
+      Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>,
+      -1 /* dimensionality at runtime */,
+      nanoflann::metric_L2>;
 
-  // InvGraphConstructor
-  auto* g = graph_constructor.BuildCompGraph(graph_storage_type);
-  return g;
-}
-
-robin::IGraph* robin::MakeRotAvgInvGraph(const std::vector<Eigen::Matrix3d>& measurements,
-                                         So3Distance dist_type, double noise_bound,
-                                         GraphsStorageType graph_storage_type) {
-  // Measurements sets
-  So3Y y_set(measurements);
+  kd_tree_t tree(dim, std::cref(data_rowmajor), {10 /* max leaf size */});
 
-  // InvFunc
-  switch (dist_type) {
-  case So3Distance::CHORDAL_DISTANCE: {
-    // Distance function
-    So3ChordalDist inv_func;
+  // nanoflann radiusSearch uses squared L2 distance
+  const double search_radius_sq = noise_bound * noise_bound;
 
-    // Compatibility check function
-    robin::SraCompCheck<So3ChordalDist> comp_check(&inv_func, noise_bound);
+  std::vector<std::vector<size_t>> half_adj_list(N);
 
-    // InvGraphConstructor
-    robin::CompGraphConstructor<So3Y, robin::SraCompCheck<So3ChordalDist>, 2> graph_constructor;
-    graph_constructor.SetCompCheckFunction(&comp_check);
-    graph_constructor.SetMeasurements(&y_set);
+  nanoflann::SearchParameters params;
+  params.sorted = false;
 
-    // Build graph
-    auto* g = graph_constructor.BuildCompGraph(graph_storage_type);
+#pragma omp parallel for schedule(dynamic, 64)
+  for (size_t i = 0; i < N; ++i) {
+    std::vector<nanoflann::ResultItem<size_t, double>> matches;
+    tree.index->radiusSearch(data_rowmajor.row(i).data(), search_radius_sq, matches, params);
 
-    return g;
+    for (const auto& match : matches) {
+      if (match.first > i) {
+        half_adj_list[i].push_back(match.first);
+      }
+    }
   }
-  case robin::So3Distance::GEODESIC_DISTANCE: {
-    // InvFunc
-    So3GeodesicDist inv_func;
 
-    // Compatibility check function
-    robin::SraCompCheck<So3GeodesicDist> comp_check(&inv_func, noise_bound);
+  return BuildGraphFromHalfEdges(half_adj_list, graph_storage_type);
+}
 
-    // InvGraphConstructor
-    robin::CompGraphConstructor<So3Y, robin::SraCompCheck<So3GeodesicDist>, 2> graph_constructor;
-    graph_constructor.SetCompCheckFunction(&comp_check);
-    graph_constructor.SetMeasurements(&y_set);
+robin::IGraph* robin::MakeRotAvgInvGraph(const std::vector<Eigen::Matrix3d>& measurements,
+                                         So3Distance dist_type, double noise_bound,
+                                         GraphsStorageType graph_storage_type) {
+  const size_t N = measurements.size();
 
-    // Build graph
-    auto* g = graph_constructor.BuildCompGraph(graph_storage_type);
+  // Embed rotations as 9D vectors (flatten each 3x3 matrix)
+  Eigen::Matrix<double, Eigen::Dynamic, 9, Eigen::RowMajor> embedded(N, 9);
+  for (size_t i = 0; i < N; ++i) {
+    Eigen::Map<const Eigen::Matrix<double, 1, 9, Eigen::RowMajor>> flat(measurements[i].data());
+    embedded.row(i) = flat;
+  }
 
-    return g;
+  // Build KD-tree in R^9
+  using kd_tree_t = nanoflann::KDTreeEigenMatrixAdaptor<
+      Eigen::Matrix<double, Eigen::Dynamic, 9, Eigen::RowMajor>,
+      9,
+      nanoflann::metric_L2>;
+
+  kd_tree_t tree(9, std::cref(embedded), {10 /* max leaf size */});
+
+  // L2 in R^9 equals the Frobenius (chordal) distance.
+  // For geodesic mode, use chordal upper bound as a pre-filter:
+  // chordal = 2*sqrt(2)*sin(geodesic/2), then refine with exact geodesic check.
+  double search_radius_sq;
+  if (dist_type == So3Distance::GEODESIC_DISTANCE) {
+    double chordal_ub = 2.0 * std::sqrt(2.0) * std::sin(noise_bound / 2.0);
+    search_radius_sq = chordal_ub * chordal_ub;
+  } else {
+    search_radius_sq = noise_bound * noise_bound;
   }
+
+  std::vector<std::vector<size_t>> half_adj_list(N);
+  So3GeodesicDist geodesic_fn;
+
+  nanoflann::SearchParameters params;
+  params.sorted = false;
+
+#pragma omp parallel for schedule(dynamic, 64)
+  for (size_t i = 0; i < N; ++i) {
+    std::vector<nanoflann::ResultItem<size_t, double>> matches;
+    tree.index->radiusSearch(embedded.row(i).data(), search_radius_sq, matches, params);
+
+    for (const auto& match : matches) {
+      if (match.first > i) {
+        bool keep = true;
+        if (dist_type == So3Distance::GEODESIC_DISTANCE) {
+          keep = (geodesic_fn(measurements[i], measurements[match.first]) <= noise_bound);
+        }
+        if (keep) {
+          half_adj_list[i].push_back(match.first);
+        }
+      }
+    }
   }
+
+  return BuildGraphFromHalfEdges(half_adj_list, graph_storage_type);
 }
 
 robin::IGraph* robin::Make3dRegInvGraph(const Eigen::Matrix3Xd& src_3d_points,