add polyline in polygon (#13)

* add

* init

* works

* add polyline in polygon

* add polyline in polygon

* fix

src/main.cpp

@@ -16,6 +16,7 @@
 #include <pybind11/pybind11.h>
 
 #include "fast_crossing.hpp"
+#include "polyline_in_polygon.hpp"
 #include "pybind11_fast_crossing.hpp"
 #include "pybind11_flatbush.hpp"
 #include "pybind11_nanoflann_kdtree.hpp"
@@ -52,6 +53,8 @@ PYBIND11_MODULE(_pybind11_fast_crossing, m)
     m.def("densify_polyline", &cubao::densify_polyline, //
           "polyline"_a, py::kw_only(), "max_gap"_a,
           "densify polyline, interpolate to satisfy max_gap");
+    m.def("polyline_in_polygon", &cubao::polyline_in_polygon, //
+          "polyline"_a, "polygon"_a, py::kw_only(), "is_wgs84"_a = false);
 
 #ifdef VERSION_INFO
     m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO);

src/polyline_in_polygon.hpp

@@ -0,0 +1,70 @@
+#ifndef CUBAO_POLYLINE_IN_POLYGON
+#define CUBAO_POLYLINE_IN_POLYGON
+
+#include "fast_crossing.hpp"
+#include "point_in_polygon.hpp"
+
+namespace cubao
+{
+using PolylineChunks = std::map<std::tuple<int,    // seg_idx
+                                           double, // t
+                                           double, // range
+                                           int,    // seg_idx,
+                                           double, // t
+                                           double  // range
+                                           >,
+                                RowVectors>;
+PolylineChunks
+polyline_in_polygon(const RowVectors &polyline, //
+                    const Eigen::Ref<const RowVectorsNx2> &polygon,
+                    bool is_wgs84 = false)
+{
+    auto fc = FastCrossing(is_wgs84);
+    fc.add_polyline(polygon);
+    auto intersections = fc.intersections(polyline);
+    // pt, (t, s), cur_label=(poly1, seg1), tree_label=(poly2, seg2)
+    auto ruler = PolylineRuler(polyline, is_wgs84);
+    // 0.0, [r1, r2, ..., length]
+    const int N = intersections.size() + 1;
+    Eigen::VectorXd ranges(N);
+    int idx = -1;
+    for (auto &inter : intersections) {
+        int seg_idx = std::get<2>(inter)[1];
+        double t = std::get<1>(inter)[0];
+        double r = ruler.range(seg_idx, t);
+        ranges[++idx] = r;
+    }
+    ranges[++idx] = ruler.length();
+    RowVectorsNx2 midpoints(N, 3);
+    {
+        idx = 0;
+        double r = 0.0;
+        while (idx < N) {
+            double rr = ranges[idx];
+            midpoints.row(idx) = ruler.at((r + rr) / 2.0).head(2);
+            r = rr;
+            ++idx;
+        }
+    }
+    auto mask = point_in_polygon(midpoints, polygon);
+    PolylineChunks ret;
+    {
+        idx = 0;
+        double r = 0.0;
+        while (idx < N) {
+            if (mask[idx] && ranges[idx] > r) {
+                auto [seg1, t1] = ruler.segment_index_t(r);
+                auto [seg2, t2] = ruler.segment_index_t(ranges[idx]);
+                ret.emplace(std::make_tuple(seg1, t1, r, seg2, t2, ranges[idx]),
+                            ruler.lineSliceAlong(r, ranges[idx]));
+            }
+            r = ranges[idx];
+            ++idx;
+        }
+    }
+    return ret;
+}
+
+} // namespace cubao
+
+#endif

tests/test_basic.py

@@ -12,6 +12,7 @@
     Quiver,
     densify_polyline,
     point_in_polygon,
+    polyline_in_polygon,
     tf,
 )
 
@@ -764,6 +765,62 @@ def test_nearst_wgs84():
     # assert len(idx) == 5
 
 
+def test_polyline_in_polygon():
+    """
+           2    4
+           *    *
+      A   /|   /|
+       o---+--/-+--------------o D
+       |/  | /  |              |
+       /   |/   |              |
+      /|   *    *              |
+    1* |   3    5              |
+      Bo-----------------------o
+                               C
+    """
+    polygon_ABCD = np.array(
+        [
+            [0.0, 0.0],  # A
+            [0.0, -10.0],  # B
+            [20.0, -10.0],  # C
+            [20.0, 0.0],  # D
+            [0.0, 0.0],  # A
+        ]
+    )
+    polyline_12345 = np.array(
+        [
+            [-2.0, -9.0, 0.0],  # 1
+            [3.0, 7.0, 1.0],  # 2
+            [3.0, -7.0, 2.0],  # 3
+            [8.0, 7.0, 3.0],  # 4
+            [8.0, -7.0, 4.0],  # 5
+        ]
+    )
+    chunks = polyline_in_polygon(polyline_12345, polygon_ABCD)
+    ranges = []
+    for (seg1, t1, r1, seg2, t2, r2), polyline in chunks.items():
+        ranges.append(r2 - r1)
+        print(f"\n#length: {r2 - r1:.3f}")
+        print(seg1, t1, r1)
+        print(seg2, t2, r2)
+        print(polyline)
+    expected_ranges = [2.72883, 14.4676666, 7.01783]
+    np.testing.assert_allclose(ranges, expected_ranges, atol=1e-4)
+
+    anchor_lla = [123.4, 56.7, 8.9]
+    chunks = polyline_in_polygon(
+        tf.enu2lla(polyline_12345, anchor_lla=anchor_lla),
+        tf.enu2lla(
+            np.c_[polygon_ABCD, np.zeros(len(polygon_ABCD))], anchor_lla=anchor_lla
+        )[:, :2],
+        is_wgs84=True,
+    )
+    ranges = []
+    for _, _, r1, _, _, r2 in chunks.keys():
+        ranges.append(r2 - r1)
+    np.testing.assert_allclose(ranges, expected_ranges, atol=1e-4)
+
+
 def pytest_main(dir: str, *, test_file: str = None):
 
     os.chdir(dir)
@@ -782,5 +839,6 @@ def pytest_main(dir: str, *, test_file: str = None):
 
 
 if __name__ == "__main__":
+    np.set_printoptions(suppress=True)
     pwd = os.path.abspath(os.path.dirname(__file__))
     pytest_main(pwd, test_file=os.path.basename(__file__))