nenesub.hpp

Go to the documentation of this file.

#ifndef NENESUB_HPP
#define NENESUB_HPP
 
#include <vector>
#include <queue>
#include <cstddef>
#include <algorithm>
#include <type_traits>
 
#include "knncolle/knncolle.hpp"
#include "sanisizer/sanisizer.hpp"
 
namespace nenesub {
 
struct Options {
    int num_neighbors = 20;
 
    int min_remaining = 10;
 
    int num_threads = 10;
};
 
template<typename Input_>
std::remove_cv_t<std::remove_reference_t<Input_> > I(Input_ x) {
    return x;
}
template<typename Index_, class GetNeighbors_, class GetIndex_, class GetMaxDistance_>
void compute(const Index_ num_obs, const GetNeighbors_ get_neighbors, const GetIndex_ get_index, const GetMaxDistance_ get_max_distance, const Options& options, std::vector<Index_>& selected) {
    typedef decltype(I(get_max_distance(0))) Distance;
    struct Payload {
        Payload(const Index_ identity, const Index_ remaining, const Distance max_distance) : remaining(remaining), identity(identity), max_distance(max_distance) {}
        Index_ remaining;
        Index_ identity;
        Distance max_distance;
    };
 
    const auto cmp = [](const Payload& left, const Payload& right) -> bool {
        if (left.remaining == right.remaining) {
            if (left.max_distance == right.max_distance) {
                return left.identity > right.identity; // smallest identities show up first.
            }
            return left.max_distance > right.max_distance; // smallest distances show up first.
        }
        return left.remaining < right.remaining; // largest remaining show up first.
    };
    std::priority_queue<Payload, std::vector<Payload>, decltype(I(cmp))> store(
        cmp,
        [&]{
            std::vector<Payload> container;
            container.reserve(num_obs);
            return container;
        }()
    );
 
    auto reverse_map = sanisizer::create<std::vector<std::vector<Index_> > >(num_obs);
    auto remaining = sanisizer::create<std::vector<Index_> >(num_obs);
    for (Index_ c = 0; c < num_obs; ++c) {
        const auto& neighbors = get_neighbors(c);
        const Index_ nneighbors = neighbors.size();
 
        if (nneighbors) { // protect get_max_distance just in case there are no neighbors.
            store.emplace(c, nneighbors, get_max_distance(c));
            for (Index_ n = 0; n < nneighbors; ++n) {
                reverse_map[get_index(neighbors, n)].push_back(c);
            }
            remaining[c] = nneighbors;
        }
    }
 
    selected.clear();
    auto tainted = sanisizer::create<std::vector<unsigned char> >(num_obs);
    Index_ min_remaining = options.min_remaining;
    while (!store.empty()) {
        auto payload = store.top();
        store.pop();
        if (tainted[payload.identity]) { // it's already in another selected point's local neighborhood.
            continue;
        }
 
        const Index_ new_remaining = remaining[payload.identity];
        if (new_remaining < min_remaining) { // it can't be valid so we skip it.
            continue;
        }
 
        payload.remaining = new_remaining;
        if (!store.empty() && cmp(payload, store.top())) { // cycling it back into the queue with an updated remaining count, if it's not better than the queue's best.
            store.push(payload);
            continue;
        }
 
        selected.push_back(payload.identity);
        tainted[payload.identity] = 1;
        for (const auto x : reverse_map[payload.identity]) {
            --remaining[x];
        }
 
        const auto& neighbors = get_neighbors(payload.identity);
        const Index_ nneighbors = neighbors.size();
        for (Index_ n = 0; n < nneighbors; ++n) {
            const auto current = get_index(neighbors, n);
            if (tainted[current]) {
                continue;
            }
            tainted[current] = 1;
            for (const auto x : reverse_map[current]) {
                --remaining[x];
            }
        }
    }
 
    std::sort(selected.begin(), selected.end());
}
 
template<typename Index_, typename Distance_>
std::vector<Index_> compute(const knncolle::NeighborList<Index_, Distance_>& neighbors, const Options& options) {
    std::vector<Index_> output;
    compute(
        static_cast<Index_>(neighbors.size()),
        [&](const Index_ i) -> const auto& { return neighbors[i]; }, 
        [](const std::vector<std::pair<Index_, Distance_> >& x, const Index_ n) -> Index_ { return x[n].first; }, 
        [&](const Index_ i) -> Distance_ { return neighbors[i].back().second; }, 
        options,
        output
    );
    return output;
}
 
template<typename Index_, typename Input_, typename Distance_>
std::vector<Index_> compute(const knncolle::Prebuilt<Index_, Input_, Distance_>& prebuilt, const Options& options) {
    const int k = options.num_neighbors;
    if (k < options.min_remaining) {
        throw std::runtime_error("number of neighbors is less than 'min_remaining'");
    }
 
    const Index_ nobs = prebuilt.num_observations();
    const auto capped_k = knncolle::cap_k(k, nobs);
    auto nn_indices = sanisizer::create<std::vector<std::vector<Index_> > >(nobs);
    auto max_distance = sanisizer::create<std::vector<Distance_> >(nobs);
 
    knncolle::parallelize(options.num_threads, nobs, [&](const int, const Index_ start, const Index_ length) -> void {
        const auto sptr = prebuilt.initialize();
        std::vector<Distance_> nn_distances;
        for (Index_ i = start, end = start + length; i < end; ++i) {
            sptr->search(i, capped_k, &(nn_indices[i]), &nn_distances);
            max_distance[i] = (capped_k ? nn_distances.back() : 0);
        }
    });
 
    std::vector<Index_> output;
    compute(
        nobs,
        [&](const Index_ i) -> const std::vector<Index_>& { return nn_indices[i]; }, 
        [](const std::vector<Index_>& x, const Index_ n) -> Index_ { return x[n]; }, 
        [&](const Index_ i) -> Distance_ { return max_distance[i]; },
        options,
        output
    );
    return output;
}
 
template<typename Index_, typename Input_, typename Distance_, class Matrix_ = knncolle::Matrix<Index_, Input_> >
std::vector<Index_> compute(
    const std::size_t num_dims, 
    const Index_ num_obs, 
    const Input_* data, 
    const knncolle::Builder<Index_, Input_, Distance_, Matrix_>& knn_method,
    const Options& options) 
{
    const auto prebuilt = knn_method.build_unique(knncolle::SimpleMatrix<Index_, Input_>(num_dims, num_obs, data));
    return compute(*prebuilt, options);
}
 
}
 
#endif