diff --git a/.gitignore b/.gitignore index 51aa0d3..5931277 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ wheelhouse site data/* !data/Makefile +*.whl diff --git a/3rdparty/packedrtree.cpp b/3rdparty/packedrtree.cpp index 6007420..d41993f 100644 --- a/3rdparty/packedrtree.cpp +++ b/3rdparty/packedrtree.cpp @@ -180,9 +180,12 @@ void hilbertSort(std::vector> &items) }); } -void hilbertSort(std::vector &items) +void hilbertSort(std::vector &items) { + hilbertSort(items, calcExtent(items)); +} + +void hilbertSort(std::vector &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 index 5829673..9a09e2a 100644 --- a/3rdparty/packedrtree.h +++ b/3rdparty/packedrtree.h @@ -71,6 +71,11 @@ struct NodeItem std::vector 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 void hilbertSort(std::deque &items) } void hilbertSort(std::vector &items); +void hilbertSort(std::vector &items, const NodeItem &extent); NodeItem calcExtent(const std::vector> &items); NodeItem calcExtent(const std::vector &rects); diff --git a/3rdparty/optional.hpp b/3rdparty/tl/optional.hpp similarity index 100% rename from 3rdparty/optional.hpp rename to 3rdparty/tl/optional.hpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 0da4dbf..a1d6a6a 100644 --- 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 index 71dbebb..fe5360a 100644 --- 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 index 1ba5643..402e8f0 100644 --- 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 new file mode 100644 index 0000000..13221f9 --- /dev/null +++ b/src/bindings/benchmarks.cpp @@ -0,0 +1,68 @@ +#include +#include +#include +#include +#include + +#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(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 new file mode 100644 index 0000000..5598d0f --- /dev/null +++ b/src/bindings/network.cpp @@ -0,0 +1,21 @@ +#include +#include +#include +#include +#include + +#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_(m, "Network", py::module_local()) // + // + ; +} +} // namespace nano_fmm diff --git a/src/bindings/packedrtree.cpp b/src/bindings/packedrtree.cpp new file mode 100644 index 0000000..acb8374 --- /dev/null +++ b/src/bindings/packedrtree.cpp @@ -0,0 +1,189 @@ +#include +#include +#include +#include +#include + +#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_(m, "NodeItem", py::module_local()) // + .def(py::init<>()) + .def(py::init(), // + "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_(m, "Item", py::module_local()) // + .def(py::init<>()) + .def_readwrite("nodeItem", &Item::nodeItem) + // + ; + py::class_(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(&hilbert), // + "x"_a, "y"_a) + // + .def( + "hilbertSort", + [](const std::vector &items) { + auto sorted = items; + hilbertSort(sorted); + return sorted; + }, + "items"_a) + .def( + "calcExtent", + [](const std::vector &rects) { + return calcExtent(rects); + }, + "rects"_a) + // + ; + + py::class_(m, "PackedRTree", py::module_local()) // + .def(py::init &, 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 &bbox_min, + const Eigen::Ref &bbox_max, + const uint16_t node_size) { + auto extent = NodeItem::create(); + const uint64_t N = bbox_min.rows(); + auto nodes = std::vector(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 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 diff --git a/src/bindings/polyline.cpp b/src/bindings/polyline.cpp index be1830e..8b019b6 100644 --- a/src/bindings/polyline.cpp +++ b/src/bindings/polyline.cpp @@ -19,11 +19,17 @@ void bind_polyline(py::module &m) "A"_a, "B"_a) .def("distance", &LineSegment::distance, "P"_a) .def("distance2", &LineSegment::distance2, "P"_a) + .def("nearest", &LineSegment::nearest, "P"_a) + .def("t", &LineSegment::t, "P"_a) + .def("interpolate", &LineSegment::interpolate, "t"_a) .def_property_readonly( - "length", - [](const LineSegment &self) { return std::sqrt(self.len2); }) + "length", [](const LineSegment &self) { return self.length(); }) .def_property_readonly( "length2", [](const LineSegment &self) { return self.len2; }) + .def_property_readonly("dir", + [](const LineSegment &self) -> Eigen::Vector3d { + return self.dir(); + }) .def_property_readonly( "A", [](const LineSegment &self) -> const Eigen::Vector3d & { @@ -46,14 +52,26 @@ void bind_polyline(py::module &m) ; py::class_(m, "Polyline", py::module_local()) // - .def(py::init &, - const std::optional>(), // - "coords"_a, py::kw_only(), "k"_a = std::nullopt) + .def(py::init &, bool>(), "coords"_a, + py::kw_only(), "is_wgs84"_a = false) // .def("polyline", &Polyline::polyline, rvp::reference_internal) - .def("scale", &Polyline::scale) + .def("k", &Polyline::k) .def("is_wgs84", &Polyline::is_wgs84) // + .def("range", &Polyline::range, "seg_idx"_a, py::kw_only(), "t"_a) + .def("segment_index_t", &Polyline::segment_index_t, "range"_a) + .def("length", &Polyline::length) + .def("along", &Polyline::along, "range"_a, py::kw_only(), + "extend"_a = false) + .def("snap", &Polyline::snap, "point"_a) + .def("slice", &Polyline::slice, py::kw_only(), "min"_a = std::nullopt, + "max"_a = std::nullopt) + // + .def("segment", &Polyline::segment, "index"_a, rvp::reference_internal) + // .def("segments", &Polyline::segments) + .def("ranges", &Polyline::ranges, rvp::reference_internal) + // ; } } // namespace nano_fmm diff --git a/src/bindings/utils.cpp b/src/bindings/utils.cpp new file mode 100644 index 0000000..9b8899f --- /dev/null +++ b/src/bindings/utils.cpp @@ -0,0 +1,33 @@ +#include +#include +#include +#include +#include + +#include "nano_fmm/utils.hpp" + +namespace nano_fmm +{ +namespace py = pybind11; +using namespace pybind11::literals; +using rvp = py::return_value_policy; + +void bind_utils(py::module &m) +{ + m // + .def("cheap_ruler_k", &utils::cheap_ruler_k, "latitude"_a) + .def("lla2enu", &utils::lla2enu, // + "llas"_a, py::kw_only(), // + "anchor_lla"_a = std::nullopt, // + "k"_a = std::nullopt) + .def("enu2lla", &utils::enu2lla, // + "enus"_a, py::kw_only(), // + "anchor_lla"_a, // + "k"_a = std::nullopt) + .def("index2mask", &utils::index2mask, "index"_a, "N"_a) + .def("mask2index", &utils::mask2index, "mask"_a) + .def("to_Nx3", &utils::to_Nx3, "coords"_a) + // + ; +} +} // namespace nano_fmm diff --git a/src/main.cpp b/src/main.cpp index 6ff70a3..c86d806 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -9,8 +9,12 @@ namespace py = pybind11; namespace nano_fmm { +void bind_benchmarks(py::module &m); +void bind_network(py::module &m); +void bind_packedrtree(py::module &m); void bind_polyline(py::module &m); -} +void bind_utils(py::module &m); +} // namespace nano_fmm PYBIND11_MODULE(_nano_fmm, m) { @@ -26,5 +30,13 @@ PYBIND11_MODULE(_nano_fmm, m) m.attr("__version__") = "dev"; #endif + auto utils = m.def_submodule("utils"); + auto benchmarks = m.def_submodule("benchmarks"); + auto flatbush = m.def_submodule("flatbush"); + + nano_fmm::bind_benchmarks(benchmarks); + nano_fmm::bind_network(m); + nano_fmm::bind_packedrtree(flatbush); nano_fmm::bind_polyline(m); + nano_fmm::bind_utils(utils); } diff --git a/src/nano_fmm/config.cpp b/src/nano_fmm/config.cpp new file mode 100755 index 0000000..53ba6fc --- /dev/null +++ b/src/nano_fmm/config.cpp @@ -0,0 +1 @@ +#include "nano_fmm/config.hpp" diff --git a/src/nano_fmm/network.cpp b/src/nano_fmm/network.cpp index 4159358..5fb77a7 100644 --- a/src/nano_fmm/network.cpp +++ b/src/nano_fmm/network.cpp @@ -1,4 +1,5 @@ #include "nano_fmm/network.hpp" +#include "nano_fmm/utils.hpp" namespace nano_fmm { @@ -11,12 +12,14 @@ std::shared_ptr Network::config() return config_; } -void Network::add(const Eigen::Ref &polyline, int64_t id) +void Network::config(std::shared_ptr config) { config_ = config; } + +void Network::add_node(int64_t node_id, const Eigen::Ref &polyline) { // } -void Network::add(const std::vector &polylines, - std::optional ids) +void Network::add_edge(int64_t edge_id, int64_t source_node, + int64_t target_node) { // } @@ -24,9 +27,12 @@ void Network::add(const std::vector &polylines, int Network::load(const std::string &path) { // + return 0; } bool Network::dump(const std::string &path) const { // + return false; } + } // namespace nano_fmm diff --git a/src/nano_fmm/network.hpp b/src/nano_fmm/network.hpp index 9fcb2bf..8bee972 100644 --- a/src/nano_fmm/network.hpp +++ b/src/nano_fmm/network.hpp @@ -21,13 +21,13 @@ struct ProjectedPoint struct Network { - Network(bool is_wgs84 = true) : is_wgs84_(is_wgs84) {} + Network(bool is_wgs84 = false) : is_wgs84_(is_wgs84) {} std::shared_ptr config(); + void config(std::shared_ptr config); - void add(const Eigen::Ref &polyline, int64_t id = -1); - void add(const std::vector &polylines, - std::optional ids = std::nullopt); + void add_node(int64_t node_id, const Eigen::Ref &polyline); + void add_edge(int64_t edge_id, int64_t source_node, int64_t target_node); std::vector query(const Eigen::Vector3d &position, @@ -39,6 +39,8 @@ struct Network bool build_ubodt(std::optional thresh) const; + Network to_2d() const; + private: bool is_wgs84_{true}; std::vector ids_; diff --git a/src/nano_fmm/polyline.hpp b/src/nano_fmm/polyline.hpp index 9f8a63b..7b72d7b 100644 --- a/src/nano_fmm/polyline.hpp +++ b/src/nano_fmm/polyline.hpp @@ -13,7 +13,7 @@ struct LineSegment const double len2, inv_len2; LineSegment(const Eigen::Vector3d &a, const Eigen::Vector3d &b) : A(a), B(b), AB(b - a), // - len2((b - a).squaredNorm()), inv_len2(1.0 / len2) + len2(AB.squaredNorm()), inv_len2(1.0 / len2) { } double distance2(const Eigen::Vector3d &P) const @@ -32,13 +32,31 @@ struct LineSegment { return std::sqrt(distance2(P)); } + // return P', distance, t + std::tuple + nearest(const Eigen::Vector3d &P) const + { + double dot = (P - A).dot(AB); + if (dot <= 0) { + return std::make_tuple(A, (P - A).squaredNorm(), 0.0); + } else if (dot >= len2) { + return std::make_tuple(B, (P - B).squaredNorm(), 1.0); + } + Eigen::Vector3d PP = A + (dot * inv_len2 * AB); + return std::make_tuple(PP, (PP - P).squaredNorm(), dot * inv_len2); + } + double t(const Eigen::Vector3d &P) const + { + return (P - A).dot(AB) * inv_len2; + } + Eigen::Vector3d interpolate(double t) const { - // 0 -> A, 1 -> B - return A; + return A * (1.0 - t) + B * t; } - // dist, t, dot + double length() const { return std::sqrt(len2); } + Eigen::Vector3d dir() const { return AB / length(); } }; // https://github.com/cubao/headers/blob/main/include/cubao/polyline_ruler.hpp @@ -46,31 +64,24 @@ struct Polyline { EIGEN_MAKE_ALIGNED_OPERATOR_NEW Polyline(const Eigen::Ref &polyline, - const std::optional scale = {}) + bool is_wgs84 = false) : polyline_(polyline), // N_(polyline.rows()), // - scale_(scale) + is_wgs84_(is_wgs84), + k_(is_wgs84 ? cheap_ruler_k(polyline(0, 1)) : Eigen::Vector3d::Ones()) { } const RowVectors &polyline() const { return polyline_; } - std::optional scale() const { return scale_; } - bool is_wgs84() const { return (bool)scale_; } + Eigen::Vector3d k() const { return k_; } + bool is_wgs84() const { return is_wgs84_; } - double range(int seg_idx) const { return ranges()[seg_idx]; } - double range(int seg_idx, double t) const + double range(int seg_idx, double t = 0.0) const { auto &ranges = this->ranges(); return ranges[seg_idx] * (1.0 - t) + ranges[seg_idx + 1] * t; } - int segment_index(double range) const - { - const double *ranges = this->ranges().data(); - int I = std::upper_bound(ranges, ranges + N_, range) - ranges; - return std::min(std::max(0, I - 1), N_ - 2); - } - std::pair segment_index_t(double range) const { const double *ranges = this->ranges().data(); @@ -106,17 +117,31 @@ struct Polyline return a + (b - a) * t; } - private: - const RowVectors polyline_; - const int N_; - const std::optional scale_; - - mutable std::optional> segments_; - mutable std::optional ranges_; - + const LineSegment &segment(int index) const + { + index = index < 0 ? index + N_ - 1 : index; + return segments()[index]; + } const std::vector &segments() const { - // + if (segments_) { + return *segments_; + } + segments_ = std::vector{}; + segments_->reserve(N_ - 1); + if (!is_wgs84_) { + for (int i = 1; i < N_; ++i) { + segments_->emplace_back(polyline_.row(i - 1), polyline_.row(i)); + } + } else { + for (int i = 1; i < N_; ++i) { + Eigen::Vector3d A = polyline_.row(i - 1) - polyline_.row(0); + Eigen::Vector3d B = polyline_.row(i) - polyline_.row(0); + A.array() *= k_.array(); + B.array() *= k_.array(); + segments_->emplace_back(A, B); + } + } return *segments_; } const Eigen::VectorXd &ranges() const @@ -125,13 +150,38 @@ struct Polyline return *ranges_; } Eigen::VectorXd ranges(N_); + ranges.setZero(); int idx = 0; for (auto &seg : segments()) { - ranges[idx++] = std::sqrt(seg.len2); + ranges[idx + 1] = ranges[idx] + seg.length(); + ++idx; } ranges_ = std::move(ranges); return *ranges_; } + + private: + const RowVectors polyline_; + const int N_; + const bool is_wgs84_; + const Eigen::Vector3d k_; + + mutable std::optional> segments_; + mutable std::optional ranges_; + + // same as utils.hpp/cheap_ruler_k + inline Eigen::Vector3d cheap_ruler_k(double latitude) + { + static constexpr double RE = 6378.137; + static constexpr double FE = 1.0 / 298.257223563; + static constexpr double E2 = FE * (2 - FE); + static constexpr double RAD = M_PI / 180.0; + static constexpr double MUL = RAD * RE * 1000.; + double coslat = std::cos(latitude * RAD); + double w2 = 1.0 / (1.0 - E2 * (1.0 - coslat * coslat)); + double w = std::sqrt(w2); + return Eigen::Vector3d(MUL * w * coslat, MUL * w * w2 * (1 - E2), 1.0); + } }; } // namespace nano_fmm diff --git a/src/nano_fmm/types.hpp b/src/nano_fmm/types.hpp index 1318938..01478a3 100644 --- a/src/nano_fmm/types.hpp +++ b/src/nano_fmm/types.hpp @@ -17,4 +17,6 @@ using RowVectorsNx3 = RowVectors; // without-z using RowVectorsNx2 = Eigen::Matrix; +using VectorUi64 = Eigen::Matrix; + } // namespace nano_fmm diff --git a/src/nano_fmm/utils.hpp b/src/nano_fmm/utils.hpp index 6b3c0e5..6a5c942 100644 --- a/src/nano_fmm/utils.hpp +++ b/src/nano_fmm/utils.hpp @@ -18,7 +18,7 @@ inline Eigen::Vector3d cheap_ruler_k(double latitude) static constexpr double RAD = M_PI / 180.0; static constexpr double MUL = RAD * RE * 1000.; double coslat = std::cos(latitude * RAD); - double w2 = 1 / (1 - E2 * (1 - coslat * coslat)); + double w2 = 1.0 / (1.0 - E2 * (1.0 - coslat * coslat)); double w = std::sqrt(w2); return Eigen::Vector3d(MUL * w * coslat, MUL * w * w2 * (1 - E2), 1.0); } @@ -50,10 +50,12 @@ inline RowVectors enu2lla(const Eigen::Ref &enus, if (!enus.rows()) { return RowVectors(0, 3); } - auto k = cheap_ruler_k(anchor_lla[1]); + if (!k) { + k = cheap_ruler_k(anchor_lla[1]); + } RowVectors llas = enus; for (int i = 0; i < 3; ++i) { - llas.col(i).array() /= k[i]; + llas.col(i).array() /= (*k)[i]; llas.col(i).array() += anchor_lla[i]; } return llas; @@ -62,7 +64,7 @@ inline RowVectors enu2lla(const Eigen::Ref &enus, // https://github.com/cubao/headers/blob/main/include/cubao/eigen_helpers.hpp inline Eigen::VectorXi -indexes2mask(const Eigen::Ref &indexes, int N) +index2mask(const Eigen::Ref &indexes, int N) { Eigen::VectorXi mask(N); mask.setZero(); @@ -72,8 +74,7 @@ indexes2mask(const Eigen::Ref &indexes, int N) return mask; } -inline Eigen::VectorXi -mask2indexes(const Eigen::Ref &mask) +inline Eigen::VectorXi mask2index(const Eigen::Ref &mask) { Eigen::VectorXi indexes(mask.sum()); for (int i = 0, j = 0, N = mask.size(); i < N; ++i) { diff --git a/tests/test_basic.py b/tests/test_basic.py index 5ef75ae..9d1b61a 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -1,5 +1,211 @@ -import nano_fmm as m +import time +import numpy as np -def test_main(): - assert m.add(1, 2) == 3 +import nano_fmm as fmm +from nano_fmm import LineSegment +from nano_fmm import flatbush as fb + + +def test_add(): + assert fmm.add(1, 2) == 3 + + +def test_segment(): + seg = LineSegment([0, 0, 0], [10, 0, 0]) + assert 4.0 == seg.distance([5.0, 4.0, 0.0]) + assert 5.0 == seg.distance([-4.0, 3.0, 0.0]) + assert 5.0 == seg.distance([14.0, 3.0, 0.0]) + assert 25.0 == seg.distance2([14.0, 3.0, 0.0]) + + assert seg.length == 10.0 + assert seg.length2 == 100.0 + assert np.all(seg.A == [0, 0, 0]) + assert np.all(seg.B == [10, 0, 0]) + assert np.all(seg.AB == [10, 0, 0]) + assert np.all(seg.dir == [1, 0, 0]) + assert np.all(seg.interpolate(0.4) == [4, 0, 0]) + assert seg.t([4, 0, 0]) == 0.4 + PP, dist, t = seg.nearest([4, 1, 0]) + assert np.all(PP == [4, 0, 0]) + assert dist == 1.0 + assert t == 0.4 + + seg = LineSegment([0, 0, 0], [0, 0, 0]) + assert 5.0 == seg.distance([3.0, 4.0, 0.0]) + assert 5.0 == seg.distance([-4.0, 3.0, 0.0]) + assert 13.0 == seg.distance([5.0, 12.0, 0.0]) + + +def test_utils(): + k0 = fmm.utils.cheap_ruler_k(0.0) + k1 = fmm.utils.cheap_ruler_k(30.0) + assert k0[0] > k1[0] + + anchor = [123.4, 5.6, 7.8] + enus = [[1, 2, 3], [4, 5, 6]] + llas = fmm.utils.enu2lla(enus, anchor_lla=anchor) + enus2 = fmm.utils.lla2enu(llas, anchor_lla=anchor) + assert np.max(enus2 - enus) < 1e-9 + + +def test_polyline(): + enus = [[0, 0, 0], [10, 0, 0], [13, 4, 0]] + polyline = fmm.Polyline(enus) + assert polyline.segment(-1) == polyline.segment(1) + assert polyline.segment(-1) != polyline.segment(0) + assert np.all(polyline.ranges() == [0, 10, 15]) + + anchor = [123.4, 5.6, 7.8] + k = fmm.utils.cheap_ruler_k(anchor[1]) + llas = fmm.utils.enu2lla(enus, anchor_lla=anchor) + polyline2 = fmm.Polyline(llas, is_wgs84=True) + assert np.fabs(polyline2.k() - k).max() < 100 + for i in range(2): + seg1 = polyline.segment(i) + seg2 = polyline2.segment(i) + assert np.max(np.fabs(seg1.A - seg2.A)) < 1e-9 + assert np.max(np.fabs(seg1.B - seg2.B)) < 1e-9 + + +def test_cheap_ruler_k(): + N = 100000 + tic = time.time() + fmm.benchmarks.cheap_ruler_k(N) + toc = time.time() + print(toc - tic, "secs") + tic = time.time() + fmm.benchmarks.cheap_ruler_k_lookup_table(N) + toc = time.time() + print(toc - tic, "secs (with lookup)") + print() + + +def test_geobuf_rtree(): + n = fb.NodeItem() + assert n.minX == n.minY == n.maxX == n.maxY == 0.0 + assert n.offset == 0 + assert n.width() == n.height() == 0.0 + + n = fb.NodeItem.sum(fb.NodeItem(0, 1, 2, 3, 4), fb.NodeItem(0, 10, 20, 30, 40)) + + tree = fb.PackedRTree( + [fb.NodeItem(1, 1, 9, 9, 0), fb.NodeItem(5, 5, 8, 8, 0)], + extent=fb.NodeItem(0, 0, 10, 10, 0), + ) + assert len(tree.to_bytes()) + + bboxes = [ + [0, 0, 10, 10], + [1, 1, 5, 5], + [3, 3, 7, 7], + [2, 2, 9, 3], + ] + bboxes = np.array(bboxes, dtype=np.float64) + tree = fb.PackedRTree(bboxes[:, :2], bboxes[:, 2:]) + + tree.search(0, 0, 3, 3) + tree.searchIndex(0, 0, 1, 1) + print() + + +def test_cpp_migrated_1(): + """ + migrated test: https://github.com/flatgeobuf/flatgeobuf/blob/master/src/cpp/test/packedrtree.h + PackedRTree 2 item one dimension + """ + nodes = [ + fb.NodeItem(0, 0, 0, 0), + fb.NodeItem(0, 0, 0, 0), + ] + assert nodes[0] == nodes[1] + extent = fb.calcExtent(nodes) + assert nodes[0].intersects(fb.NodeItem(0, 0, 0, 0)) + nodes = fb.hilbertSort(nodes) + offset = 0 + for node in nodes: + node.offset = offset + offset += fb.NodeItem._size_() + tree = fb.PackedRTree(nodes, extent) + hits = tree.search(0, 0, 0, 0) + assert len(hits) == 2 + assert nodes[hits[0].index].intersects(0, 0, 0, 0) + + +def test_cpp_migrated_2(): + """ + migrated test: https://github.com/flatgeobuf/flatgeobuf/blob/master/src/cpp/test/packedrtree.h + PackedRTree 2 items 2 + """ + nodes = [ + fb.NodeItem(0, 0, 1, 1), + fb.NodeItem(2, 2, 3, 3), + ] + extent = fb.calcExtent(nodes) + assert nodes[0].intersects(0, 0, 1, 1) + assert nodes[1].intersects(2, 2, 3, 3) + nodes = fb.hilbertSort(nodes) + offset = 0 + for node in nodes: + node.offset = offset + offset += fb.NodeItem._size_() + assert nodes[1].intersects(0, 0, 1, 1) + assert nodes[0].intersects(2, 2, 3, 3) + + tree = fb.PackedRTree(nodes, extent) + data = tree.to_bytes() + assert len(data) == 120 + assert type(data) == bytes + + hits = tree.search(0, 0, 1, 1) + assert len(hits) == 1 + assert nodes[hits[0].index].intersects(0, 0, 1, 1) + + +def test_cpp_migrated_3(): + """ + migrated test: https://githuconstexpr bool operator==(point const& lhs, point const& rhs) + PackedRTree 19 items + roundtrip + streamSearch + """ + nodes = [ + fb.NodeItem(0, 0, 1, 1), + fb.NodeItem(2, 2, 3, 3), + fb.NodeItem(10, 10, 11, 11), + fb.NodeItem(100, 100, 110, 110), + fb.NodeItem(101, 101, 111, 111), + fb.NodeItem(102, 102, 112, 112), + fb.NodeItem(103, 103, 113, 113), + fb.NodeItem(104, 104, 114, 114), + fb.NodeItem(10010, 10010, 10110, 10110), + fb.NodeItem(10010, 10010, 10110, 10110), + fb.NodeItem(10010, 10010, 10110, 10110), + fb.NodeItem(10010, 10010, 10110, 10110), + fb.NodeItem(10010, 10010, 10110, 10110), + fb.NodeItem(10010, 10010, 10110, 10110), + fb.NodeItem(10010, 10010, 10110, 10110), + fb.NodeItem(10010, 10010, 10110, 10110), + fb.NodeItem(10010, 10010, 10110, 10110), + fb.NodeItem(10010, 10010, 10110, 10110), + fb.NodeItem(10010, 10010, 10110, 10110), + fb.NodeItem(10010, 10010, 10110, 10110), + ] + extent = fb.calcExtent(nodes) + nodes = fb.hilbertSort(nodes) + offset = 0 + for node in nodes: + node.offset = offset + offset += fb.NodeItem._size_() + + tree = fb.PackedRTree(nodes, extent) + hits = tree.search(102, 102, 103, 103) + assert len(hits) == 4 + for h in hits: + node = nodes[h.index] + print(h, node) + assert node.intersects(102, 102, 103, 103) + + data = tree.to_bytes() + tree2 = fb.PackedRTree(data, len(nodes)) + hits2 = tree2.search(102, 102, 103, 103) + assert hits == hits2 + assert hits != hits2[::-1]