making knn and octree work on more types than double, allowing knn to do nearest neighbor searches between two different point clouds (#1313)

removing a capture issue with gcc 4.8.4 in knn

Co-authored-by: Jérémie Dumas <[email protected]>

include/igl/knn.cpp

@@ -1,88 +1,120 @@
 #include "knn.h"
+#include "sort.h"
 #include "parallel_for.h"
 
 #include <cmath>
 #include <queue>
+#include <set>
+#include <algorithm>
 
 namespace igl {
-  template <typename DerivedP, typename KType, typename IndexType,
+  template <typename DerivedP, typename IndexType,
   typename DerivedCH, typename DerivedCN, typename DerivedW,
   typename DerivedI>
   IGL_INLINE void knn(const Eigen::MatrixBase<DerivedP>& P,
-                      const KType & k,
+                      size_t k,
                       const std::vector<std::vector<IndexType> > & point_indices,
                       const Eigen::MatrixBase<DerivedCH>& CH,
                       const Eigen::MatrixBase<DerivedCN>& CN,
                       const Eigen::MatrixBase<DerivedW>& W,
-                      Eigen::PlainObjectBase<DerivedI> & I)
-  {
+                      Eigen::PlainObjectBase<DerivedI> & I) {
+      knn(P,P,k,point_indices,CH,CN,W,I);
+  }
+
+  template <typename DerivedP, typename DerivedV, typename IndexType,
+  typename DerivedCH, typename DerivedCN, typename DerivedW,
+  typename DerivedI>
+      IGL_INLINE void knn(
+              const Eigen::MatrixBase<DerivedP>& P,
+              const Eigen::MatrixBase<DerivedV>& V,
+              size_t k,
+              const std::vector<std::vector<IndexType> > & point_indices,
+              const Eigen::MatrixBase<DerivedCH>& CH,
+              const Eigen::MatrixBase<DerivedCN>& CN,
+              const Eigen::MatrixBase<DerivedW>& W,
+              Eigen::PlainObjectBase<DerivedI> & I) {
     typedef typename DerivedCN::Scalar CentersType;
     typedef typename DerivedW::Scalar WidthsType;
-    
-    typedef Eigen::Matrix<typename DerivedP::Scalar, 1, 3> RowVector3PType;
-    
-    int n = P.rows();
-    const KType real_k = std::min(n,k);
-    
-    auto distance_to_width_one_cube = [](RowVector3PType point){
-      return std::sqrt(std::pow(std::max(std::abs(point(0))-1,0.0),2)
-                       + std::pow(std::max(std::abs(point(1))-1,0.0),2)
-                       + std::pow(std::max(std::abs(point(2))-1,0.0),2));
+
+    using Scalar = typename DerivedP::Scalar;
+    typedef Eigen::Matrix<Scalar, 1, 3> RowVector3PType;
+
+
+    const size_t Psize = P.rows();
+    const size_t Vsize = V.rows();
+    if(Vsize <= k) {
+        I.resize(Psize,Vsize);
+        for(size_t i = 0; i < Psize; ++i) {
+            Eigen::Matrix<Scalar,Eigen::Dynamic,1> D = (V.rowwise() - P.row(i)).rowwise().norm();
+            Eigen::Matrix<Scalar,Eigen::Dynamic,1> S;
+            Eigen::VectorXi R;
+            igl::sort(D,1,true,S,R);
+            I.row(i) = R.transpose();
+        }
+        return;
+    }
+
+    I.resize(Psize,k);
+
+
+    auto distance_to_width_one_cube = [](const RowVector3PType& point) -> Scalar {
+      return std::sqrt(std::pow<Scalar>(std::max<Scalar>(std::abs(point(0))-1,0.0),2)
+                       + std::pow<Scalar>(std::max<Scalar>(std::abs(point(1))-1,0.0),2)
+                       + std::pow<Scalar>(std::max<Scalar>(std::abs(point(2))-1,0.0),2));
     };
-    
+
     auto distance_to_cube = [&distance_to_width_one_cube]
-              (RowVector3PType point,
+              (const RowVector3PType& point,
                Eigen::Matrix<CentersType,1,3> cube_center,
-               WidthsType cube_width){
+               WidthsType cube_width) -> Scalar {
       RowVector3PType transformed_point = (point-cube_center)/cube_width;
       return cube_width*distance_to_width_one_cube(transformed_point);
     };
-    
-    I.resize(n,real_k);
-    
-    igl::parallel_for(n,[&](int i)
+
+
+    igl::parallel_for(Psize,[&](size_t i)
     {
       int points_found = 0;
       RowVector3PType point_of_interest = P.row(i);
-      
+
       //To make my priority queue take both points and octree cells,
       //I use the indices 0 to n-1 for the n points,
       // and the indices n to n+m-1 for the m octree cells
-      
+
       // Using lambda to compare elements.
-      auto cmp = [&point_of_interest, &P, &CN, &W,
-                  &n, &distance_to_cube](int left, int right) {
-        double leftdistance, rightdistance;
-        if(left < n){ //left is a point index
-          leftdistance = (P.row(left) - point_of_interest).norm();
+      auto cmp = [&point_of_interest, &V, &CN, &W,
+                  Vsize, &distance_to_cube](int left, int right) {
+        Scalar leftdistance, rightdistance;
+        if(left < Vsize){ //left is a point index
+          leftdistance = (V.row(left) - point_of_interest).norm();
         } else { //left is an octree cell
           leftdistance = distance_to_cube(point_of_interest,
-                                            CN.row(left-n),
-                                            W(left-n));
+                                            CN.row(left-Vsize),
+                                            W(left-Vsize));
         }
-      
-        if(right < n){ //left is a point index
-          rightdistance = (P.row(right) - point_of_interest).norm();
+
+        if(right < Vsize){ //left is a point index
+          rightdistance = (V.row(right) - point_of_interest).norm();
         } else { //left is an octree cell
           rightdistance = distance_to_cube(point_of_interest,
-                                             CN.row(right-n),
-                                             W(right-n));
+                                             CN.row(right-Vsize),
+                                             W(right-Vsize));
         }
         return leftdistance > rightdistance;
       };
-      
+
       std::priority_queue<IndexType, std::vector<IndexType>,
         decltype(cmp)> queue(cmp);
-      
-      queue.push(n); //This is the 0th octree cell (ie the root)
-      while(points_found < real_k){
+
+      queue.push(Vsize); //This is the 0th octree cell (ie the root)
+      while(points_found < k){
         IndexType curr_cell_or_point = queue.top();
         queue.pop();
-        if(curr_cell_or_point < n){ //current index is for is a point
+        if(curr_cell_or_point < Vsize){ //current index is for is a point
           I(i,points_found) = curr_cell_or_point;
           points_found++;
         } else {
-          IndexType curr_cell = curr_cell_or_point - n;
+          IndexType curr_cell = curr_cell_or_point - Vsize;
           if(CH(curr_cell,0) == -1){ //In the case of a leaf
             if(point_indices.at(curr_cell).size() > 0){
               //Assumption: Leaves either have one point, or none
@@ -91,7 +123,7 @@ namespace igl {
           } else { //Not a leaf
             for(int j = 0; j < 8; j++){
               //+n to adjust for the octree cells
-              queue.push(CH(curr_cell,j)+n);
+              queue.push(CH(curr_cell,j)+Vsize);
             }
           }
         }
@@ -105,6 +137,12 @@ namespace igl {
 #ifdef IGL_STATIC_LIBRARY
 // Explicit template instantiation
 // generated by autoexplicit.sh
-template void igl::knn<Eigen::Matrix<double, -1, -1, 0, -1, -1>, int, int, Eigen::Matrix<int, -1, -1, 0, -1, -1>, Eigen::Matrix<double, -1, -1, 0, -1, -1>, Eigen::Matrix<double, -1, 1, 0, -1, 1>, Eigen::Matrix<int, -1, -1, 0, -1, -1> >(Eigen::MatrixBase<Eigen::Matrix<double, -1, -1, 0, -1, -1> > const&, int const&, std::vector<std::vector<int, std::allocator<int> >, std::allocator<std::vector<int, std::allocator<int> > > > const&, Eigen::MatrixBase<Eigen::Matrix<int, -1, -1, 0, -1, -1> > const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, -1, 0, -1, -1> > const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, 1, 0, -1, 1> > const&, Eigen::PlainObjectBase<Eigen::Matrix<int, -1, -1, 0, -1, -1> >&);
-template void igl::knn<Eigen::Matrix<double, -1, -1, 0, -1, -1>, int, int, Eigen::Matrix<int, -1, 8, 0, -1, 8>, Eigen::Matrix<double, -1, 3, 0, -1, 3>, Eigen::Matrix<double, -1, 1, 0, -1, 1>, Eigen::Matrix<int, -1, -1, 0, -1, -1> >(Eigen::MatrixBase<Eigen::Matrix<double, -1, -1, 0, -1, -1> > const&, int const&, std::vector<std::vector<int, std::allocator<int> >, std::allocator<std::vector<int, std::allocator<int> > > > const&, Eigen::MatrixBase<Eigen::Matrix<int, -1, 8, 0, -1, 8> > const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, 3, 0, -1, 3> > const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, 1, 0, -1, 1> > const&, Eigen::PlainObjectBase<Eigen::Matrix<int, -1, -1, 0, -1, -1> >&);
+
+template void igl::knn<Eigen::Matrix<double, -1, -1, 0, -1, -1>, Eigen::Matrix<double, -1, -1, 0, -1, -1>, int, Eigen::Matrix<int, -1, 8, 0, -1, 8>, Eigen::Matrix<double, -1, 3, 0, -1, 3>, Eigen::Matrix<double, -1, 1, 0, -1, 1>, Eigen::Matrix<int, -1, -1, 0, -1, -1> >(Eigen::MatrixBase<Eigen::Matrix<double, -1, -1, 0, -1, -1> > const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, -1, 0, -1, -1> > const&, unsigned long, std::vector<std::vector<int, std::allocator<int> >, std::allocator<std::vector<int, std::allocator<int> > > > const&, Eigen::MatrixBase<Eigen::Matrix<int, -1, 8, 0, -1, 8> > const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, 3, 0, -1, 3> > const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, 1, 0, -1, 1> > const&, Eigen::PlainObjectBase<Eigen::Matrix<int, -1, -1, 0, -1, -1> >&);
+template void igl::knn<Eigen::Matrix<double, -1, -1, 0, -1, -1>, int, Eigen::Matrix<int, -1, -1, 0, -1, -1>, Eigen::Matrix<double, -1, -1, 0, -1, -1>, Eigen::Matrix<double, -1, 1, 0, -1, 1>, Eigen::Matrix<int, -1, -1, 0, -1, -1> >(Eigen::MatrixBase<Eigen::Matrix<double, -1, -1, 0, -1, -1> > const&, unsigned long, std::vector<std::vector<int, std::allocator<int> >, std::allocator<std::vector<int, std::allocator<int> > > > const&, Eigen::MatrixBase<Eigen::Matrix<int, -1, -1, 0, -1, -1> > const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, -1, 0, -1, -1> > const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, 1, 0, -1, 1> > const&, Eigen::PlainObjectBase<Eigen::Matrix<int, -1, -1, 0, -1, -1> >&);
+#ifdef WIN32
+template void igl::knn<Eigen::Matrix<double,-1,-1,0,-1,-1>,int,Eigen::Matrix<int,-1,-1,0,-1,-1>,Eigen::Matrix<double,-1,-1,0,-1,-1>,Eigen::Matrix<double,-1,1,0,-1,1>,Eigen::Matrix<int,-1,-1,0,-1,-1> >(Eigen::MatrixBase<Eigen::Matrix<double,-1,-1,0,-1,-1> > const &,unsigned __int64,std::vector<std::vector<int,std::allocator<int> >,std::allocator<std::vector<int,std::allocator<int> > > > const &,Eigen::MatrixBase<Eigen::Matrix<int,-1,-1,0,-1,-1> > const &,Eigen::MatrixBase<Eigen::Matrix<double,-1,-1,0,-1,-1> > const &,Eigen::MatrixBase<Eigen::Matrix<double,-1,1,0,-1,1> > const &,Eigen::PlainObjectBase<Eigen::Matrix<int,-1,-1,0,-1,-1> > &);
+template void igl::knn<Eigen::Matrix<double,-1,-1,0,-1,-1>,Eigen::Matrix<double,-1,-1,0,-1,-1>,int,Eigen::Matrix<int,-1,8,0,-1,8>,Eigen::Matrix<double,-1,3,0,-1,3>,Eigen::Matrix<double,-1,1,0,-1,1>,Eigen::Matrix<int,-1,-1,0,-1,-1> >(Eigen::MatrixBase<Eigen::Matrix<double,-1,-1,0,-1,-1> > const &,Eigen::MatrixBase<Eigen::Matrix<double,-1,-1,0,-1,-1> > const &,unsigned __int64,std::vector<std::vector<int,std::allocator<int> >,std::allocator<std::vector<int,std::allocator<int> > > > const &,Eigen::MatrixBase<Eigen::Matrix<int,-1,8,0,-1,8> > const &,Eigen::MatrixBase<Eigen::Matrix<double,-1,3,0,-1,3> > const &,Eigen::MatrixBase<Eigen::Matrix<double,-1,1,0,-1,1> > const &,Eigen::PlainObjectBase<Eigen::Matrix<int,-1,-1,0,-1,-1> > &);
+#endif
+
 #endif

include/igl/knn.h

@@ -34,11 +34,40 @@ namespace igl
   //          of the ith octree cell
   // Outputs:
   //   I  #P by k list of k-nearest-neighbor indices into P
-  template <typename DerivedP, typename KType, typename IndexType,
+  template <typename DerivedP, typename IndexType,
     typename DerivedCH, typename DerivedCN, typename DerivedW,
     typename DerivedI>
-  IGL_INLINE void knn(const Eigen::MatrixBase<DerivedP>& P,
-    const KType & k,
+  IGL_INLINE void knn(
+          const Eigen::MatrixBase<DerivedP>& P,
+    size_t k,
+    const std::vector<std::vector<IndexType> > & point_indices,
+    const Eigen::MatrixBase<DerivedCH>& CH,
+    const Eigen::MatrixBase<DerivedCN>& CN,
+    const Eigen::MatrixBase<DerivedW>& W,
+    Eigen::PlainObjectBase<DerivedI> & I);
+
+
+  // Inputs:
+  //   P  #P by 3 list of point locations for which which we want the neighbors of
+  //   V  #V by 3 list of point locations for which may be neighbors 
+  //   k  number of neighbors to find
+  //   point_indices  a vector of vectors, where the ith entry is a vector of
+  //                  the indices into P that are the ith octree cell's points
+  //   CH     #OctreeCells by 8, where the ith row is the indices of
+  //          the ith octree cell's children
+  //   CN     #OctreeCells by 3, where the ith row is a 3d row vector
+  //          representing the position of the ith cell's center
+  //   W      #OctreeCells, a vector where the ith entry is the width
+  //          of the ith octree cell
+  // Outputs:
+  //   I  #P by k list of k-nearest-neighbor indices into V
+  template <typename DerivedP, typename DerivedV, typename IndexType,
+    typename DerivedCH, typename DerivedCN, typename DerivedW,
+    typename DerivedI>
+  IGL_INLINE void knn(
+          const Eigen::MatrixBase<DerivedP>& P,
+          const Eigen::MatrixBase<DerivedV>& V,
+    size_t k,
     const std::vector<std::vector<IndexType> > & point_indices,
     const Eigen::MatrixBase<DerivedCH>& CH,
     const Eigen::MatrixBase<DerivedCN>& CN,
@@ -49,4 +78,4 @@ namespace igl
 #  include "knn.cpp"
 #endif
 #endif
-
+

include/igl/octree.cpp

@@ -19,8 +19,9 @@ namespace igl {
     typedef typename DerivedCH::Scalar ChildrenType;
     typedef typename DerivedCN::Scalar CentersType;
     typedef typename DerivedW::Scalar WidthsType;
+    typedef typename DerivedP::Scalar PointScalar;
     typedef Eigen::Matrix<ChildrenType,8,1> Vector8i;
-    typedef Eigen::Matrix<typename DerivedP::Scalar, 1, 3> RowVector3PType;
+    typedef Eigen::Matrix<PointScalar, 1, 3> RowVector3PType;
     typedef Eigen::Matrix<CentersType, 1, 3>       RowVector3CentersType;
     
     std::vector<Eigen::Matrix<ChildrenType,8,1>,
@@ -29,8 +30,8 @@ namespace igl {
         Eigen::aligned_allocator<Eigen::Matrix<CentersType,1,3> > > centers;
     std::vector<WidthsType> widths;
     
-    auto get_octant = [](RowVector3PType location,
-                         RowVector3CentersType center){
+    auto get_octant = [](const RowVector3PType& location,
+                         const RowVector3CentersType& center){
       // We use a binary numbering of children. Treating the parent cell's
       // center as the origin, we number the octants in the following manner:
       // The first bit is 1 iff the octant's x coordinate is positive
@@ -85,7 +86,7 @@ namespace igl {
   
     // Useful list of number 0..7
     const Vector8i zero_to_seven = (Vector8i()<<0,1,2,3,4,5,6,7).finished();
-    const Vector8i neg_ones = (Vector8i()<<-1,-1,-1,-1,-1,-1,-1,-1).finished();
+    const Vector8i neg_ones = Vector8i::Constant(-1);
   
     std::function< void(const ChildrenType, const int) > helper;
     helper = [&helper,&translate_center,&get_octant,&m,
@@ -138,17 +139,10 @@ namespace igl {
     children.emplace_back(neg_ones);
   
     //Get the minimum AABB for the points
-    RowVector3PType backleftbottom(P.col(0).minCoeff(),
-                                   P.col(1).minCoeff(),
-                                   P.col(2).minCoeff());
-    RowVector3PType frontrighttop(P.col(0).maxCoeff(),
-                                  P.col(1).maxCoeff(),
-                                  P.col(2).maxCoeff());
-    RowVector3CentersType aabb_center = (backleftbottom+frontrighttop)/2.0;
-    WidthsType aabb_width = std::max(std::max(
-                                          frontrighttop(0) - backleftbottom(0),
-                                          frontrighttop(1) - backleftbottom(1)),
-                                          frontrighttop(2) - backleftbottom(2));
+    RowVector3PType backleftbottom = P.colwise().minCoeff();
+    RowVector3PType frontrighttop = P.colwise().maxCoeff();
+    RowVector3CentersType aabb_center = (backleftbottom+frontrighttop)/PointScalar(2.0);
+    WidthsType aabb_width = (frontrighttop - backleftbottom).maxCoeff();
     centers.emplace_back( aabb_center );
   
     //Widths are the side length of the cube, (not half the side length):

tests/include/igl/knn.cpp

@@ -0,0 +1,78 @@
+#include <test_common.h>
+#include <igl/list_to_matrix.h>
+#include <igl/knn.h>
+#include <igl/octree.h>
+
+
+
+TEST_CASE("knn", "[igl]")
+{
+
+    auto run_knn = [](int k, const Eigen::MatrixXd & P, const Eigen::MatrixXd& V, typename Eigen::MatrixXi& I) {
+        std::vector<std::vector<int> > point_indices;
+        Eigen::Matrix<int,Eigen::Dynamic,8> CH;
+        Eigen::Matrix<double,Eigen::Dynamic,3> CN;
+        Eigen::Matrix<double,Eigen::Dynamic,1> W;
+        Eigen::Matrix<double,Eigen::Dynamic,1> A;
+
+        igl::octree(V,point_indices,CH,CN,W);
+
+        igl::knn(P,V,k,point_indices,CH,CN,W,I);
+    };
+    Eigen::Matrix<double,Eigen::Dynamic,3> V;
+    Eigen::Matrix<double,Eigen::Dynamic,3> P;
+    Eigen::MatrixXi I;
+    Eigen::MatrixXi answer;
+    {
+        //Test some simple points off a unit cube
+        V.resize(8,3);
+        V << 
+            0,0,1,
+            0,1,0,
+            0,1,1,
+            0,0,0,
+            1,0,0,
+            1,1,0,
+            1,1,1,
+            1,0,1;
+
+        P.resize(3,3);
+
+        P << 
+            0 ,0 ,0.6,
+              0.3 ,0.1 ,0.2,
+              .7,.6,0;
+
+
+        answer.resize(3,3);
+        answer << 0,3,2,
+               3,4,0,
+               5,4,1;
+        run_knn(3,P,V,I);
+        REQUIRE (I == answer);
+
+    }
+    {
+        //Test whether the kdtree works when passed things of different size
+
+        V.resize(2,3);
+        V << 0,0,0,
+          1,1,1;
+
+        P << 0,0,0,
+            -1,-1,-1,
+            2,2,2;
+
+        run_knn(10,P,V,I);
+        answer.resize(3,2);
+        answer << 
+            0,1,
+            0,1,
+            1,0;
+
+        REQUIRE (I == answer);
+    }
+
+
+
+}