compute.hpp

Go to the documentation of this file.

#ifndef MNNCORRECT_COMPUTE_HPP
#define MNNCORRECT_COMPUTE_HPP
 
#include <algorithm>
#include <vector>
#include <numeric>
#include <stdexcept>
#include <cstdint>
 
#include "knncolle/knncolle.hpp"
 
#include "AutomaticOrder.hpp"
#include "CustomOrder.hpp"
#include "Options.hpp"
#include "restore_order.hpp"
 
namespace mnncorrect {
 
struct Details {
    Details() = default;
 
    Details(std::vector<size_t> merge_order, std::vector<size_t> num_pairs) : merge_order(std::move(merge_order)), num_pairs(std::move(num_pairs)) {}
    std::vector<size_t> merge_order;
 
    std::vector<size_t> num_pairs;
};
 
namespace internal {
 
template<typename Dim_, typename Index_, typename Float_>
Details compute(size_t num_dim, const std::vector<size_t>& num_obs, const std::vector<const Float_*>& batches, Float_* output, const Options<Dim_, Index_, Float_>& options) {
    auto builder = options.builder;
    if (!builder) {
        builder.reset(new knncolle::VptreeBuilder<knncolle::EuclideanDistance, knncolle::SimpleMatrix<Dim_, Index_, Float_>, Float_>);
    }
 
    if (!options.order.empty()) {
        CustomOrder<Dim_, Index_, Float_> runner(num_dim, num_obs, batches, output, *builder, options.num_neighbors, options.order, options.mass_cap, options.num_threads);
        runner.run(options.num_mads, options.robust_iterations, options.robust_trim);
        return Details(runner.get_order(), runner.get_num_pairs());
 
    } else if (options.automatic_order) {
        AutomaticOrder<Dim_, Index_, Float_> runner(num_dim, num_obs, batches, output, *builder, options.num_neighbors, options.reference_policy, options.mass_cap, options.num_threads);
        runner.run(options.num_mads, options.robust_iterations, options.robust_trim);
        return Details(runner.get_order(), runner.get_num_pairs());
 
    } else {
        std::vector<size_t> trivial_order(num_obs.size());
        std::iota(trivial_order.begin(), trivial_order.end(), 0);
        CustomOrder<Dim_, Index_, Float_> runner(num_dim, num_obs, batches, output, *builder, options.num_neighbors, trivial_order, options.mass_cap, options.num_threads);
        runner.run(options.num_mads, options.robust_iterations, options.robust_trim);
        return Details(std::move(trivial_order), runner.get_num_pairs());
    }
}
 
}
template<typename Dim_, typename Index_, typename Float_>
Details compute(size_t num_dim, const std::vector<size_t>& num_obs, const std::vector<const Float_*>& batches, Float_* output, const Options<Dim_, Index_, Float_>& options) {
    auto stats = internal::compute(num_dim, num_obs, batches, output, options);
    internal::restore_order(num_dim, stats.merge_order, num_obs, output);
    return stats;
}
 
template<typename Dim_, typename Index_, typename Float_>
Details compute(size_t num_dim, const std::vector<size_t>& num_obs, const Float_* input, Float_* output, const Options<Dim_, Index_, Float_>& options) {
    std::vector<const Float_*> batches;
    batches.reserve(num_obs.size());
    for (auto n : num_obs) {
        batches.push_back(input);
        input += n * num_dim; // already size_t's, so no need to worry about overflow.
    }
    return compute(num_dim, num_obs, batches, output, options);
}
 
template<typename Dim_, typename Index_, typename Float_, typename Batch_>
Details compute(size_t num_dim, size_t num_obs, const Float_* input, const Batch_* batch, Float_* output, const Options<Dim_, Index_, Float_>& options) {
    const size_t nbatches = (num_obs ? static_cast<size_t>(*std::max_element(batch, batch + num_obs)) + 1 : 0);
    std::vector<size_t> sizes(nbatches);
    for (size_t o = 0; o < num_obs; ++o) {
        ++sizes[batch[o]];
    }
 
    // Avoiding the need to allocate a temporary buffer
    // if we're already dealing with contiguous batches.
    bool already_sorted = true;
    for (size_t o = 1; o < num_obs; ++o) {
       if (batch[o] < batch[o-1]) {
           already_sorted = false;
           break;
       }
    }
    if (already_sorted) {
        return compute(num_dim, sizes, input, output, options);
    }
 
    size_t accumulated = 0;
    std::vector<size_t> offsets(nbatches);
    for (size_t b = 0; b < nbatches; ++b) {
        offsets[b] = accumulated;
        accumulated += sizes[b];
    }
 
    // Dumping everything by order into another vector.
    std::vector<Float_> tmp(num_dim * num_obs);
    std::vector<const Float_*> ptrs(nbatches);
    for (size_t b = 0; b < nbatches; ++b) {
        ptrs[b] = tmp.data() + offsets[b] * num_dim;
    }
 
    for (size_t o = 0; o < num_obs; ++o) {
        auto current = input + o * num_dim;
        auto& offset = offsets[batch[o]];
        auto destination = tmp.data() + num_dim * offset; // already size_t's, so no need to cast to avoid overflow.
        std::copy_n(current, num_dim, destination);
        ++offset;
    }
 
    auto stats = internal::compute(num_dim, sizes, ptrs, output, options);
    internal::restore_order(num_dim, stats.merge_order, sizes, batch, output);
    return stats;
}
 
}
 
#endif