utils.hpp

Go to the documentation of this file.

#ifndef MNNCORRECT_UTILS_HPP
#define MNNCORRECT_UTILS_HPP
 
#include <vector>
#include <algorithm>
#include <memory>
#include <cstddef>
#include <type_traits>
 
#include "knncolle/knncolle.hpp"
 
#ifndef MNNCORRECT_CUSTOM_PARALLEL
#include "subpar/subpar.hpp"
#endif
 
namespace mnncorrect {
 
typedef std::size_t BatchIndex;
 
enum class MergePolicy : char { INPUT, SIZE, VARIANCE, RSS };
 
template<typename Task_, class Run_>
void parallelize(const int num_workers, const Task_ num_tasks, Run_ run_task_range) {
#ifndef MNNCORRECT_CUSTOM_PARALLEL
    // Methods could allocate or throw, so nothrow_ = false is safest.
    subpar::parallelize_range<false>(num_workers, num_tasks, std::move(run_task_range));
#else
    MNNCORRECT_CUSTOM_PARALLEL(num_workers, num_tasks, run_task_range);
#endif
}
 
template<typename Index_>
struct Batch {
    Index_ start = 0;
 
    Index_ size = 0;
};
 
template<typename Index_, typename Distance_>
using NeighborSet = std::vector<std::vector<std::pair<Index_, Distance_> > >;
 
template<typename Index_, typename Float_>
struct MetaBatch {
    // Each uncorrected metabatch has an original contiguous set of observations.
    // Once corrected, the metabatch ceases to exist as it becomes part of the destination metabatch.
    std::unique_ptr<knncolle::Prebuilt<Index_, Float_, Float_> > original_index;
    Batch<Index_> original_ids;
 
    // Corrected observations from other (meta)batches that have been redistributed into this meta batch.
    struct CorrectedBatch {
        CorrectedBatch() = default;
        CorrectedBatch(std::unique_ptr<knncolle::Prebuilt<Index_, Float_, Float_> > index, std::vector<Index_> ids) : index(std::move(index)), ids(std::move(ids)) {}
        std::unique_ptr<knncolle::Prebuilt<Index_, Float_, Float_> > index;
        std::vector<Index_> ids;
    };
    std::vector<CorrectedBatch> corrected;
};
 
template<typename Input_>
using I = std::remove_cv_t<std::remove_reference_t<Input_> >;
 
// Putting this here so that we can re-use it in the tests.
template<typename Index_, typename Float_, class Matrix_>
std::unique_ptr<knncolle::Prebuilt<Index_, Float_, Float_> > subset_and_index(
    const std::size_t num_dim,
    const std::vector<Index_>& subset,
    const Float_* const data,
    const knncolle::Builder<Index_, Float_, Float_, Matrix_>& builder,
    Float_* const buffer
) {
    const auto num_subset = subset.size();
    for (I<decltype(num_subset)> f = 0; f < num_subset; ++f) {
        const auto curdata = data + sanisizer::product_unsafe<std::size_t>(subset[f], num_dim);
        std::copy_n(curdata, num_dim, buffer + sanisizer::product_unsafe<std::size_t>(f, num_dim));
    }
    return builder.build_unique(knncolle::SimpleMatrix<Index_, Float_>(num_dim, num_subset, buffer));
}
}
 
#endif