find_nearest_neighbors.hpp

Go to the documentation of this file.

#ifndef KNNCOLLE_FIND_NEAREST_NEIGHBORS_HPP
#define KNNCOLLE_FIND_NEAREST_NEIGHBORS_HPP
 
#include <vector>
#include <utility>
#include <type_traits>
 
#include "Prebuilt.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 constexpr(std::is_signed<Index_>::value) {
        if (k < num_observations) {
            return k;
        }
    } else {
        if (static_cast<typename std::make_unsigned<Index_>::type>(k) < num_observations) {
            return k;
        }
    }
    if (num_observations) {
        return num_observations - 1;
    }
    return 0;
}
 
template<typename Index_>
int cap_k_query(int k, Index_ num_observations) {
    if constexpr(std::is_signed<Index_>::value) {
        if (k <= num_observations) {
            return k;
        }
    } else {
        if (static_cast<typename std::make_unsigned<Index_>::type>(k) <= num_observations) {
            return k;
        }
    }
    return 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) {
    Index_ nobs = index.num_observations();
    k = cap_k(k, nobs);
    NeighborList<Index_, Distance_> output(nobs);
 
    parallelize(num_threads, nobs, [&](int, Index_ start, Index_ length) -> void {
        auto sptr = index.initialize();
        std::vector<Index_> indices;
        std::vector<Distance_> distances;
        for (Index_ i = start, end = start + length; i < end; ++i) {
            sptr->search(i, k, &indices, &distances);
            int actual_k = indices.size();
            output[i].reserve(actual_k);
            for (int 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) {
    Index_ nobs = index.num_observations();
    k = cap_k(k, nobs);
    std::vector<std::vector<Index_> > output(nobs);
 
    parallelize(num_threads, nobs, [&](int, Index_ start, Index_ length) -> void {
        auto sptr = index.initialize();
        for (Index_ i = start, end = start + length; i < end; ++i) {
            sptr->search(i, k, &(output[i]), NULL);
        }
    });
 
    return output;
}
 
}
 
#endif