test packed rtree (#4)

* move source * fix * fix * line segment * fix utils * lint * test k * better polyline * not ready * add network * lookup table is faster * fix * init packed rtree * add flatbush * bindings * bind more * fix * test search * not ready * migrate test * should fix bytes * double free?? * fix init, should return raw pointer * fix * fix * disable lookup --------- Co-authored-by: TANG ZHIXIONG <[email protected]>
cubao · Aug 20, 2023 · 0310bfa · 0310bfa
1 parent 80b4a95
commit 0310bfa
Show file tree

Hide file tree

Showing 20 changed files with 703 additions and 56 deletions.
diff --git a/.gitignore b/.gitignore
@@ -10,3 +10,4 @@ wheelhouse
 site
 data/*
 !data/Makefile
+*.whl
diff --git a/3rdparty/packedrtree.cpp b/3rdparty/packedrtree.cpp
@@ -180,9 +180,12 @@ void hilbertSort(std::vector<std::shared_ptr<Item>> &items)
               });
 }
 
-void hilbertSort(std::vector<NodeItem> &items)
+void hilbertSort(std::vector<NodeItem> &items) {
+    hilbertSort(items, calcExtent(items));
+}
+
+void hilbertSort(std::vector<NodeItem> &items, const NodeItem &extent)
 {
-    NodeItem extent = calcExtent(items);
     const double minX = extent.minX;
     const double minY = extent.minY;
     const double width = extent.width();

diff --git a/3rdparty/packedrtree.h b/3rdparty/packedrtree.h
@@ -71,6 +71,11 @@ struct NodeItem
     std::vector<double> toVector();
 };
 
+inline bool operator==(const NodeItem& lhs, const NodeItem& rhs)
+{
+    return lhs.minX == rhs.minX && lhs.minY == rhs.minY && lhs.maxX == rhs.maxX && lhs.maxY == rhs.maxY && lhs.offset == rhs.offset;
+}
+
 struct Item
 {
     NodeItem nodeItem;
@@ -82,6 +87,11 @@ struct SearchResultItem
     uint64_t index;
 };
 
+inline bool operator==(const SearchResultItem& lhs, const SearchResultItem& rhs)
+{
+    return lhs.index == rhs.index && lhs.offset == rhs.offset;
+}
+
 std::ostream &operator<<(std::ostream &os, NodeItem const &value);
 
 uint32_t hilbert(uint32_t x, uint32_t y);
@@ -119,6 +129,7 @@ template <class ITEM_TYPE> void hilbertSort(std::deque<ITEM_TYPE> &items)
 }
 
 void hilbertSort(std::vector<NodeItem> &items);
+void hilbertSort(std::vector<NodeItem> &items, const NodeItem &extent);
 NodeItem calcExtent(const std::vector<std::shared_ptr<Item>> &items);
 NodeItem calcExtent(const std::vector<NodeItem> &rects);
 

diff --git a/3rdparty/optional.hpp → 3rdparty/tl/optional.hpp b/3rdparty/optional.hpp → 3rdparty/tl/optional.hpp
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -1,8 +1,9 @@
-cmake_minimum_required(VERSION 3.4...3.18)
+cmake_minimum_required(VERSION 3.5...3.18)
 project(nano_fmm)
 
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+set(CMAKE_INCLUDE_CURRENT_DIR ON)
 set(CMAKE_CXX_STANDARD 17)
 
 if(NOT CMAKE_ARCHIVE_OUTPUT_DIRECTORY)
@@ -15,6 +16,24 @@ if(NOT CMAKE_RUNTIME_OUTPUT_DIRECTORY)
   set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR}/bin)
 endif()
 
+# set(CMAKE_BUILD_TYPE "Debug")
+
+if(NOT CMAKE_BUILD_TYPE OR CMAKE_BUILD_TYPE STREQUAL "")
+  set(CMAKE_BUILD_TYPE
+      "Release"
+      CACHE STRING "" FORCE)
+  message(STATUS "Set build type to default: ${CMAKE_BUILD_TYPE}")
+else()
+  message(STATUS "Your build type: ${CMAKE_BUILD_TYPE}")
+endif()
+if(CMAKE_BUILD_TYPE STREQUAL "Debug")
+  set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -O0 -ggdb")
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O0 -ggdb")
+elseif(CMAKE_BUILD_TYPE STREQUAL "Release")
+  set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -O3")
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3")
+endif()
+
 include_directories(3rdparty src)
 
 set(PYBIND11_CPP_STANDARD -std=c++17)

diff --git a/MANIFEST.in b/MANIFEST.in
@@ -1,5 +1,7 @@
-include README.md LICENSE pybind11/LICENSE
+global-include CMakeLists.txt *.cmake
+graft 3rdparty
+graft nano_fmm
 graft pybind11/include
 graft pybind11/tools
 graft src
-global-include CMakeLists.txt *.cmake
+include README.md LICENSE pybind11/LICENSE
diff --git a/Makefile b/Makefile
@@ -54,11 +54,13 @@ test_in_dev_container:
 
 PYTHON ?= python3
 python_install:
-	$(PYTHON) setup.py install
+	$(PYTHON) setup.py install --force
 python_build:
 	$(PYTHON) setup.py bdist_wheel
 python_sdist:
 	$(PYTHON) setup.py 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))'

diff --git a/src/bindings/benchmarks.cpp b/src/bindings/benchmarks.cpp
@@ -0,0 +1,68 @@
+#include <pybind11/eigen.h>
+#include <pybind11/iostream.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/stl_bind.h>
+
+#include "nano_fmm/utils.hpp"
+
+namespace nano_fmm
+{
+namespace py = pybind11;
+using namespace pybind11::literals;
+using rvp = py::return_value_policy;
+
+inline Eigen::Vector3d cheap_ruler_k_lookup_table(double latitude)
+{
+    // lookup table
+#ifdef K
+#undef K
+#endif
+#define K(lat) utils::cheap_ruler_k((double)lat)
+    const static Eigen::Vector3d Ks[] = {
+        // clang-format off
+        K(0),K(1),K(2),K(3),K(4),K(5),K(6),K(7),K(8),K(9),K(10),K(11),K(12),K(13),
+        K(14),K(15),K(16),K(17),K(18),K(19),K(20),K(21),K(22),K(23),K(24),K(25),K(26),
+        K(27),K(28),K(29),K(30),K(31),K(32),K(33),K(34),K(35),K(36),K(37),K(38),K(39),
+        K(40),K(41),K(42),K(43),K(44),K(45),K(46),K(47),K(48),K(49),K(50),K(51),K(52),
+        K(53),K(54),K(55),K(56),K(57),K(58),K(59),K(60),K(61),K(62),K(63),K(64),K(65),
+        K(66),K(67),K(68),K(69),K(70),K(71),K(72),K(73),K(74),K(75),K(76),K(77),K(78),
+        K(79),K(80),K(81),K(82),K(83),K(84),K(85),K(86),K(87),K(88),K(89),K(90)
+        // clang-format on
+    };
+    int idx =
+        std::min(90, static_cast<int>(std::floor(std::fabs(latitude) + 0.5)));
+    return Ks[idx];
+}
+
+void bind_benchmarks(py::module &m)
+{
+    m //
+        .def(
+            "cheap_ruler_k",
+            [](int round) {
+                Eigen::Vector3d k(0, 0, 0);
+                for (int i = 0; i < round; ++i) {
+                    for (double l = 0; l < 90.0; l += 0.5) {
+                        k += utils::cheap_ruler_k(l);
+                    }
+                }
+                return k;
+            },
+            "round"_a = 1000)
+        .def(
+            "cheap_ruler_k_lookup_table",
+            [](int round) {
+                Eigen::Vector3d k(0, 0, 0);
+                for (int i = 0; i < round; ++i) {
+                    for (double l = 0; l < 90.0; l += 0.5) {
+                        k += cheap_ruler_k_lookup_table(l);
+                    }
+                }
+                return k;
+            },
+            "round"_a = 1000)
+        //
+        ;
+}
+} // namespace nano_fmm
diff --git a/src/bindings/network.cpp b/src/bindings/network.cpp
@@ -0,0 +1,21 @@
+#include <pybind11/eigen.h>
+#include <pybind11/iostream.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/stl_bind.h>
+
+#include "nano_fmm/network.hpp"
+
+namespace nano_fmm
+{
+namespace py = pybind11;
+using namespace pybind11::literals;
+using rvp = py::return_value_policy;
+
+void bind_network(py::module &m)
+{
+    py::class_<Network>(m, "Network", py::module_local()) //
+                                                          //
+        ;
+}
+} // namespace nano_fmm
diff --git a/src/bindings/packedrtree.cpp b/src/bindings/packedrtree.cpp
@@ -0,0 +1,189 @@
+#include <pybind11/eigen.h>
+#include <pybind11/iostream.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/stl_bind.h>
+
+#include "packedrtree.h"
+#include "spdlog/spdlog.h"
+#include "nano_fmm/types.hpp"
+
+namespace nano_fmm
+{
+namespace py = pybind11;
+using namespace pybind11::literals;
+using rvp = py::return_value_policy;
+
+void bind_packedrtree(py::module &m)
+{
+    using namespace FlatGeobuf;
+    py::class_<NodeItem>(m, "NodeItem", py::module_local()) //
+        .def(py::init<>())
+        .def(py::init<double, double, double, double, uint64_t>(), //
+             "minX"_a, "minY"_a,                                   //
+             "maxX"_a, "maxY"_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(
+            py::init([](const Eigen::Vector4d &bbox, uint64_t offset) {
+                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)
+        .def_readwrite("maxY", &NodeItem::maxY)
+        .def_readwrite("offset", &NodeItem::offset)
+        .def("width", &NodeItem::width)
+        .def("height", &NodeItem::height)
+        .def_static("sum", &NodeItem::sum, "a"_a, "b"_a)
+        .def_static("create", &NodeItem::create, "offset"_a = 0)
+        .def("expand", &NodeItem::expand, "r"_a)
+        .def("intersects", &NodeItem::intersects, "r"_a)
+        .def(
+            "intersects",
+            [](const NodeItem &self, double minX, double minY, double maxX,
+               double maxY) {
+                return self.intersects({minX, minY, maxX, maxY});
+            },
+            "minX"_a, "minY"_a, "maxX"_a, "maxY"_a)
+        .def("toVector", &NodeItem::toVector)
+        .def("to_numpy",
+             [](const NodeItem &self) {
+                 return Eigen::Vector4d(self.minX, self.minY, //
+                                        self.maxX, self.maxY);
+             })
+        .def("__repr__",
+             [](const NodeItem &n) {
+                 return fmt::format(
+                     "NodeItem(min=[{},{}],max=[{},{}],offset={})", n.minX,
+                     n.minY, n.maxX, n.maxY, n.offset);
+             })
+        //
+        .def(py::self == py::self)
+        //
+        .def_static("_size_", []() { return sizeof(NodeItem); });
+    py::class_<Item>(m, "Item", py::module_local()) //
+        .def(py::init<>())
+        .def_readwrite("nodeItem", &Item::nodeItem)
+        //
+        ;
+    py::class_<SearchResultItem>(m, "SearchResultItem", py::module_local()) //
+        .def(py::init<>())
+        .def_readwrite("offset", &SearchResultItem::offset)
+        .def_readwrite("index", &SearchResultItem::index)
+        .def("__repr__",
+             [](const SearchResultItem &self) {
+                 return fmt::format("SearchResultItem(offset={},index={})",
+                                    self.offset, self.index);
+             })
+        //
+        .def(py::self == py::self)
+        //
+        ;
+
+    m                                                                    //
+        .def("hilbert", py::overload_cast<uint32_t, uint32_t>(&hilbert), //
+             "x"_a, "y"_a)
+        //
+        .def(
+            "hilbertSort",
+            [](const std::vector<NodeItem> &items) {
+                auto sorted = items;
+                hilbertSort(sorted);
+                return sorted;
+            },
+            "items"_a)
+        .def(
+            "calcExtent",
+            [](const std::vector<NodeItem> &rects) {
+                return calcExtent(rects);
+            },
+            "rects"_a)
+        //
+        ;
+
+    py::class_<PackedRTree>(m, "PackedRTree", py::module_local()) //
+        .def(py::init<const std::vector<NodeItem> &, const NodeItem &,
+                      const uint16_t>(),
+             "nodes"_a, "extent"_a, "nodeSize"_a = 16)
+        .def(py::init([](py::buffer buf, const uint64_t numItems,
+                         const uint16_t nodeSize) {
+                 py::buffer_info info = buf.request();
+                 return new PackedRTree(info.ptr, numItems, nodeSize);
+             }),
+             "data"_a, "numItems"_a, "nodeSize"_a = 16)
+        .def(py::init([](const Eigen::Ref<const RowVectorsNx2> &bbox_min,
+                         const Eigen::Ref<const RowVectorsNx2> &bbox_max,
+                         const uint16_t node_size) {
+                 auto extent = NodeItem::create();
+                 const uint64_t N = bbox_min.rows();
+                 auto nodes = std::vector<NodeItem>(N);
+                 for (uint64_t i = 0; i < N; ++i) {
+                     nodes[i].minX = bbox_min(i, 0);
+                     nodes[i].minY = bbox_min(i, 1);
+                     nodes[i].maxX = bbox_max(i, 0);
+                     nodes[i].maxY = bbox_max(i, 1);
+                     extent.expand(nodes[i]);
+                 }
+                 hilbertSort(nodes, extent);
+                 uint64_t offset = 0;
+                 for (auto &node : nodes) {
+                     node.offset = offset;
+                     offset += sizeof(NodeItem);
+                 }
+                 return new PackedRTree(nodes, extent, node_size);
+             }),
+             "min"_a, "max"_a, "nodeSize"_a = 16)
+        .def("search", &PackedRTree::search, //
+             "minX"_a, "minY"_a,             //
+             "maxX"_a, "maxY"_a)
+        .def(
+            "searchIndex",
+            [](const PackedRTree &self,  //
+               double minX, double minY, //
+               double maxX, double maxY, //
+               bool use_offset) {
+                auto hits = self.search(minX, minY, maxX, maxY);
+                const size_t N = hits.size();
+                VectorUi64 idx(N);
+                if (use_offset) {
+                    for (size_t i = 0; i < N; ++i) {
+                        idx[i] = hits[i].offset;
+                    }
+                } else {
+                    for (size_t i = 0; i < N; ++i) {
+                        idx[i] = hits[i].index;
+                    }
+                }
+                return idx;
+            },
+            "minX"_a, "minY"_a, //
+            "maxX"_a, "maxY"_a, py::kw_only(), "use_offset"_a = true)
+        .def_static("generateLevelBounds", &PackedRTree::generateLevelBounds,
+                    "numItems"_a, "nodeSize"_a)
+        .def("size", py::overload_cast<>(&PackedRTree::size, py::const_))
+        .def("getExtent", &PackedRTree::getExtent)
+        .def("to_bytes",
+             [](PackedRTree &self) {
+                 std::vector<uint8_t> bytes;
+                 self.streamWrite([&bytes](uint8_t *buf, size_t size) {
+                     std::copy(buf, buf + size, std::back_inserter(bytes));
+                 });
+                 return py::bytes((const char *)bytes.data(), bytes.size());
+             })
+
+        //
+        ;
+}
+} // namespace nano_fmm