sorting in C++ side

Author: eloyfelix

FPSim2/FPSim2.py

@@ -2,7 +2,7 @@
 import numpy as np
 import tables as tb
 from .io import S_INDEXS, load_fps, load_query, get_bounds_range
-from .FPSim2lib import _similarity_search, _substructure_search
+from .FPSim2lib import _similarity_search, _substructure_search, sort_results
 
 
 def on_disk_search(query, fp_filename, chunk_indexes, threshold, s_index):
@@ -172,8 +172,6 @@ def _base_search(self, query_string, threshold, search_func, chunk_size, s_index
         if fp_range:
             if n_workers == 1 and not on_disk:
                 np_res = search_func(query, self.fps.fps, threshold, fp_range[0], fp_range[1])
-                if s_index != 'substructure':
-                    np_res[::-1].sort(order='coeff')
             else:
                 results = self._parallel_run(query,
                                              search_func,
@@ -187,7 +185,7 @@ def _base_search(self, query_string, threshold, search_func, chunk_size, s_index
                 if results:
                     np_res = np.concatenate(results)
                     if s_index != 'substructure':
-                        np_res[::-1].sort(order='coeff')
+                        sort_results(np_res)
                 else:
                     np_res = empty_np
         else:

FPSim2/src/sim.cpp

@@ -108,6 +108,15 @@ py::array_t<uint32_t> _substructure_search(py::array_t<unsigned long long> pyque
     return py::array_t<uint32_t>(results->size(), results->data(), capsule);
 }
 
+bool cmp(const Result &l, const Result &r) { return l.coeff > r.coeff; }
+
+void sort_results(py::array_t<Result> pyres)
+{
+    py::buffer_info buf = pyres.request();
+    Result *ptr = (Result *) buf.ptr;
+    std::sort(&ptr[0], &ptr[buf.shape[0]], cmp);
+}
+
 py::array_t<Result> _similarity_search(py::array_t<unsigned long long> pyquery,
                                        py::array_t<unsigned long long> pydb,
                                        float threshold,
@@ -153,7 +162,9 @@ py::array_t<Result> _similarity_search(py::array_t<unsigned long long> pyquery,
         int_count = 0;
         i++;
     }
+    // set final size and sort
     results->resize(total_sims);
+    std::sort(results->begin(), results->end(), cmp);
 
     // acquire the GIL
     py::gil_scoped_acquire acquire;

FPSim2/src/sim.hpp

@@ -17,6 +17,8 @@ __inline float tanimoto_coeff(uint32_t int_count,
 
 uint32_t py_popcount(py::array_t<unsigned long long> pyquery);
 
+void sort_results(py::array_t<Result> pyres);
+
 py::array_t<uint32_t> _substructure_search(py::array_t<unsigned long long> pyquery,
                                            py::array_t<unsigned long long> pydb,
                                            float threshold,

FPSim2/src/wraps.cpp

@@ -7,7 +7,8 @@ namespace py = pybind11;
 PYBIND11_MODULE(FPSim2lib, m)
 {
     PYBIND11_NUMPY_DTYPE(Result, mol_id, coeff);
+    m.def("py_popcount", &py_popcount, py::call_guard<py::gil_scoped_release>());
+    m.def("sort_results", &sort_results);
     m.def("_substructure_search", &_substructure_search);
     m.def("_similarity_search", &_similarity_search);
-    m.def("py_popcount", &py_popcount, py::call_guard<py::gil_scoped_release>());
 }