Integrate KdQuiver, tested nearest with filter options (#12)

* sorted->sort, misc updates

* export quiver & bush

* export arrows

* before add filter logic

* add xyz filter

* filter by xyz

* filter

* filter by angle

* okay

* integrate nearest filter

* fix position

* fix nearest

* release

* update upwards

* fix test

Author: district10

docs/about/release-notes.md

@@ -10,6 +10,10 @@ To upgrade `fast-crossing` to the latest version, use pip:
 pip install -U fast-crossing
 ```
 
+## Version 0.0.6 (2023-03-12)
+
+*   Integrate KdQuiver (kdtree + dirs)
+
 ## Version 0.0.5 (2023-03-08)
 
 *   Add nanaflann KdTree, use identical interface to `scipy.spatial.cKDTree`

fast_crossing/spatial.py

@@ -71,8 +71,8 @@ def query_ball_point(
             ii, dd = self.tree.nearest(
                 xyz,
                 radius=r,
+                sort=return_sorted,
                 return_squared_l2=True,
-                sorted=return_sorted,
             )
             if return_length:
                 return len(ii)
@@ -87,8 +87,8 @@ def query_ball_point(
             ii, dd = self.tree.nearest(
                 xyz,
                 radius=rr,
+                sort=return_sorted,
                 return_squared_l2=True,
-                sorted=return_sorted,
             )
             ret_ii.append(ii.tolist())
         if return_length:

setup.py

@@ -124,7 +124,7 @@ def build_extension(self, ext):
 # logic and declaration, and simpler if you include description/version in a file.
 setup(
     name="fast_crossing",
-    version="0.0.5",
+    version="0.0.6",
     author="tzx",
     author_email="dvorak4tzx@gmail.com",
     url="https://fast-crossing.readthedocs.io",

src/fast_crossing.hpp

@@ -76,8 +76,6 @@ struct FastCrossing
             return;
         }
         if (!quiver_) {
-            bush_ = FlatBush();
-            bush_->Finish();
             return;
         }
         auto &polylines = quiver_->polylines();
@@ -562,37 +560,32 @@ struct FastCrossing
         return nearest(point_index(index[0], index[1]), return_squared_l2);
     }
 
-    std::pair<IntNx2, Eigen::VectorXd>
-    nearest(const Eigen::Vector3d &position, //
-            std::optional<int> k = std::nullopt,
-            std::optional<double> radius = std::nullopt,
-            bool sorted = true, //
-            bool return_squared_l2 = false) const
+    std::pair<IntNx2, Eigen::VectorXd> nearest(
+        const Eigen::Vector3d &position, //
+        std::optional<int> k = std::nullopt,
+        std::optional<double> radius = std::nullopt,
+        bool sort = true, //
+        bool return_squared_l2 = false,
+        std::optional<std::pair<Eigen::Vector3d, Quiver::FilterParams>> filter =
+            std::nullopt) const
     {
-        if (k) {
-            auto [ii, dd] =
-                quiver_->nearest(position, *k, sorted, return_squared_l2);
-            auto vec_ii = point_index(ii);
-            return {
-                Eigen::Map<const IntNx2>(vec_ii[0].data(), vec_ii.size(), 2),
-                dd};
-        } else if (radius) {
-            auto [ii, dd] =
-                quiver_->nearest(position, *radius, sorted, return_squared_l2);
-            auto vec_ii = point_index(ii);
-            return {
-                Eigen::Map<const IntNx2>(vec_ii[0].data(), vec_ii.size(), 2),
-                dd};
-        } else {
+        if (!k && !radius) {
             throw std::invalid_argument("should specify k or radius");
         }
+        auto [ii, dd] =
+            k ? quiver_->nearest(position, *k, sort, return_squared_l2)
+              : quiver_->nearest(position, *radius, sort, return_squared_l2);
+        if (filter) {
+            auto [dir, params] = *filter;
+            auto ii_dd = quiver_->filter(ii, dd, Arrow(position, dir), params);
+            ii = ii_dd.first;
+            dd = ii_dd.second;
+        }
+        auto vec_ii = point_index(ii);
+        return {Eigen::Map<const IntNx2>(vec_ii[0].data(), vec_ii.size(), 2),
+                dd};
     }
 
-    const FlatBush &bush() const
-    {
-        finish();
-        return *bush_;
-    }
     bool is_wgs84() const { return is_wgs84_; }
     int num_poylines() const
     {
@@ -614,6 +607,18 @@ struct FastCrossing
         return quiver_->polyline(label);
     }
 
+    const FlatBush *export_bush(bool autobuild = true) const
+    {
+        if (autobuild) {
+            finish();
+        }
+        return bush_ ? &*bush_ : nullptr;
+    }
+    const KdQuiver *export_quiver() const
+    {
+        return quiver_ ? &*quiver_ : nullptr;
+    }
+
   private:
     const bool is_wgs84_{false};
 
@@ -631,6 +636,11 @@ struct FastCrossing
 
     // auto rebuild flatbush
     mutable std::optional<FlatBush> bush_;
+    const FlatBush &bush() const
+    {
+        finish();
+        return *bush_;
+    }
 };
 } // namespace cubao
 

src/kd_quiver.hpp

@@ -72,26 +72,28 @@ struct KdQuiver : Quiver
     std::pair<Eigen::VectorXi, Eigen::VectorXd>
     nearest(const Eigen::Vector3d &position, //
             int k,                           //
-            bool sorted = true,              //
+            bool sort = true,                //
             bool return_squared_l2 = false) const
     {
         return tree().nearest(is_wgs84_ ? lla2enu(position) : position, //
                               k,                                        //
-                              sorted,                                   //
+                              sort,                                     //
                               return_squared_l2);
     }
     std::pair<Eigen::VectorXi, Eigen::VectorXd>
     nearest(const Eigen::Vector3d &position, //
             double radius,                   //
-            bool sorted = true,              //
+            bool sort = true,                //
             bool return_squared_l2 = false) const
     {
         return tree().nearest(is_wgs84_ ? lla2enu(position) : position, //
                               radius,                                   //
-                              sorted,                                   //
+                              sort,                                     //
                               return_squared_l2);
     }
 
+    Eigen::Map<const RowVectors> positions() const { return tree().points(); }
+
     RowVectors positions(const Eigen::VectorXi &hits) const
     {
         const int N = hits.size();
@@ -118,7 +120,7 @@ struct KdQuiver : Quiver
         }
         return coords;
     }
-    std::vector<Arrow> arrows(const Eigen::VectorXi &hits)
+    std::vector<Arrow> arrows(const Eigen::VectorXi &hits) const
     {
         const int N = hits.size();
         std::vector<Arrow> arrows;
@@ -127,41 +129,12 @@ struct KdQuiver : Quiver
             auto _ = index(hits[i]);
             int line_idx = _[0];
             int pt_idx = _[1];
-            auto &ruler = polylines_.at(line_idx);
-            arrows.push_back(Arrow(ruler.at(pt_idx), ruler.dir(pt_idx)));
+            arrows.push_back(arrow(line_idx, pt_idx));
         }
         return arrows;
     }
 
-    static Eigen::VectorXi
-    filter(const std::vector<Arrow> &arrows, const Arrow &arrow,
-           // angle
-           std::optional<Eigen::VectorXd> angle_slots = std::nullopt,
-           // positions
-           std::optional<Eigen::VectorXd> x_slots = std::nullopt,
-           std::optional<Eigen::VectorXd> y_slots = std::nullopt,
-           std::optional<Eigen::VectorXd> z_slots = std::nullopt,
-           bool is_wgs84 = false)
-    {
-        Eigen::VectorXi mask(arrows.size());
-        return mask;
-    }
-
-    Eigen::VectorXi
-    filter(const Eigen::VectorXi &hits, const Arrow &arrow,
-           std::optional<Eigen::VectorXd> angle_slots = std::nullopt,
-           std::optional<Eigen::VectorXd> x_slots = std::nullopt,
-           std::optional<Eigen::VectorXd> y_slots = std::nullopt,
-           std::optional<Eigen::VectorXd> z_slots = std::nullopt)
-    {
-        return filter(arrows(hits), arrow,       //
-                      angle_slots,               //
-                      x_slots, y_slots, z_slots, //
-                      is_wgs84_);
-    }
-
-    // helpers
-    Arrow arrow(const PolylineRuler &ruler, int segment_index, double t)
+    Arrow arrow(const PolylineRuler &ruler, int segment_index, double t) const
     {
 
         auto &polyline_ = ruler.polyline();
@@ -177,11 +150,108 @@ struct KdQuiver : Quiver
         }
         return arrow;
     }
-    Arrow arrow(const PolylineRuler &ruler, double range)
+    Arrow arrow(const PolylineRuler &ruler, double range) const
     {
         auto [seg_idx, t] = ruler.segment_index_t(range);
         return arrow(ruler, seg_idx, t);
     }
+    Arrow arrow(int point_index) const
+    {
+        Eigen::Vector2i ii = this->index(point_index);
+        return arrow(ii[0], ii[1]);
+    }
+    Arrow arrow(int polyline_index, int segment_index) const
+    {
+        auto ruler = polyline(polyline_index);
+        return Arrow(ruler->at(segment_index), ruler->dir(segment_index)) //
+            .polyline_index(polyline_index)                               //
+            .segment_index(segment_index)                                 //
+            .t(0.0)                                                       //
+            .range(ruler->range(segment_index))                           //
+            ;
+    }
+    Arrow arrow(int polyline_index, int segment_index, double t) const
+    {
+        auto ruler = polyline(polyline_index);
+        return Arrow(ruler->at(segment_index, t), ruler->dir(segment_index))
+            .polyline_index(polyline_index)        //
+            .segment_index(segment_index)          //
+            .t(t)                                  //
+            .range(ruler->range(segment_index, t)) //
+            ;
+    }
+    Arrow arrow(int polyline_index, double range) const
+    {
+        auto ruler = polyline(polyline_index);
+        auto [xyz, dir] = ruler->arrow(range, false);
+        auto [idx, t] = ruler->segment_index_t(range);
+        return Arrow(xyz, dir)              //
+            .polyline_index(polyline_index) //
+            .segment_index(idx)             //
+            .t(t)                           //
+            .range(range)                   //
+            ;
+    }
+
+    static Eigen::VectorXi filter(const std::vector<Arrow> &arrows,
+                                  const Arrow &arrow,
+                                  const Quiver::FilterParams &params,
+                                  bool is_wgs84 = false)
+    {
+        if (is_wgs84) {
+            Quiver quiver(arrow.position());
+            return quiver.filter(arrows, arrow, params);
+        } else {
+            Quiver quiver;
+            return quiver.filter(arrows, arrow, params);
+        }
+    }
+
+    static Eigen::VectorXi select_by_mask(const Eigen::VectorXi &indexes,
+                                          const Eigen::VectorXi &mask)
+    {
+        std::vector<int> ret;
+        ret.reserve(mask.sum());
+        for (int i = 0, N = indexes.size(); i < N; ++i) {
+            if (mask[i]) {
+                ret.push_back(indexes[i]);
+            }
+        }
+        return Eigen::VectorXi::Map(&ret[0], ret.size());
+    }
+    static Eigen::VectorXd select_by_mask(const Eigen::VectorXd &norms,
+                                          const Eigen::VectorXi &mask)
+    {
+        std::vector<double> ret;
+        ret.reserve(mask.sum());
+        for (int i = 0, N = norms.size(); i < N; ++i) {
+            if (mask[i]) {
+                ret.push_back(norms[i]);
+            }
+        }
+        return Eigen::VectorXd::Map(&ret[0], ret.size());
+    }
+
+    Eigen::VectorXi filter(const Eigen::VectorXi &hits, //
+                           const Arrow &arrow,          //
+                           const Quiver::FilterParams &params) const
+    {
+        auto mask = filter(arrows(hits), arrow, //
+                           params, is_wgs84_);
+        return select_by_mask(hits, mask);
+    }
+
+    std::pair<Eigen::VectorXi, Eigen::VectorXd>
+    filter(const Eigen::VectorXi &hits,  //
+           const Eigen::VectorXd &norms, //
+           const Arrow &arrow,           //
+           const Quiver::FilterParams &params) const
+    {
+        auto mask = filter(arrows(hits), arrow, //
+                           params, is_wgs84_);
+        return {select_by_mask(hits, mask), //
+                select_by_mask(norms, mask)};
+    }
 
     void reset() { reset_index(); }
 

src/main.cpp

@@ -21,7 +21,10 @@
 #include "pybind11_nanoflann_kdtree.hpp"
 #include "pybind11_quiver.hpp"
 
+#define CUBAO_ARGV_DEFAULT_NONE(argv) py::arg_v(#argv, std::nullopt, "None")
+
 #include "pybind11_polyline_ruler.hpp"
+#include "pybind11_crs_transform.hpp"
 
 #include "point_in_polygon.hpp"
 #include "densify_polyline.hpp"
@@ -40,6 +43,9 @@ PYBIND11_MODULE(_pybind11_fast_crossing, m)
     cubao::bind_quiver(m);
     cubao::bind_polyline_ruler(m);
 
+    auto tf = m.def_submodule("tf");
+    cubao::bind_crs_transform(tf);
+
     m.def("point_in_polygon", &cubao::point_in_polygon, //
           py::kw_only(), "points"_a, "polygon"_a,
           "point-in-polygon test, returns 0-1 mask");

src/nanoflann_kdtree.hpp

@@ -122,15 +122,15 @@ struct KdTree : PointCloud
     std::pair<Eigen::VectorXi, Eigen::VectorXd>
     nearest(const Eigen::Vector3d &position, //
             int k,                           //
-            bool sorted = true,              //
+            bool sort = true,                //
             bool return_squared_l2 = false) const
     {
         auto indexes = std::vector<int>(k);
         auto distances = std::vector<double>(k);
         auto resultSet = nanoflann::KNNResultSet<double, int, int>(k);
         resultSet.init(&indexes[0], &distances[0]);
         auto params = nanoflann::SearchParams();
-        params.sorted = sorted;
+        params.sorted = sort;
         index().findNeighbors(resultSet, position.data(), params);
         const int N = resultSet.size();
         return std::make_pair(
@@ -142,11 +142,11 @@ struct KdTree : PointCloud
     std::pair<Eigen::VectorXi, Eigen::VectorXd>
     nearest(const Eigen::Vector3d &position, //
             double radius,                   //
-            bool sorted = true,              //
+            bool sort = true,                //
             bool return_squared_l2 = false) const
     {
         auto params = nanoflann::SearchParams();
-        params.sorted = sorted;
+        params.sorted = sort;
         std::vector<std::pair<int, double>> indices_dists;
         index().radiusSearch(position.data(), //
                              radius * radius, //