KDTree: Use nanoflann

.github/workflows/build.yml

@@ -11,7 +11,7 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - run: |
-        brew install openblas swig boost python3 tbb nlopt cminpack ceres-solver bison flex hdf5 ipopt primesieve spectra pagmo cuba
+        brew install openblas swig boost python3 tbb nlopt cminpack ceres-solver bison flex hdf5 ipopt primesieve spectra pagmo cuba nanoflann
         pip3 install matplotlib scipy chaospy pandas dill --break-system-packages
     - run: |
         cmake \
@@ -37,7 +37,7 @@ jobs:
       - name: Install
         shell: cmd
         run: |
-          conda install -y cmake swig ninja winflexbison "libblas=*=*netlib" "liblapack=*=*netlib" "liblapacke=*=*netlib" mpc boost-cpp libxml2 hdf5 primesieve tbb-devel cminpack ceres-solver nlopt dlib-cpp ipopt pagmo-devel spectralib zlib scipy pandas matplotlib-base dill
+          conda install -y cmake swig ninja winflexbison "libblas=*=*netlib" "liblapack=*=*netlib" "liblapacke=*=*netlib" mpc boost-cpp libxml2 hdf5 primesieve tbb-devel cminpack ceres-solver nlopt dlib-cpp ipopt pagmo-devel spectralib zlib nanoflann scipy pandas matplotlib-base dill
       - uses: ilammy/msvc-dev-cmd@v1
       - name: Build
         shell: cmd

CMakeLists.txt

