switching to nclist search

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## Version 0.5.4
+
+- Switch to the nearest implementation from [LTLA/nclist-cpp](https://github.com/LTLA/nclist-cpp).
+- Improved worksharing across threads for `find_overlaps_groups` and `nearest_groups`. The idea is to ensure that all threads are occupied, even when the groups are of different sizes. This is done with some more fine-grained parallelization during both index building and overlap identification.
+
 ## Version 0.5.0 - 0.5.3
 
 - Switch to [LTLA/nclist-cpp](https://github.com/LTLA/nclist-cpp) for overlap and intersection operations.

lib/CMakeLists.txt

@@ -11,8 +11,8 @@ include(FetchContent)
 FetchContent_Declare(
   nclist
   GIT_REPOSITORY https://github.com/LTLA/nclist-cpp
-  GIT_TAG v0.2.0
-  # GIT_TAG master
+  # GIT_TAG v0.2.0
+  GIT_TAG master
 )
 
 FetchContent_MakeAvailable(nclist)

lib/src/nclssearch.cpp

@@ -30,34 +30,31 @@ struct NCListSearchHandler {
 
         self_starts.assign(starts_ptr, starts_ptr + n);
         self_ends.assign(ends_ptr, ends_ptr + n);
-
-        sorted_starts.resize(n);
-        sorted_ends.resize(n);
-        for (Index i = 0; i < n; ++i) {
-            sorted_starts[i] = {self_starts[i], i};
-            sorted_ends[i] = {self_ends[i], i};
-        }
-
-        std::sort(sorted_starts.begin(), sorted_starts.end());
-        std::sort(sorted_ends.begin(), sorted_ends.end());
     }
 
     nclist::Nclist<Index, Position> nclist_obj;
 
     std::vector<Position> self_starts;
     std::vector<Position> self_ends;
-    std::vector<std::pair<Position, Index>> sorted_starts;
-    std::vector<std::pair<Position, Index>> sorted_ends;
 };
 
 py::object perform_follow(
     NCListSearchHandler &self,
     py::array_t<Position> query_starts,
     const std::string& select,
     int num_threads = 1) {
+
     auto q_starts_ptr = static_cast<const Position*>(query_starts.request().ptr);
     Index n_queries = query_starts.request().shape[0];
 
+    std::vector<std::pair<Position, Index>> sorted_ends;
+    sorted_ends.resize(self.self_ends.size());
+    for (Index i = 0; i < self.self_ends.size(); ++i) {
+        sorted_ends[i] = {self.self_ends[i], i};
+    }
+
+    std::sort(sorted_ends.begin(), sorted_ends.end());
+
     std::vector<Index> results(n_queries);
     std::vector<std::thread> workers;
     workers.reserve(num_threads);
@@ -77,9 +74,9 @@ py::object perform_follow(
                 Position q_start = q_starts_ptr[i];
 
                 // Binary search on sorted ends to find the last subject ending before the query starts
-                auto it = std::upper_bound(self.sorted_ends.begin(), self.sorted_ends.end(), std::make_pair(q_start, Index(-1)));
+                auto it = std::upper_bound(sorted_ends.begin(), sorted_ends.end(), std::make_pair(q_start, Index(-1)));
 
-                if (it == self.sorted_ends.begin()) {
+                if (it == sorted_ends.begin()) {
                     results[i] = -1; // No preceding range
                 } else {
                     --it;
@@ -113,9 +110,18 @@ py::object perform_precede(
     py::array_t<Position> query_ends,
     const std::string& select,
     int num_threads = 1) {
+
     auto q_ends_ptr = static_cast<const Position*>(query_ends.request().ptr);
     Index n_queries = query_ends.request().shape[0];
 
+    std::vector<std::pair<Position, Index>> sorted_starts;
+    sorted_starts.resize(self.self_starts.size());
+    for (Index i = 0; i < self.self_starts.size(); ++i) {
+        sorted_starts[i] = {self.self_starts[i], i};
+    }
+
+    std::sort(sorted_starts.begin(), sorted_starts.end());
+
     std::vector<Index> results(n_queries);
     std::vector<std::thread> workers;
     workers.reserve(num_threads);
@@ -135,9 +141,9 @@ py::object perform_precede(
                 Position q_end = q_ends_ptr[i];
 
                 // Binary search on sorted starts to find the first subject starting after the query ends
-                auto it = std::lower_bound(self.sorted_starts.begin(), self.sorted_starts.end(), std::make_pair(q_end, Index(-1)));
+                auto it = std::lower_bound(sorted_starts.begin(), sorted_starts.end(), std::make_pair(q_end, Index(-1)));
 
-                if (it == self.sorted_starts.end()) {
+                if (it == sorted_starts.end()) {
                     results[i] = -1; // No following range
                 } else {
                     results[i] = it->second;
@@ -170,11 +176,18 @@ py::object perform_nearest(
     py::array_t<Position> query_starts,
     py::array_t<Position> query_ends,
     const std::string& select,
-    int num_threads=1) {
+    int num_threads=1,
+    bool adjacent_equals_overlap=false) {
+
     auto q_starts_ptr = static_cast<const Position*>(query_starts.request().ptr);
     auto q_ends_ptr = static_cast<const Position*>(query_ends.request().ptr);
     Index n_queries = query_starts.request().shape[0];
 
+    bool quit_on_first = (select == "arbitrary");
+    if (select != "all" && select != "arbitrary" && !quit_on_first) {
+        throw std::runtime_error("Invalid 'select' parameter. Must be 'all' or 'arbitrary'.");
+    }
+
     std::vector<std::vector<Index>> all_results(n_queries);
     std::vector<std::thread> workers;
     workers.reserve(num_threads);
@@ -190,57 +203,23 @@ py::object perform_nearest(
         }
 
         workers.emplace_back([&](int first, int length) -> void {
-            nclist::OverlapsAnyWorkspace<Index> ws_any;
-            nclist::OverlapsAnyParameters<Position> ol_params;
-            ol_params.min_overlap = 1;
+            std::vector<Index> nearest_matches;
 
             for (Index i = first, last = first + length; i < last; ++i) {
-                Position q_start = q_starts_ptr[i];
-                Position q_end = q_ends_ptr[i];
-
-                std::vector<std::pair<Position, Index>> candidates;
-
-                std::vector<Index> overlaps;
-                nclist::overlaps_any(self.nclist_obj, q_start, q_end, ol_params, ws_any, overlaps);
-                for (const auto& ol_idx : overlaps) {
-                    candidates.emplace_back(0, ol_idx);
-                }
-
-                auto it_b = std::upper_bound(self.sorted_ends.begin(), self.sorted_ends.end(), std::make_pair(q_start, std::numeric_limits<Index>::max()));
-                if (it_b != self.sorted_ends.begin()) {
-                    --it_b;
-                    candidates.emplace_back(q_start - it_b->first, it_b->second);
-                }
-
-                auto it_a = std::lower_bound(self.sorted_starts.begin(), self.sorted_starts.end(), std::make_pair(q_end, Index(-1)));
-                if (it_a != self.sorted_starts.end()) {
-                    candidates.emplace_back(it_a->first - q_end, it_a->second);
-                }
-
-                if (candidates.empty()) {
-                    continue;
-                }
-
-                Position min_dist = candidates[0].first;
-                for (size_t k = 1; k < candidates.size(); ++k) {
-                    if (candidates[k].first < min_dist) {
-                        min_dist = candidates[k].first;
+                nclist::NearestWorkspace<Index> ws_nearest;
+                nclist::NearestParameters<Position> params;
+                params.quit_on_first = quit_on_first;
+                params.adjacent_equals_overlap = adjacent_equals_overlap;
+
+                nclist::nearest(self.nclist_obj, q_starts_ptr[i], q_ends_ptr[i], params, ws_nearest, nearest_matches);
+
+                if (!nearest_matches.empty()) {
+                    if (select == "arbitrary" && !quit_on_first) {
+                        all_results[i] = {nearest_matches.back()};
+                    } else {
+                        all_results[i] = nearest_matches;
                     }
                 }
-
-                std::vector<Index> best_hits;
-                for (const auto& cand : candidates) {
-                    if (cand.first == min_dist) {
-                        best_hits.push_back(cand.second);
-                    }
-                }
-
-                if (select == "arbitrary") {
-                    std::sort(best_hits.begin(), best_hits.end());
-                    all_results[i] = {best_hits[0]};
-                } else {
-                    all_results[i] = best_hits;
-                }
             }
         }, jobs_so_far, current_jobs);
         jobs_so_far += current_jobs;
@@ -280,6 +259,166 @@ py::object perform_nearest(
     }
 }
 
+struct NearestGroupInfo {
+    const Index* ptr;
+    std::size_t size;
+};
+
+pybind11::tuple perform_nearest_groups(
+    py::array_t<Position> self_starts,
+    py::array_t<Position> self_ends,
+    const std::vector<py::array_t<Index>>& self_groups,
+    py::array_t<Position> query_starts,
+    py::array_t<Position> query_ends,
+    const std::vector<py::array_t<Index>>& query_groups,
+    const std::string& select,
+    int num_threads=1,
+    bool adjacent_equals_overlap=false) {
+
+    auto s_starts_ptr = static_cast<const Position*>(self_starts.request().ptr);
+    auto s_ends_ptr = static_cast<const Position*>(self_ends.request().ptr);
+    auto q_starts_ptr = static_cast<const Position*>(query_starts.request().ptr);
+    auto q_ends_ptr = static_cast<const Position*>(query_ends.request().ptr);
+    std::size_t n_groups = self_groups.size();
+    if (n_groups != query_groups.size()) {
+        throw std::runtime_error("The number of self/subject groups must be equal to the number of query groups.");
+    }
+
+    bool quit_on_first = (select == "arbitrary");
+    if (select != "all" && select != "arbitrary" && !quit_on_first) {
+        throw std::runtime_error("Invalid 'select' parameter. Must be 'all' or 'arbitrary'.");
+    }
+
+    std::vector<NearestGroupInfo> self_group_info(n_groups);
+    std::vector<NearestGroupInfo> query_group_info(n_groups);
+    for (std::size_t i = 0; i < n_groups; ++i) {
+        auto s_req = self_groups[i].request();
+        self_group_info[i] = {static_cast<const Index*>(s_req.ptr), static_cast<std::size_t>(s_req.shape[0])};
+        auto q_req = query_groups[i].request();
+        query_group_info[i] = {static_cast<const Index*>(q_req.ptr), static_cast<std::size_t>(q_req.shape[0])};
+    }
+
+    // Building the indices in parallel with a simple worker pool.
+    std::vector<nclist::Nclist<Index, Position> > built(n_groups);
+    {
+        std::vector<std::thread> workers;
+        workers.reserve(num_threads);
+        std::mutex mut;
+        std::size_t group = 0;
+
+        for (int i = 0; i < num_threads; ++i) {
+            workers.emplace_back([&]() -> void {
+                while (1) {
+                    std::unique_lock lck(mut);
+                    if (group == n_groups) {
+                        break;
+                    }
+                    const auto curgroup = group;
+                    ++group;
+                    lck.unlock();
+
+                    const auto& s_info = self_group_info[curgroup];
+                    if (s_info.size != 0) {
+                        built[curgroup] = nclist::build<Index, Position>(s_info.size, s_info.ptr, s_starts_ptr, s_ends_ptr);
+                    }
+                }
+            });
+        }
+
+        for (auto& worker : workers) {
+            worker.join();
+        }
+    }
+
+    // Now running through all groups to find overlaps in parallel.
+    std::vector<std::vector<std::vector<Index> > > all_group_results(n_groups);
+    std::vector<std::thread> workers;
+    workers.reserve(num_threads);
+    std::vector<std::size_t> total_hits_per_thread(num_threads);
+
+    for (std::size_t g = 0; g < n_groups; ++g) {
+        const auto& s_info = self_group_info[g];
+        const auto& q_info = query_group_info[g];
+        if (s_info.size == 0 || q_info.size == 0) {
+            continue;
+        }
+
+        const std::size_t num_jobs = q_info.size / num_threads;
+        const std::size_t num_remaining = q_info.size % num_threads;
+        std::size_t jobs_so_far = 0;
+
+        const auto& nclist_obj = built[g];
+        auto& current_group_results = all_group_results[g];
+        current_group_results.resize(q_info.size);
+
+        workers.clear();
+        for (int i = 0; i < num_threads; ++i) {
+            const std::size_t current_jobs = num_jobs + (i < num_remaining);
+            if (current_jobs == 0) {
+                break;
+            }
+
+            workers.emplace_back([&](int thread, std::size_t first, std::size_t length) -> void {
+                std::size_t current_total_hits = 0;
+
+                for (std::size_t k = first, last = first + length; k < last; ++k) {
+                    const Index original_query_idx = q_info.ptr[k];
+                    auto& single_query_matches = current_group_results[k];
+
+                    nclist::NearestWorkspace<Index> ws_nearest;
+                    nclist::NearestParameters<Position> params;
+                    params.quit_on_first = quit_on_first;
+                    params.adjacent_equals_overlap = adjacent_equals_overlap;
+
+                    nclist::nearest(nclist_obj, q_starts_ptr[original_query_idx], q_ends_ptr[original_query_idx], params, ws_nearest, single_query_matches);
+
+                    if (!single_query_matches.empty()) {
+                        if (select == "arbitrary" && !quit_on_first) {
+                            single_query_matches.front() = single_query_matches.back();
+                            single_query_matches.resize(1);
+                        }
+                    }
+
+                    current_total_hits += single_query_matches.size();
+                }
+
+                // Don't add directly to this value in the inner loop, to reduce the risk of false sharing.
+                total_hits_per_thread[thread] += current_total_hits;
+            }, i, jobs_so_far, current_jobs);
+
+            jobs_so_far += current_jobs;
+        }
+
+        for (auto& worker : workers) {
+            worker.join();
+        }
+    }
+
+    const std::size_t total_hits = std::accumulate(total_hits_per_thread.begin(), total_hits_per_thread.end(), static_cast<std::size_t>(0));
+    py::array_t<Index> query_hits(total_hits);
+    py::array_t<Index> self_hits(total_hits);
+    auto q_res_ptr = static_cast<Index*>(query_hits.request().ptr);
+    auto s_res_ptr = static_cast<Index*>(self_hits.request().ptr);
+
+    std::size_t current_pos = 0;
+    for (std::size_t group_idx = 0, all_group_size=all_group_results.size(); group_idx < all_group_size; ++group_idx) {
+        const auto& group_res = all_group_results[group_idx];
+        const auto& q_info = query_group_info[group_idx];
+
+        for (std::size_t query_in_group_idx = 0, group_res_size = group_res.size(); query_in_group_idx < group_res_size; ++query_in_group_idx) {
+            const auto& relative_subject_matches = group_res[query_in_group_idx];
+            if (!relative_subject_matches.empty()) {
+                Index original_query_idx = q_info.ptr[query_in_group_idx];
+                std::copy(relative_subject_matches.begin(), relative_subject_matches.end(), s_res_ptr + current_pos);
+                std::fill(q_res_ptr + current_pos, q_res_ptr + current_pos + relative_subject_matches.size(), original_query_idx);
+                current_pos += relative_subject_matches.size();
+            }
+        }
+    }
+
+    return py::make_tuple(query_hits, self_hits);
+}
+
 
 void init_nclistsearch(pybind11::module &m){
 
@@ -304,5 +443,18 @@ void init_nclistsearch(pybind11::module &m){
             py::arg("query_ends"),
             py::arg("select") = "arbitrary",
             py::arg("num_threads") = 1,
+            py::arg("adjacent_equals_overlap") = false,
             "Find nearest ranges in both directions.");
+
+    m.def("nearest_groups", &perform_nearest_groups,
+        py::arg("self_starts"),
+        py::arg("self_ends"),
+        py::arg("self_groups"),
+        py::arg("query_starts"),
+        py::arg("query_ends"),
+        py::arg("query_groups"),
+        py::arg("select") = "arbitrary",
+        py::arg("num_threads") = 1,
+        py::arg("adjacent_equals_overlap") = false,
+        "Find nearest ranges in both directions, respecting group boundaries.");
 }