find_nearest_neighbors.hpp

Go to the documentation of this file.

#ifndef KNNCOLLE_FIND_NEAREST_NEIGHBORS_HPP
#define KNNCOLLE_FIND_NEAREST_NEIGHBORS_HPP
 
#include "Prebuilt.hpp"
 
#include <vector>
#include <utility>
#include <type_traits>
 
#include "sanisizer/sanisizer.hpp"
#ifndef KNNCOLLE_CUSTOM_PARALLEL
#include "subpar/subpar.hpp"
#endif
 
 
namespace knncolle {
 
template<typename Task_, class Run_>
void parallelize(int num_workers, Task_ num_tasks, Run_ run_task_range) {
#ifndef KNNCOLLE_CUSTOM_PARALLEL
    // Don't make this nothrow_ = true, as the derived methods could do anything...
    subpar::parallelize(num_workers, num_tasks, std::move(run_task_range));
#else
    KNNCOLLE_CUSTOM_PARALLEL(num_workers, num_tasks, run_task_range);
#endif
}
 
template<typename Index_, typename Distance_>
using NeighborList = std::vector<std::vector<std::pair<Index_, Distance_> > >;
 
template<typename Index_>
int cap_k(int k, Index_ num_observations) {
    if (sanisizer::is_less_than(sanisizer::attest_gez(k), sanisizer::attest_gez(num_observations))) {
        return k;
    } else if (num_observations) {
        return num_observations - 1;
    } else {
        return 0;
    }
}
 
template<typename Index_>
int cap_k_query(int k, Index_ num_observations) {
    return sanisizer::min(sanisizer::attest_gez(k), sanisizer::attest_gez(num_observations));
}
 
template<typename Index_, typename Data_, typename Distance_>
NeighborList<Index_, Distance_> find_nearest_neighbors(const Prebuilt<Index_, Data_, Distance_>& index, int k, int num_threads = 1) {
    const Index_ nobs = index.num_observations();
    k = cap_k(k, nobs);
    auto output = sanisizer::create<NeighborList<Index_, Distance_> >(sanisizer::attest_gez(nobs));
 
    parallelize(num_threads, nobs, [&](int, Index_ start, Index_ length) -> void {
        auto sptr = index.initialize_known();
        std::vector<Index_> indices;
        std::vector<Distance_> distances;
        for (Index_ i = start, end = start + length; i < end; ++i) {
            sptr->search(i, k, &indices, &distances);
            const auto actual_k = indices.size();
            output[i].reserve(actual_k);
            for (I<decltype(actual_k)> j = 0; j < actual_k; ++j) {
                output[i].emplace_back(indices[j], distances[j]);
            }
        }
    });
 
    return output;
}
 
template<typename Index_, typename Data_, typename Distance_>
std::vector<std::vector<Index_> > find_nearest_neighbors_index_only(const Prebuilt<Index_, Data_, Distance_>& index, int k, int num_threads = 1) {
    const Index_ nobs = index.num_observations();
    k = cap_k(k, nobs);
    auto output = sanisizer::create<std::vector<std::vector<Index_> > >(sanisizer::attest_gez(nobs));
 
    parallelize(num_threads, nobs, [&](int, Index_ start, Index_ length) -> void {
        auto sptr = index.initialize_known();
        for (Index_ i = start, end = start + length; i < end; ++i) {
            sptr->search(i, k, &(output[i]), NULL);
        }
    });
 
    return output;
}
 
}
 
#endif