@@ -35,6 +35,7 @@ option (USE_PRIMESIEVE               "Use primesieve for prime numbers generatio
 option (USE_OPENMP                   "Use OpenMP to disable threading"                                       ON)
 option (USE_OPENBLAS                 "Use OpenBLAS to disable threading"                                     ON)
 option (USE_CXX17                    "Use C++17 standard"                                                    ON)
+option (USE_NANOFLANN                "Use nanoflann for Nearest Neighbor search"                             ON)
 
 option (BUILD_PYTHON                 "Build the python module for the library"                               ON)
 option (BUILD_SHARED_LIBS            "Build shared libraries"                                                ON)
@@ -391,6 +392,16 @@ if (USE_OPENMP)
   endif ()
 endif ()
 
+if (USE_NANOFLANN)
+  find_package (nanoflann CONFIG)
+  if (nanoflann_FOUND)
+    message (STATUS "Found nanoflann: ${nanoflann_DIR} (found version \"${nanoflann_VERSION}\")")
+    set (OPENTURNS_HAVE_NANOFLANN TRUE)
+    list (APPEND OPENTURNS_PRIVATE_LIBRARIES nanoflann::nanoflann)
+    list (APPEND OPENTURNS_ENABLED_FEATURES "nanoflann")
+  endif ()
+endif ()
+
 if (MSVC)
   # Disable some warnings
   set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /wd4197 /wd4244 /wd4251 /wd4267 /wd4275 /wd4996")

lib/etc/openturns.conf.in

@@ -48,6 +48,10 @@
   <!-- OT::DesignProxy parameters -->
   <DesignProxy-DefaultCacheSize value_int="16777216" />
 
+  <!-- OT::KDTree parameters -->
+  <KDTree-leaf_max_size value_int="10" />
+  <!--KDTree-n_thread_build value_int="1" /-->
+
   <!-- OT::KFold parameters -->
   <KFold-DefaultK value_int="10" />
 

lib/include/OTconfig.h.in

@@ -82,6 +82,9 @@
 /* Define to 1 if you have the CMinpack library. */
 #cmakedefine OPENTURNS_HAVE_CMINPACK
 
+/* Define to 1 if you have the nanoflann library. */
+#cmakedefine OPENTURNS_HAVE_NANOFLANN
+
 /* Define to 1 if you have the OpenMP specification. */
 #cmakedefine OPENTURNS_HAVE_OPENMP
 

lib/src/Base/Algo/KDTree.cxx

@@ -25,13 +25,25 @@
 #include "openturns/Indices.hxx"
 #include "openturns/SobolSequence.hxx"
 #include "openturns/PersistentObjectFactory.hxx"
+#include "openturns/ResourceMap.hxx"
+#include "openturns/OTconfig.hxx"
+
+#ifdef OPENTURNS_HAVE_NANOFLANN
+#include <nanoflann.hpp>
+#if NANOFLANN_VERSION < 0x150
+namespace nanoflann {
+  using SearchParameters = SearchParams;
+}
+#endif
+#endif
 
 BEGIN_NAMESPACE_OPENTURNS
 
 CLASSNAMEINIT(KDTree)
 
 static const Factory<KDTree> Factory_KDTree;
 
+#ifndef OPENTURNS_HAVE_NANOFLANN
 /**
  * @class KDNearestNeighboursFinder
  *
@@ -125,7 +137,7 @@ class KDNearestNeighboursFinder
       Scalar squaredDistanceBoundingBox = 0.0;
       for(UnsignedInteger i = 0; i < dimension; ++i)
       {
-        Scalar difference = std::max(0.0, std::max(x[i] - upperBoundingBox[i], lowerBoundingBox[i] - x[i]));
+        const Scalar difference = std::max(0.0, std::max(x[i] - upperBoundingBox[i], lowerBoundingBox[i] - x[i]));
         squaredDistanceBoundingBox += difference * difference;
       }
       if (squaredDistanceBoundingBox < values_[0])
@@ -176,7 +188,7 @@ class KDNearestNeighboursFinder
       Scalar squaredDistanceBoundingBox = 0.0;
       for(UnsignedInteger i = 0; i < dimension; ++i)
       {
-        Scalar difference = std::max(0.0, std::max(x[i] - upperBoundingBox[i], lowerBoundingBox[i] - x[i]));
+        const Scalar difference = std::max(0.0, std::max(x[i] - upperBoundingBox[i], lowerBoundingBox[i] - x[i]));
         squaredDistanceBoundingBox += difference * difference;
       }
       if (squaredDistanceBoundingBox < values_[0])
@@ -241,6 +253,8 @@ class KDNearestNeighboursFinder
 
 }; /* class KDNearestNeighboursFinder */
 
+#endif
+
 /**
  * @class KDTree
  *
@@ -250,19 +264,13 @@ class KDNearestNeighboursFinder
 /* Default constructor */
 KDTree::KDTree()
   : NearestNeighbourAlgorithmImplementation()
-  , points_(0, 0)
-  , boundingBox_()
-  , tree_()
 {
   // Nothing to do
 }
 
 /* Parameters constructor */
 KDTree::KDTree(const Sample & points)
   : NearestNeighbourAlgorithmImplementation()
-  , points_(0, 0)
-  , boundingBox_()
-  , tree_()
 {
   // Build the tree
   setSample(points);
@@ -274,11 +282,103 @@ Sample KDTree::getSample() const
   return points_;
 }
 
+#ifdef OPENTURNS_HAVE_NANOFLANN
+class KDTreeSampleAdaptor
+{
+public:
+  explicit KDTreeSampleAdaptor(const Sample & points)
+  : data_(points.getImplementation()->data())
+  , size_(points.getSize())
+  , dimension_(points.getDimension())
+  {
+  }
+
+  inline size_t kdtree_get_point_count() const
+  {
+    return size_;
+  }
+
+  inline Scalar kdtree_get_pt(const size_t idx, const size_t dim) const
+  {
+    return data_[dim + idx * dimension_];
+  }
+
+  template <class BBOX>
+  bool kdtree_get_bbox(BBOX & /*bb*/) const
+  {
+    return false;
+  }
+
+private:
+  const Scalar *data_ = nullptr;
+  UnsignedInteger size_ = 0;
+  UnsignedInteger dimension_ = 0;
+};
+
+
+using nano_kd_tree_t = nanoflann::KDTreeSingleIndexAdaptor<
+        nanoflann::L2_Simple_Adaptor<Scalar, KDTreeSampleAdaptor >,
+        KDTreeSampleAdaptor, -1>;
+
+
+// insulation class for nanoflann implementation
+class KDTreeImplementation
+{
+public:
+  explicit KDTreeImplementation(const Sample & points)
+  {
+    const UnsignedInteger dimension = points.getDimension();
+    sampleAdaptor_ = new KDTreeSampleAdaptor(points);
+    nanoflann::KDTreeSingleIndexAdaptorParams indexParameters;
+    indexParameters.leaf_max_size = ResourceMap::GetAsUnsignedInteger("KDTree-leaf_max_size");
+#if NANOFLANN_VERSION >= 0x150
+    indexParameters.n_thread_build = ResourceMap::GetAsUnsignedInteger("KDTree-n_thread_build");
+#endif
+    indexAdaptor_ = new nano_kd_tree_t(dimension, *sampleAdaptor_, indexParameters);
+#if NANOFLANN_VERSION < 0x140
+    // newer versions do this in the ctor
+    indexAdaptor_->buildIndex();
+#endif
+  }
+
+  UnsignedInteger query(const Point & x) const
+  {
+    UnsignedInteger result = 0;
+    Scalar distance = 0.0;
+    nanoflann::KNNResultSet<Scalar, UnsignedInteger, UnsignedInteger> resultSet(1);
+    resultSet.init(&result, &distance);
+    nanoflann::SearchParameters searchParameters;
+    indexAdaptor_->findNeighbors(resultSet, x.data(), searchParameters);
+    return result;
+  }
+
+  Indices queryK(const Point & x, const UnsignedInteger k, const Bool sorted) const
+  {
+    Indices result(k);
+    Point distance(k);
+    nanoflann::KNNResultSet<Scalar, UnsignedInteger, UnsignedInteger> resultSet(k);
+    resultSet.init(const_cast<UnsignedInteger*>(result.data()), const_cast<Scalar*>(distance.data()));
+    nanoflann::SearchParameters searchParameters;
+    searchParameters.sorted = sorted;
+    indexAdaptor_->findNeighbors(resultSet, x.data(), searchParameters);
+    return result;
+  }
+
+private:
+  Pointer<KDTreeSampleAdaptor> sampleAdaptor_;
+  Pointer<nano_kd_tree_t> indexAdaptor_;
+};
+
+#endif
+
 void KDTree::setSample(const Sample & points)
 {
   if (points == points_) return;
 
   points_ = points;
+#ifdef OPENTURNS_HAVE_NANOFLANN
+  p_implementation_ = new KDTreeImplementation(points);
+#else
   boundingBox_ = Interval(points_.getDimension());
   tree_ = Indices(3 * (points_.getSize() + 1));
 
@@ -297,6 +397,7 @@ void KDTree::setSample(const Sample & points)
   }
   boundingBox_.setLowerBound(points_.getMin());
   boundingBox_.setUpperBound(points_.getMax());
+#endif
 }
 
 /* Virtual constructor */
