fix packedrtree.hpp (#13)

* not ready * not ready * not ready * packedrtree works * fix !!! * fix * fix --------- Co-authored-by: TANG ZHIXIONG <[email protected]>
cubao · Oct 21, 2023 · b65d359 · b65d359
1 parent 6e79f7c
commit b65d359
Show file tree

Hide file tree

Showing 13 changed files with 925 additions and 20 deletions.
diff --git a/3rdparty/naive_svg.hpp b/3rdparty/naive_svg.hpp
diff --git a/3rdparty/packedrtree.hpp b/3rdparty/packedrtree.hpp
@@ -1,5 +1,6 @@
 // sync with
 //  https://github.com/cubao/headers/blob/main/include/packedrtree.hpp
+//  https://github.com/cubao/nano-fmm/blob/master/3rdparty/packedrtree.hpp
 
 /******************************************************************************
  *
@@ -249,7 +250,7 @@ inline void hilbertSort(std::vector<NodeItem> &items)
 class PackedRTree
 {
     NodeItem _extent;
-    NodeItem *_nodeItems = nullptr;
+    std::vector<NodeItem> _nodeItems;
     uint64_t _numItems;
     uint64_t _numNodes;
     uint16_t _nodeSize;
@@ -264,7 +265,7 @@ class PackedRTree
                              static_cast<uint16_t>(65535));
         _levelBounds = generateLevelBounds(_numItems, _nodeSize);
         _numNodes = _levelBounds.front().second;
-        _nodeItems = new NodeItem[static_cast<size_t>(_numNodes)];
+        _nodeItems.resize(_numNodes);
     }
     void generateNodes()
     {
@@ -293,11 +294,6 @@ class PackedRTree
     }
 
   public:
-    ~PackedRTree()
-    {
-        if (_nodeItems != nullptr)
-            delete[] _nodeItems;
-    }
     PackedRTree(const std::vector<NodeItem> &nodes, const NodeItem &extent,
                 const uint16_t nodeSize = 16)
         : _extent(extent), _numItems(nodes.size())
@@ -407,7 +403,7 @@ class PackedRTree
 
     void streamWrite(const std::function<void(uint8_t *, size_t)> &writeData)
     {
-        writeData(reinterpret_cast<uint8_t *>(_nodeItems),
+        writeData(reinterpret_cast<uint8_t *>(&_nodeItems[0]),
                   static_cast<size_t>(_numNodes * sizeof(NodeItem)));
     }
 

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -37,6 +37,16 @@ endif()
 
 include_directories(3rdparty src)
 
+option(BUILD_EXAMPLES "Build examples" OFF)
+if(BUILD_EXAMPLES)
+  file(GLOB SRCS examples/*.cpp)
+  foreach(src ${SRCS})
+    string(REGEX REPLACE "(^.*/|.cpp$)" "" exe ${src})
+    add_executable(${exe} ${src})
+    install(TARGETS ${exe} RUNTIME DESTINATION bin)
+  endforeach(src)
+endif()
+
 set(PYBIND11_CPP_STANDARD -std=c++17)
 add_subdirectory(pybind11)
 file(GLOB_RECURSE SRCS src/*.cpp)

diff --git a/Makefile b/Makefile
@@ -62,7 +62,6 @@ python_sdist:
 python_sdist_wheel:
 	$(PYTHON) -m pip wheel dist/nano_fmm-*.tar.gz --no-deps
 python_test:
-	$(PYTHON) -c 'from pybind11_rdp import rdp; print(rdp([[1, 1], [2, 2], [3, 3], [4, 4]]))'
 	$(PYTHON) -c 'import nano_fmm; print(nano_fmm.add(1, 2))'
 	$(PYTHON) -m nano_fmm add 1 2
 	$(PYTHON) -m nano_fmm pure_python_func --arg1=43234

diff --git a/examples/demo_packedrtree.cpp b/examples/demo_packedrtree.cpp
@@ -0,0 +1,98 @@
+#define DBG_MACRO_NO_WARNING
+#include <dbg.h>
+
+#include <stdexcept>
+#include <iostream>
+#include <random>
+#include <set>
+#include <unordered_set>
+
+#include "naive_svg.hpp"
+#include "packedrtree.hpp"
+
+using namespace FlatGeobuf;
+using namespace cubao;
+
+size_t unix_time()
+{
+    using namespace std;
+    using namespace std::chrono;
+    return duration_cast<milliseconds>(system_clock::now().time_since_epoch())
+        .count();
+}
+
+int main(int argc, char **argv)
+{
+    const int num_boxes = 800;
+    std::random_device rnd_dev;
+    std::default_random_engine engine(rnd_dev());
+
+    int width = 1000, height = 1000;
+    int padding = 10;
+
+    std::uniform_int_distribution<int> dist_x(padding, width - padding);
+    std::uniform_int_distribution<int> dist_y(padding, width - padding);
+    std::uniform_int_distribution<int> dist_w(1, padding);
+    std::uniform_int_distribution<int> dist_h(1, padding);
+
+    auto nodes = std::vector<NodeItem>();
+    nodes.reserve(num_boxes);
+    for (int i = 0; i < num_boxes; ++i) {
+        double x = (double)dist_x(engine);
+        double y = (double)dist_y(engine);
+        double w = (double)dist_w(engine);
+        double h = (double)dist_h(engine);
+        nodes.push_back({x, y, x + w, y + h, (uint64_t)i});
+    }
+    auto extent = calcExtent(nodes);
+    hilbertSort(nodes, extent);
+    auto rtree = PackedRTree(nodes, extent);
+
+    double xmin = (double)dist_x(engine);
+    double xmax = (double)dist_x(engine);
+    double ymin = (double)dist_y(engine);
+    double ymax = (double)dist_y(engine);
+    if (xmin > xmax) {
+        std::swap(xmin, xmax);
+    }
+    if (ymin > ymax) {
+        std::swap(ymin, ymax);
+    }
+    xmin -= padding;
+    xmax += padding;
+    ymin -= padding;
+    ymax += padding;
+
+    SVG svg(width, height);
+    svg.add_polygon({
+                        {xmin, ymin},
+                        {xmax, ymin},
+                        {xmax, ymax},
+                        {xmin, ymax},
+                        {xmin, ymin},
+                    })
+        .stroke(SVG::COLOR(0xff0000))
+        .stroke_width(3.0);
+
+    std::set<int64_t> hits;
+    for (auto h : rtree.search(xmin, ymin, xmax, ymax)) {
+        hits.insert(h.offset);
+    }
+    dbg(hits.size());
+    for (auto &node : nodes) {
+        auto &box = svg.add_polygon({
+            {node.minX, node.minY},
+            {node.maxX, node.minY},
+            {node.maxX, node.maxY},
+            {node.minX, node.maxY},
+            {node.minX, node.minY},
+        });
+        if (!hits.count(node.offset)) {
+            box.stroke(SVG::Color(0x00ff00));
+        } else {
+            box.stroke(SVG::Color(0xff0000)).fill(SVG::Color(0xff0000).a(0.2));
+        }
+    }
+    svg.dump("packedrtree-" + std::to_string(unix_time()) + ".svg");
+    return 0;
+}
diff --git a/src/bindings/pybind11_network.cpp b/src/bindings/pybind11_network.cpp
@@ -55,8 +55,9 @@ void bind_network(py::module &m)
         .def("__repr__", [](const ProjectedPoint &self) {
             auto &p = self.position();
             auto &d = self.direction();
-            return fmt::format("ProjectedPoint(pos=[{},{},{}],dir=[{},{},{}],"
-                               "dist={},road={},offset={})",
+            return fmt::format("ProjectedPoint(pos=[{:.7f},{:.7f},{:.2f}],"
+                               "dir=[{:.2f},{:.2f},{:.1f}],"
+                               "dist={:.2f},road={},offset={:.2f})",
                                p[0], p[1], p[2], //
                                d[0], d[1], d[2], //
                                self.distance(), self.road_id(), self.offset());
@@ -108,6 +109,7 @@ void bind_network(py::module &m)
     py::class_<Network>(m, "Network", py::module_local()) //
                                                           //
         .def(py::init<bool>(), py::kw_only(), "is_wgs84"_a)
+        .def("is_wgs84", &Network::is_wgs84)
         //
         .def("add_road", &Network::add_road, "geom"_a, py::kw_only(), "id"_a)
         .def("add_link", &Network::add_link, "source_road"_a, "target_road"_a,

diff --git a/src/bindings/pybind11_packedrtree.cpp b/src/bindings/pybind11_packedrtree.cpp
@@ -34,11 +34,6 @@ void bind_packedrtree(py::module &m)
                 return new NodeItem{bbox[0], bbox[1], bbox[2], bbox[3], offset};
             }),
             "bbox"_a, "offset"_a = 0)
-        .def(py::init([](const Eigen::Vector2d &min, const Eigen::Vector2d &max,
-                         uint64_t offset) {
-                 return new NodeItem{min[0], min[1], max[0], max[1], offset};
-             }),
-             "min"_a, "max"_a, "offset"_a = 0)
         .def_readwrite("minX", &NodeItem::minX)
         .def_readwrite("minY", &NodeItem::minY)
         .def_readwrite("maxX", &NodeItem::maxX)

diff --git a/src/bindings/pybind11_polyline.cpp b/src/bindings/pybind11_polyline.cpp
@@ -61,7 +61,8 @@ void bind_polyline(py::module &m)
                       const Eigen::Vector3d>(),
              "coords"_a, py::kw_only(), "k"_a)
         //
-        .def("polyline", &Polyline::polyline, rvp::reference_internal)
+        .def("as_numpy", &Polyline::polyline, rvp::reference_internal)
+        .def("N", &Polyline::N)
         .def("k", &Polyline::k)
         .def("is_wgs84", &Polyline::is_wgs84)
         //

diff --git a/src/nano_fmm/network.cpp b/src/nano_fmm/network.cpp
@@ -373,6 +373,13 @@ FlatGeobuf::PackedRTree &Network::rtree() const
     using namespace FlatGeobuf;
 
     auto nodes = std::vector<NodeItem>{};
+    int N = 0;
+    for (auto &pair : roads_) {
+        N += pair.second.N() - 1;
+    }
+    nodes.reserve(N);
+    segs_.reserve(N);
+
     uint64_t ii = 0;
     for (auto &pair : roads_) {
         int64_t poly_idx = pair.first;
@@ -397,7 +404,7 @@ FlatGeobuf::PackedRTree &Network::rtree() const
         }
     }
     auto extent = calcExtent(nodes);
-    hilbertSort(nodes);
+    hilbertSort(nodes, extent);
     rtree_ = FlatGeobuf::PackedRTree(nodes, extent);
     return *rtree_;
 }

diff --git a/src/nano_fmm/network.hpp b/src/nano_fmm/network.hpp
@@ -22,6 +22,7 @@ namespace nano_fmm
 struct Network
 {
     Network(bool is_wgs84) : is_wgs84_(is_wgs84) {}
+    bool is_wgs84() const { return is_wgs84_; }
 
     // road network
     bool add_road(const Eigen::Ref<const RowVectors> &geom, int64_t road_id);

diff --git a/src/nano_fmm/network/ubodt.hpp b/src/nano_fmm/network/ubodt.hpp
@@ -46,7 +46,7 @@ struct UbodtRecord
         return source_road_ == rhs.source_road_ &&
                target_road_ == rhs.target_road_ &&
                source_next_ == rhs.source_next_ &&
-               target_prev_ == rhs.target_prev_;
+               target_prev_ == rhs.target_prev_ && cost_ == rhs.cost_;
     }
 
   private:

diff --git a/src/nano_fmm/polyline.hpp b/src/nano_fmm/polyline.hpp
@@ -26,6 +26,7 @@ struct Polyline
     }
 
     const RowVectors &polyline() const { return polyline_; }
+    int N() const { return N_; }
     Eigen::Vector3d k() const { return k_; }
     bool is_wgs84() const { return is_wgs84_; }
 

diff --git a/tests/test_basic.py b/tests/test_basic.py
@@ -575,7 +575,7 @@ def test_indexer():
 # # fmm.utils.set_logging_level(6) # off
 
 
-def test_network():
+def test_network_read_write():
     network = Network.load("README.md")
     assert network is None
     network = Network.load("missing_file")
@@ -601,3 +601,80 @@ def test_network():
     rows = network.build_ubodt()
     rows = sorted(rows)
     print(rows[:5])
+
+
+def test_network_query():
+    network = Network.load("build/network.geojson")
+    assert network.is_wgs84()
+    assert len(network.roads()) == 1016
+    assert network.next_roads(1293) == {1297, 1298}
+    assert network.prev_roads(1297) == {1293}
+
+    polyline = network.road(1293)
+    assert polyline.is_wgs84()
+    assert len(polyline.as_numpy()) == polyline.N() == 6
+    assert round(polyline.length(), 3) == 225.543
+    assert round(polyline.range(0, t=0.3), 3) == 10.236
+    assert polyline.k().round(2).tolist() == [95504.26, 110869.46, 1.0]
+    seed_lla = [120.663031, 31.40531, 0]
+    lla, dist, seg_idx, t = polyline.nearest(seed_lla)
+    assert round(dist, 2) == 105.71
+    assert seg_idx == 4 and t == 1.0
+    # 0---1---2---3---4---5
+    #                     ^
+    #                   t=1.0
+    assert np.fabs(polyline.range(4, t=1.0) - polyline.length()) < 1e-15
+
+    hits = network.query(seed_lla, radius=40.0)
+    assert len(hits) == 4
+
+
+def test_network_query_enu():
+    network = Network(is_wgs84=False)
+    """
+    r0 o-------------o
+                     |
+    r1 o-------------o
+                     |
+    r2               o
+    """
+    network.add_road([[0, 5, 0], [10, 5, 0]], id=0)
+    network.add_road([[0, 0, 0], [10, 0, 0]], id=1)
+    network.add_road([[10, -5, 0], [10, 5, 0]], id=2)
+    hits = network.query([5, 0, 0], radius=3)
+    hits = [h.to_rapidjson()() for h in hits]
+    assert hits == [
+        {
+            "position": [5.0, 0.0, 0.0],
+            "direction": [1.0, 0.0, 0.0],
+            "distance": 0.0,
+            "road_id": 1,
+            "offset": 5.0,
+        },
+    ]
+
+    hits = network.query([5, 0, 0], radius=5)
+    hits = [h.to_rapidjson()() for h in hits]
+    assert hits == [
+        {
+            "position": [5.0, 0.0, 0.0],
+            "direction": [1.0, 0.0, 0.0],
+            "distance": 0.0,
+            "road_id": 1,
+            "offset": 5.0,
+        },
+        {
+            "position": [5.0, 5.0, 0.0],
+            "direction": [1.0, 0.0, 0.0],
+            "distance": 5.0,
+            "road_id": 0,
+            "offset": 5.0,
+        },
+        {
+            "position": [10.0, 0.0, 0.0],
+            "direction": [0.0, 1.0, 0.0],
+            "distance": 5.0,
+            "road_id": 2,
+            "offset": 5.0,
+        },
+    ]