nenesub.hpp

Go to the documentation of this file.

#ifndef NENESUB_HPP
#define NENESUB_HPP
 
#include <vector>
#include <queue>
#include <cstdint>
#include <algorithm>
#include "knncolle/knncolle.hpp"
 
namespace nenesub {
 
struct Options {
    int num_neighbors = 20;
 
    int min_remaining = 10;
 
    int num_threads = 10;
};
 
template<typename Index_, class GetNeighbors_, class GetIndex_, class GetMaxDistance_>
void compute(Index_ num_obs, GetNeighbors_ get_neighbors, GetIndex_ get_index, GetMaxDistance_ get_max_distance, const Options& options, std::vector<Index_>& selected) {
    typedef decltype(get_max_distance(0)) Distance_;
    struct Payload {
        Payload(Index_ identity, Index_ remaining, Distance_ max_distance) : remaining(remaining), identity(identity), max_distance(max_distance) {}
        Index_ remaining;
        Index_ identity;
        Distance_ max_distance;
    };
 
    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(cmp)> store(
        cmp,
        [&]{
            std::vector<Payload> container;
            container.reserve(num_obs);
            return container;
        }()
    );
 
    std::vector<std::vector<Index_> > reverse_map(num_obs);
    std::vector<Index_> remaining(num_obs);
    for (Index_ c = 0; c < num_obs; ++c) {
        const auto& neighbors = get_neighbors(c);
        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();
    std::vector<uint8_t> tainted(num_obs);
    Index_ min_remaining = options.min_remaining;
    while (!store.empty()) {
        auto payload = store.top();
        store.pop();
        if (tainted[payload.identity]) {
            continue;
        }
 
        const auto& neighbors = get_neighbors(payload.identity);
        Index_ new_remaining = remaining[payload.identity];
 
        if (new_remaining >= min_remaining) {
            payload.remaining = new_remaining;
            if (!store.empty() && cmp(payload, store.top())) {
                store.push(payload);
            } else {
                selected.push_back(payload.identity);
                tainted[payload.identity] = 1;
                for (auto x : reverse_map[payload.identity]) {
                    --remaining[x];
                }
 
                Index_ nneighbors = neighbors.size();
                for (Index_ n = 0; n < nneighbors; ++n) {
                    auto current = get_index(neighbors, n);
                    tainted[current] = 1;
                    for (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()),
        [&](size_t i) -> const auto& { return neighbors[i]; }, 
        [](const auto& x, Index_ n) -> Index_ { return x[n].first; }, 
        [&](size_t i) -> Distance_ { return neighbors[i].back().second; }, 
        options,
        output
    );
    return output;
}
 
template<typename Dim_, typename Index_, typename Float_>
std::vector<Index_> compute(const knncolle::Prebuilt<Dim_, Index_, Float_>& prebuilt, const Options& options) {
    int k = options.num_neighbors;
    if (k < options.min_remaining) {
        throw std::runtime_error("number of neighbors is less than 'min_remaining'");
    }
 
    Index_ nobs = prebuilt.num_observations();
    std::vector<std::vector<Index_> > nn_indices(nobs);
    std::vector<Float_> max_distance(nobs);
 
    knncolle::parallelize(options.num_threads, nobs, [&](int, Index_ start, Index_ length) -> void {
        auto sptr = prebuilt.initialize();
        std::vector<Float_> nn_distances;
        for (Index_ i = start, end = start + length; i < end; ++i) {
            sptr->search(i, k, &(nn_indices[i]), &nn_distances);
            max_distance[i] = (k ? 0 : nn_distances.back());
        }
    });
 
    std::vector<Index_> output;
    compute(
        nobs,
        [&](size_t i) -> const std::vector<Index_>& { return nn_indices[i]; }, 
        [](const std::vector<Index_>& x, Index_ n) -> Index_ { return x[n]; }, 
        [&](size_t i) -> Float_ { return max_distance[i]; },
        options,
        output
    );
    return output;
}
 
template<typename Dim_, typename Index_, typename Value_, typename Float_>
std::vector<Index_> compute(
    Dim_ num_dims, 
    Index_ num_obs, 
    const Value_* data, 
    const knncolle::Builder<knncolle::SimpleMatrix<Dim_, Index_, Value_>, Float_>& knn_method,
    const Options& options) 
{
    auto prebuilt = knn_method.build_unique(knncolle::SimpleMatrix<Dim_, Index_, Value_>(num_dims, num_obs, data));
    return compute(*prebuilt, options);
}
 
}
 
#endif