@@ -358,10 +459,15 @@ UnsignedInteger KDTree::query(const Point & x) const
   if (!points_.getSize())
     throw InvalidArgumentException(HERE) << "Cannot query KDTree with no points";
   if (points_.getSize() == 1) return 0;
+
+#ifdef OPENTURNS_HAVE_NANOFLANN
+  return p_implementation_->query(x);
+#else
   Scalar smallestDistance = SpecFunc::MaxScalar;
   Point lowerBoundingBox(boundingBox_.getLowerBound());
   Point upperBoundingBox(boundingBox_.getUpperBound());
   return getNearestNeighbourIndex(1, x, smallestDistance, lowerBoundingBox, upperBoundingBox, 0);
+#endif
 }
 
 UnsignedInteger KDTree::getNearestNeighbourIndex(const UnsignedInteger inode,
@@ -494,8 +600,12 @@ Indices KDTree::queryK(const Point & x, const UnsignedInteger k, const Bool sort
   }
   else
   {
+#ifdef OPENTURNS_HAVE_NANOFLANN
+    result = p_implementation_->queryK(x, k, sorted);
+#else
     KDNearestNeighboursFinder heap(tree_, points_, boundingBox_, k);
     result = heap.getNearestNeighboursIndices(1, x, sorted);
+#endif
   }
   return result;
 }

lib/src/Base/Algo/openturns/KDTree.hxx

@@ -27,7 +27,7 @@
 
 BEGIN_NAMESPACE_OPENTURNS
 
-class KDNearestNeighboursFinder;
+class KDTreeImplementation;
 
 /**
  * @class KDTree
@@ -39,8 +39,6 @@ class OT_API KDTree
 {
   CLASSNAME
 
-  friend class KDNearestNeighboursFinder;
-
 public:
 
   /** Default constructor */
@@ -104,6 +102,7 @@ private:
   /** The tree, stored as a list of tuples (index, leftNode, rightNode) */
   Indices tree_;
 
+  Pointer<KDTreeImplementation> p_implementation_;
 }; /* class KDTree */
 
 

lib/src/Base/Common/PlatformInfo.cxx

@@ -120,6 +120,12 @@ Description PlatformInfo::GetFeatures()
     Features_["primesieve"] = false;
 #endif
 
+#ifdef OPENTURNS_HAVE_NANOFLANN
+    Features_["nanoflann"] = true;
+#else
+    Features_["nanoflann"] = false;
+#endif
+
 #ifdef OPENTURNS_HAVE_TBB
     Features_["tbb"] = true;
 #else

lib/src/Base/Common/ResourceMap.cxx

@@ -673,6 +673,10 @@ void ResourceMap::loadDefaultConfiguration()
   // DesignProxy parameters
   addAsUnsignedInteger("DesignProxy-DefaultCacheSize", 16777216);// 2^24=16777216=128 Mio
 
