topicks.hpp

Go to the documentation of this file.

#ifndef TOPICKER_TOPICKER_HPP
#define TOPICKER_TOPICKER_HPP
 
#include <vector>
#include <algorithm>
#include <numeric>
#include <cstddef>
 
#include "sanisizer/sanisizer.hpp"
 
namespace topicks {
 
template<typename Stat_>
struct PickTopGenesOptions {
    std::optional<Stat_> bound;
 
    bool keep_ties = true;
};
 
namespace internal {
 
template<bool keep_index_, typename Index_, typename Stat_, typename Top_, class Output_, class Cmp_, class CmpEqual_>
void pick_top_genes(Index_ n, const Stat_* statistic, Top_ top, Output_& output, Cmp_ cmp, CmpEqual_ cmpeq, const PickTopGenesOptions<Stat_>& options) {
    if (top == 0) {
        if constexpr(keep_index_) {
            ; // no-op, we assume it's already empty.
        } else {
            std::fill_n(output, n, false);
        }
        return;
    }
 
    if (sanisizer::is_greater_than(top, n)) {
        if (options.bound.has_value()) {
            auto bound = *(options.bound);
            for (Index_ i = 0; i < n; ++i) {
                bool ok = cmp(statistic[i], bound);
                if constexpr(keep_index_) {
                    if (ok) {
                        output.push_back(i);
                    }
                } else {
                    output[i] = ok;
                }
            }
        } else {
            if constexpr(keep_index_) {
                output.resize(sanisizer::cast<decltype(output.size())>(n));
                std::iota(output.begin(), output.end(), static_cast<Index_>(0));
            } else {
                std::fill_n(output, n, true);
            }
        }
        return;
    }
 
    auto semi_sorted = sanisizer::create<std::vector<Index_> >(n);
    std::iota(semi_sorted.begin(), semi_sorted.end(), static_cast<Index_>(0));
    auto cBegin = semi_sorted.begin(), cMid = cBegin + top - 1, cEnd = semi_sorted.end();
    std::nth_element(cBegin, cMid, cEnd, [&](Index_ l, Index_ r) -> bool { 
        auto L = statistic[l], R = statistic[r];
        if (L == R) {
            return l < r; // always favor the earlier index for a stable sort, even if larger = false.
        } else {
            return cmp(L, R);
        }
    });
    Stat_ threshold = statistic[*cMid];
 
    if (options.keep_ties) {
        if (options.bound.has_value() && !cmp(threshold, *(options.bound))) {
            auto bound = *(options.bound);
            for (Index_ i = 0; i < n; ++i) {
                bool ok = cmp(statistic[i], bound);
                if constexpr(keep_index_) {
                    if (ok) {
                        output.push_back(i);
                    }
                } else {
                    output[i] = ok;
                }
            }
        } else {
            for (Index_ i = 0; i < n; ++i) {
                bool ok = cmpeq(statistic[i], threshold);
                if constexpr(keep_index_) {
                    if (ok) {
                        output.push_back(i);
                    }
                } else {
                    output[i] = ok;
                }
            }
        }
        return;
    }
 
    if constexpr(keep_index_) {
        output.reserve(sanisizer::cast<decltype(output.size())>(top));
    } else {
        std::fill_n(output, n, false);
    }
 
    if (options.bound.has_value()) {
        auto bound = *(options.bound);
        Index_ counter = top; // cast is known safe as top <= n.
        while (counter > 0) {
            --counter;
            auto pos = semi_sorted[counter];
            if (cmp(statistic[pos], bound)) {
                if constexpr(keep_index_) {
                    output.push_back(pos);
                } else {
                    output[pos] = true;
                }
            }
        }
    } else {
        if constexpr(keep_index_) {
            output.insert(output.end(), semi_sorted.begin(), semi_sorted.begin() + top);
        } else {
            for (decltype(top) i = 0; i < top; ++i) {
                output[semi_sorted[i]] = true;
            }
        }
    }
 
    if constexpr(keep_index_) {
        std::sort(output.begin(), output.end());
    }
}
 
}
template<typename Stat_, typename Top_, typename Bool_>
void pick_top_genes(std::size_t n, const Stat_* statistic, Top_ top, bool larger, Bool_* output, const PickTopGenesOptions<Stat_>& options) {
    if (larger) {
        internal::pick_top_genes<false>(
            n, 
            statistic, 
            top,
            output, 
            [](Stat_ l, Stat_ r) -> bool { return l > r; },
            [](Stat_ l, Stat_ r) -> bool { return l >= r; },
            options
        );
    } else {
        internal::pick_top_genes<false>(
            n, 
            statistic, 
            top,
            output, 
            [](Stat_ l, Stat_ r) -> bool { return l < r; },
            [](Stat_ l, Stat_ r) -> bool { return l <= r; },
            options
        );
    }
}
 
template<typename Bool_, typename Stat_, typename Top_>
std::vector<Bool_> pick_top_genes(std::size_t n, const Stat_* statistic, Top_ top, bool larger, const PickTopGenesOptions<Stat_>& options) {
    auto output = sanisizer::create<std::vector<Bool_> >(n
#ifdef SCRAN_VARIANCES_TEST_INIT
        , SCRAN_VARIANCES_TEST_INIT
#endif
    );
    pick_top_genes(n, statistic, top, larger, output.data(), options);
    return output;
}
 
template<typename Index_, typename Stat_, typename Top_>
std::vector<Index_> pick_top_genes_index(Index_ n, const Stat_* statistic, Top_ top, bool larger, const PickTopGenesOptions<Stat_>& options) {
    std::vector<Index_> output;
    if (larger) {
        internal::pick_top_genes<true>(
            n, 
            statistic, 
            top,
            output,
            [](Stat_ l, Stat_ r) -> bool { return l > r; },
            [](Stat_ l, Stat_ r) -> bool { return l >= r; },
            options
        );
    } else {
        internal::pick_top_genes<true>(
            n, 
            statistic, 
            top,
            output,
            [](Stat_ l, Stat_ r) -> bool { return l < r; },
            [](Stat_ l, Stat_ r) -> bool { return l <= r; },
            options
        );
    }
    return output;
}
 
}
 
#endif