Fix flann

biovault · Apr 25, 2024 · f95cdfe · f95cdfe
1 parent 5699c32
commit f95cdfe
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/hdi/dimensionality_reduction/hd_joint_probability_generator_inl.h b/hdi/dimensionality_reduction/hd_joint_probability_generator_inl.h
@@ -194,17 +194,18 @@ namespace hdi {
           else
           {
             std::vector<size_t> flannIndices;
-            flannIndices.reserve(indices.size());
-            std::transform(indices.cbegin(), indices.cend(), std::back_inserter(flannIndices),
-              [](auto value) { return static_cast<size_t>(value); });
+            flannIndices.resize(indices.size());
             flann::Matrix<size_t> indices_mat(flannIndices.data(), query.rows, nn);
             index.knnSearch(query, indices_mat, dists_mat, nn, flann_params);
+#pragma omp parallel for
+            for (std::int64_t i = 0; i < indices.size(); ++i)
+              indices[i] = static_cast<int_type>(flannIndices[i]);
           }
         }
       }
       else if (params._aknn_algorithm == hdi::dr::KNN_HNSW)
       {
-        hdi::utils::secureLog(_logger, "Computing approximated knn with HNSWLIB...");
+        hdi::utils::secureLog(_logger, "Computing approximated knn with HNSW Lib...");
 
         hnswlib::SpaceInterface<scalar> *space = NULL;
         switch (params._aknn_metric) {

diff --git a/hdi/dimensionality_reduction/hierarchical_sne_inl.h b/hdi/dimensionality_reduction/hierarchical_sne_inl.h
@@ -288,11 +288,12 @@ namespace hdi {
           else
           {
             std::vector<size_t> flannIndices;
-            flannIndices.reserve(neighborhood_graph.size());
-            std::transform(neighborhood_graph.cbegin(), neighborhood_graph.cend(), std::back_inserter(flannIndices),
-              [](auto value) { return static_cast<size_t>(value); });
+            flannIndices.resize(neighborhood_graph.size());
             flann::Matrix<size_t> indices_mat(flannIndices.data(), query.rows, nn);
             index.knnSearch(query, indices_mat, dists_mat, nn, flann_params);
+#pragma omp parallel for
+            for (std::int64_t i = 0; i < neighborhood_graph.size(); ++i)
+              neighborhood_graph[i] = static_cast<int_type>(flannIndices[i]);
           }
         }
       }