+  // KDTree parameters
+  addAsUnsignedInteger("KDTree-leaf_max_size", 10);
+  addAsUnsignedInteger("KDTree-n_thread_build", getAsUnsignedInteger("TBB-ThreadsNumber"));
+
   // KFold parameters
   addAsUnsignedInteger("KFold-DefaultK", 10);
 

lib/src/CMakeLists.txt

@@ -57,6 +57,10 @@ if (HMAT_FOUND)
   set_source_files_properties (Base/Stat/HMatrixFactory.cxx Base/Stat/HMatrixImplementation.cxx PROPERTIES SKIP_UNITY_BUILD_INCLUSION ON)
 endif ()
 
+if (nanoflann_FOUND)
+  set_source_files_properties (Base/Algo/KDTree.cxx PROPERTIES SKIP_UNITY_BUILD_INCLUSION ON)
+endif ()
+
 # SymbolicParserExprTk.cxx includes exprtk.hpp which is huge, exclude it from unity builds
 # It also causes problems on Windows
 if (OPENTURNS_HAVE_EXPRTK)

lib/test/t_KDTree_std.cxx

@@ -45,7 +45,7 @@ int main(int, char *[])
 
   const Sample sample(Normal(3).getSample(10));
   const KDTree tree(sample);
-  fullprint << "tree=" << tree << std::endl;
+  // fullprint << "tree=" << tree << std::endl;
   const Sample test(Normal(3).getSample(20));
   for (UnsignedInteger i = 0; i < test.getSize(); ++i)
   {

lib/test/t_KDTree_std.expout

@@ -1,4 +1,3 @@
-tree=class=KDTree root=class=KDNode index=5 left=class=KDNode index=4 left=NULL right=class=KDNode index=6 left=NULL right=class=KDNode index=2 left=NULL right=NULL right=class=KDNode index=7 left=class=KDNode index=0 left=NULL right=class=KDNode index=1 left=NULL right=NULL right=class=KDNode index=9 left=NULL right=class=KDNode index=8 left=NULL right=class=KDNode index=3 left=NULL right=NULL
 Nearest neighbour of class=Point name=Unnamed dimension=3 values=[-0.721533,-0.241223,-1.78796]=class=Point name=Unnamed dimension=3 values=[0.445785,-1.03808,-0.856712] (index=7)
 Nearest neighbour of class=Point name=Unnamed dimension=3 values=[0.40136,1.36783,1.00434]=class=Point name=Unnamed dimension=3 values=[-0.355007,1.43725,0.810668] (index=2)
 Nearest neighbour of class=Point name=Unnamed dimension=3 values=[0.741548,-0.0436123,0.539345]=class=Point name=Unnamed dimension=3 values=[0.473617,-0.125498,0.351418] (index=8)

python/doc/developer_guide/architecture.rst

@@ -235,6 +235,8 @@ The tools chosen for the development of the platform are:
 +---------------------------------------+---------------------------------------------------------------------------------+-------------------+
 | Multithreading (optional)             | `TBB <http://www.threadingbuildingblocks.org/>`_                                | 2017              |
 +---------------------------------------+---------------------------------------------------------------------------------+-------------------+
+| Nearest neighbor search (optional)    | `nanoflann <https://github.com/jlblancoc/nanoflann>_`                           | 1.3.2             |
++---------------------------------------+---------------------------------------------------------------------------------+-------------------+
 | Python support                        | `Python <http://www.python.org/>`_                                              | 3.6               |
 +---------------------------------------+---------------------------------------------------------------------------------+-------------------+
 | Plotting library (optional)           | `Matplotlib <http://matplotlib.org/>`_                                          | 3.0               |

python/src/KDTree_doc.i.in

@@ -3,9 +3,6 @@
 
 Allows one to store and search points fast.
 
-Available constructors:
-    KDTree(*sample*)
-
 Parameters
 ----------
 sample : 2-d sequence of float
@@ -15,6 +12,14 @@ See also
 --------
 NearestNeighbourAlgorithm
 
+Notes
+-----
+When nanoflann support is enabled, the :class:`~openturns.ResourceMap` key `KDTree-leaf_max_size`
+allows one to set the tree leaf size which involves a build vs query tradeoff: large values
+will tend to result in fast build and slow queries, and small values typically result in slow build and fast queries.
+Also when nanoflann version is at least v1.5.0, the :class:`~openturns.ResourceMap` key `KDTree-n_thread_build`
+allows one to set the number of threads used during the tree building phase. It is also decided by OPENTURNS_NUM_THREADS.
+
 Examples
 --------
 >>> import openturns as ot