From 1da69f49de1bfad699c60216b9a22ae995215103 Mon Sep 17 00:00:00 2001 From: TANG ZhiXiong Date: Sat, 18 May 2024 20:13:58 +0800 Subject: [PATCH] tested okay (#3) * fix test * init * not ready * not ready * not ready * ready * not ready * not ready * good * not ready * not ready * ready * update * not ready * not ready * update * gcc5.4 compatible * gil * good * good * good lint code * revert * good --------- Co-authored-by: TANG ZHIXIONG --- pyproject.toml | 2 +- src/fast_viterbi/__init__.py | 4 +- src/main.cpp | 333 +++++++++++++++++++++++++++++++++++ tests/test_basic.py | 4 +- 4 files changed, 338 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 89f4250..e607e4b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "scikit_build_core.build" [project] name = "fast_viterbi" -version = "0.0.1" +version = "0.0.2" description="a viterbi algo collection" readme = "README.md" authors = [ diff --git a/src/fast_viterbi/__init__.py b/src/fast_viterbi/__init__.py index 15188e5..011fc19 100644 --- a/src/fast_viterbi/__init__.py +++ b/src/fast_viterbi/__init__.py @@ -1,5 +1,5 @@ from __future__ import annotations -from ._core import __doc__, __version__, add, subtract +from ._core import FastViterbi, __doc__, __version__, add, subtract -__all__ = ["__doc__", "__version__", "add", "subtract"] +__all__ = ["__doc__", "__version__", "add", "subtract", "FastViterbi"] diff --git a/src/main.cpp b/src/main.cpp index 6ee716e..044ece0 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -1,4 +1,12 @@ #include +#include +#include + +#include +#include +#include +#include +#include #define STRINGIFY(x) #x #define MACRO_STRINGIFY(x) STRINGIFY(x) @@ -6,6 +14,312 @@ int add(int i, int j) { return i + j; } namespace py = pybind11; +using namespace pybind11::literals; + +template +struct hash_vector { + std::size_t operator()(const T &vec) const { + size_t hash_seed = std::hash()(vec.size()); + for (auto elem : vec) { + hash_seed ^= std::hash()(elem) + 0x9e3779b9 + (hash_seed << 6) + (hash_seed >> 2); + } + return hash_seed; + } +}; + +struct Seq { + const std::vector node_path; + const std::vector road_path; + Seq() = default; + Seq(std::vector &&node_path, std::vector &&road_path) + : node_path(std::move(node_path)), road_path(std::move(road_path)) {} + + bool operator==(const Seq &other) const { + if (node_path.size() != other.node_path.size()) { + return false; + } + for (int i = 0; i < node_path.size(); ++i) { + if (node_path[i] != other.node_path[i]) { + return false; + } + } + return true; + } + + Seq patch(const std::vector &more_nodes, const std::vector &more_roads) const { + auto nodes = node_path; + nodes.reserve(nodes.size() + more_nodes.size()); + nodes.insert(nodes.end(), more_nodes.begin(), more_nodes.end()); + auto roads = road_path; + roads.reserve(roads.size() + more_roads.size()); + roads.insert(roads.end(), more_roads.begin(), more_roads.end()); + return Seq(std::move(nodes), std::move(roads)); + } +}; + +namespace std { +template <> +struct hash { + size_t operator()(const Seq &s) const noexcept { return hash_vector()(s.node_path); } +}; +} // namespace std + +inline bool __equals(const std::vector &seq0, int i0, int j0, // + const std::vector &seq1, int i1, int j1) { + if (j0 - i0 != j1 - i1) { + return false; + } + while (i0 < j0) { + if (seq0[i0] != seq1[i1]) { + return false; + } + ++i0; + ++i1; + } + return true; +} + +namespace cubao { + +// https://github.com/isl-org/Open3D/blob/88693971ae7a7c3df27546ff7c5b1d91188e39cf/cpp/open3d/utility/Helper.h#L71 +constexpr double neg_inf = -std::numeric_limits::infinity(); +constexpr double pos_inf = std::numeric_limits::infinity(); + +struct FastViterbi { + using LayerIndex = int; + using CandidateIndex = int; + using NodeIndex = std::tuple; + FastViterbi(int K, int N, const std::map, double> &scores) : K_(K), N_(N) { + if (K == 0 || N < 2) { + throw std::invalid_argument("invalid K, N = " + std::to_string(K) + ", " + std::to_string(N)); + } + links_ = std::vector>(N - 1, std::vector(K)); + for (auto &pair : scores) { + auto &curr = std::get<0>(pair.first); + auto &next = std::get<1>(pair.first); + auto lidx0 = std::get<0>(curr); + auto cidx0 = std::get<1>(curr); + auto lidx1 = std::get<0>(next); + auto cidx1 = std::get<1>(next); + double score = pair.second; + if (lidx0 < 0) { + if (lidx1 == 0 && cidx1 < K) { + heads_.push_back({cidx1, score}); + scores_[-1][-1][cidx1] = score; + } + continue; + } + if (lidx0 >= N || lidx1 != lidx0 + 1 || lidx1 >= N) { + continue; + } + if (cidx0 < 0 || cidx0 >= K || cidx1 < 0 || cidx1 >= K) { + continue; + } + links_[lidx0][cidx0].push_back({cidx1, score}); + scores_[lidx0][cidx0][cidx1] = score; + } + } + + std::vector scores(const std::vector &node_path) const { + if (node_path.size() != N_) { + return {}; + } + std::vector ret; + ret.reserve(N_); + double acc = scores_.at(-1).at(-1).at(node_path[0]); + ret.push_back(acc); + for (int n = 0; n < N_ - 1; ++n) { + acc += scores_.at(n).at(node_path[n]).at(node_path[n + 1]); + ret.push_back(acc); + } + return ret; + } + + std::tuple> inference() const { + // forward + // backward + return std::make_tuple(pos_inf, std::vector{}); + } + + bool setup_roads(const std::vector> &roads) { + if (roads.size() != N_) { + std::cerr << "invalid roads, #layers=" << roads.size() << " != " << N_ << std::endl; + return false; + } + roads_ = std::vector>(N_, std::vector(K_, (int64_t)-1)); + for (int n = 0; n < N_; ++n) { + int K = roads[n].size(); + if (K > K_) { + roads_.clear(); + std::cerr << "invalid road ids at #layer=" << n << ", #candidates=" << K << std::endl; + return false; + } + for (int k = 0; k < K; ++k) { + roads_[n][k] = roads[n][k]; + } + } + return true; + } + + bool setup_shortest_road_paths(const std::map, std::vector> &sp_paths) { + if (roads_.empty()) { + std::cerr << "roads not setup" << std::endl; + return false; + } + for (auto &pair : sp_paths) { + auto &curr = std::get<0>(pair.first); + auto &next = std::get<1>(pair.first); + auto lidx0 = std::get<0>(curr); + auto cidx0 = std::get<1>(curr); + auto lidx1 = std::get<0>(next); + auto cidx1 = std::get<1>(next); + auto &path = pair.second; + if (path.empty()) { + std::cerr << "empty path" << std::endl; + sp_paths_.clear(); + return false; + } + if (lidx0 < 0) { + if (lidx1 == 0 && cidx1 < K_) { + if (path.size() == 1 && path[0] == roads_[0][cidx1]) { + sp_paths_[-1][-1][cidx1] = path; + } else { + std::cerr << "sp_path not match roads" << std::endl; + sp_paths_.clear(); + return false; + } + } + continue; + } + if (lidx0 >= N_ || lidx1 != lidx0 + 1 || lidx1 >= N_) { + continue; + } + if (cidx0 < 0 || cidx0 >= K_ || cidx1 < 0 || cidx1 >= K_) { + continue; + } + if (path.front() == roads_[lidx0][cidx0] && path.back() == roads_[lidx1][cidx1]) { + // scores_.at(lidx0).at(cidx0).at(cidx1) + sp_paths_[lidx0][cidx0][cidx1] = path; + } else { + std::cerr << "sp_path not match roads" << std::endl; + sp_paths_.clear(); + return false; + } + } + return true; + } + + std::tuple, std::vector> inference(const std::vector &road_path) const { + if (roads_.empty() || sp_paths_.empty()) { + return std::make_tuple(pos_inf, std::vector{}, std::vector{}); + } + if (road_path.empty()) { + return std::make_tuple(pos_inf, std::vector{}, std::vector{}); + } + std::vector> prev_paths(K_); + for (auto &pair : heads_) { + auto nid = pair.first; + auto rid = roads_[0][nid]; + if (rid != road_path[0]) { + continue; + } + prev_paths[nid].insert(Seq({nid}, {rid})); + } + const int N = road_path.size(); + for (int n = 0; n < N_ - 1; ++n) { + std::vector> curr_paths(K_); + auto &paths = sp_paths_.at(n); + auto &layer = links_[n]; + for (int i = 0; i < K_; ++i) { + const auto &heads = prev_paths[i]; + if (heads.empty() || layer[i].empty()) { + continue; + } + auto &p = paths.at(i); + for (auto &pair : layer[i]) { + int j = pair.first; + const auto &sig = p.at(j); + if (sig.size() == 1) { + for (auto &seq : heads) { + curr_paths[j].insert(seq.patch({j}, {})); + } + continue; + } + for (auto &seq : heads) { + auto &path = seq.road_path; + int I = path.size(); + int J = I + sig.size() - 1; + if (!__equals(sig, 1, sig.size(), road_path, I, J)) { + continue; + } + curr_paths[j].insert(seq.patch({j}, std::vector(sig.begin() + 1, sig.end()))); + } + } + } + prev_paths = std::move(curr_paths); + } + + std::vector all_paths; + for (auto &paths : prev_paths) { + for (auto &path : paths) { + if (path.road_path.size() != N) { + continue; + } + all_paths.push_back(std::move(path)); + } + } + if (all_paths.empty()) { + return std::make_tuple(pos_inf, std::vector{}, std::vector{}); + } + + double max_score = neg_inf; + int best_path = -1; + for (int i = 0; i < all_paths.size(); ++i) { + double score = calc_score(all_paths[i]); + if (score > max_score) { + max_score = score; + best_path = i; + } + } + if (best_path < 0) { + return std::make_tuple(pos_inf, std::vector{}, std::vector{}); + } + return std::make_tuple(max_score, // + all_paths[best_path].node_path, // + all_paths[best_path].road_path); + } + + private: + const int K_{-1}; + const int N_{-1}; + using Links = std::vector>; + // head layer: cidx, score + Links heads_; + // tail layers, [[cidx (in next layer), score]] + std::vector> links_; + // score map, lidx -> cidx -> next_cidx -> score + std::unordered_map>> scores_; + + // road ids, K * N + std::vector> roads_; + // sp_paths, lidx -> cidx -> next_cidx -> sp_path (road seq) + std::unordered_map>>> sp_paths_; + + double calc_score(const Seq &seq) const { + auto &nodes = seq.node_path; + if (nodes.empty()) { + return neg_inf; + } + double score = scores_.at(-1).at(-1).at(nodes[0]); + for (int n = 0; n < nodes.size() - 1; ++n) { + int i = nodes[n]; + int j = nodes[n + 1]; + score += scores_.at(n).at(i).at(j); + } + return score; + } +}; +} // namespace cubao PYBIND11_MODULE(_core, m) { m.doc() = R"pbdoc( @@ -34,6 +348,25 @@ PYBIND11_MODULE(_core, m) { Some other explanation about the subtract function. )pbdoc"); + using FastViterbi = cubao::FastViterbi; + using NodeIndex = FastViterbi::NodeIndex; + py::class_(m, "FastViterbi", py::module_local(), py::dynamic_attr()) // + .def(py::init, double> &>(), // + "K"_a, "N"_a, "scores"_a) + // + .def("scores", &FastViterbi::scores, "node_path") + // + .def("inference", py::overload_cast<>(&FastViterbi::inference, py::const_)) + // + .def("setup_roads", &FastViterbi::setup_roads, "roads"_a) + .def("setup_shortest_road_paths", &FastViterbi::setup_shortest_road_paths, "sp_paths"_a) + // + .def("inference", py::overload_cast &>(&FastViterbi::inference, py::const_), + "road_path"_a, py::call_guard()) + + // + ; + #ifdef VERSION_INFO m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO); #else diff --git a/tests/test_basic.py b/tests/test_basic.py index 93da71b..4626776 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -1,10 +1,10 @@ from __future__ import annotations -import scikit_build_example as m +import fast_viterbi as m def test_version(): - assert m.__version__ == "0.0.1" + assert m.__version__ == "0.0.2" def test_add():