grouped_sums.hpp

Go to the documentation of this file.

#ifndef TATAMI_STATS_GROUPED_SUMS_HPP
#define TATAMI_STATS_GROUPED_SUMS_HPP
 
#include "utils.hpp"
#include "tatami/tatami.hpp"
#include "sums.hpp"
#include <vector>
#include <algorithm>
 
namespace tatami_stats {
 
namespace grouped_sums {
 
struct Options {
    bool skip_nan = false;
 
    int num_threads = 1;
};
 
template<typename Value_, typename Index_, typename Group_, typename Output_>
void apply(bool row, const tatami::Matrix<Value_, Index_>* p, const Group_* group, size_t num_groups, Output_** output, const Options& sopt) {
    Index_ dim = (row ? p->nrow() : p->ncol());
    Index_ otherdim = (row ? p->ncol() : p->nrow());
 
    if (p->sparse()) {
        if (p->prefer_rows() == row) {
            tatami::parallelize([&](int, Index_ start, Index_ len) -> void {
                auto ext = tatami::consecutive_extractor<true>(p, row, start, len);
                std::vector<Value_> xbuffer(otherdim);
                std::vector<Index_> ibuffer(otherdim);
                std::vector<Output_> tmp(num_groups);
 
                for (Index_ i = 0; i < len; ++i) {
                    auto range = ext->fetch(xbuffer.data(), ibuffer.data());
                    std::fill(tmp.begin(), tmp.end(), static_cast<Output_>(0));
 
                    internal::nanable_ifelse<Value_>(
                        sopt.skip_nan,
                        [&]() {
                            for (int j = 0; j < range.number; ++j) {
                                auto val = range.value[j];
                                if (!std::isnan(val)) {
                                    tmp[group[range.index[j]]] += val;
                                }
                            }
                        },
                        [&]() {
                            for (int j = 0; j < range.number; ++j) {
                                tmp[group[range.index[j]]] += range.value[j];
                            }
                        }
                    );
 
                    for (size_t g = 0; g < num_groups; ++g) {
                        output[g][i + start] = tmp[g];
                    }
                }
            }, dim, sopt.num_threads);
 
        } else {
            // Order within each observed vector doesn't affect numerical
            // precision of the outcome, as addition order for each objective
            // vector is already well-defined for a running calculation.
            tatami::Options opt;
            opt.sparse_ordered_index = false; 
 
            tatami::parallelize([&](int thread, Index_ start, Index_ len) -> void {
                std::vector<sums::RunningSparse<Output_, Value_, Index_> > runners;
                runners.reserve(num_groups);
                std::vector<LocalOutputBuffer<Output_> > local_output;
                local_output.reserve(num_groups);
 
                for (size_t g = 0; g < num_groups; ++g) {
                    local_output.emplace_back(thread, start, len, output[g]);
                    runners.emplace_back(local_output.back().data(), sopt.skip_nan, start);
                }
 
                auto ext = tatami::consecutive_extractor<true>(p, !row, static_cast<Index_>(0), otherdim, start, len, opt);
                std::vector<Value_> xbuffer(len);
                std::vector<Index_> ibuffer(len);
 
                for (int i = 0; i < otherdim; ++i) {
                    auto range = ext->fetch(xbuffer.data(), ibuffer.data());
                    runners[group[i]].add(range.value, range.index, range.number);
                }
 
                for (size_t g = 0; g < num_groups; ++g) {
                    local_output[g].transfer();
                }
            }, dim, sopt.num_threads);
        }
 
    } else {
        if (p->prefer_rows() == row) {
            tatami::parallelize([&](int, Index_ start, Index_ len) -> void {
                auto ext = tatami::consecutive_extractor<false>(p, row, start, len);
                std::vector<Value_> xbuffer(otherdim);
                std::vector<Output_> tmp(num_groups);
 
                for (Index_ i = 0; i < len; ++i) {
                    auto ptr = ext->fetch(xbuffer.data());
                    std::fill(tmp.begin(), tmp.end(), static_cast<Output_>(0));
 
                    internal::nanable_ifelse<Value_>(
                        sopt.skip_nan,
                        [&]() {
                            for (Index_ j = 0; j < otherdim; ++j) {
                                auto val = ptr[j];
                                if (!std::isnan(val)) {
                                    tmp[group[j]] += val;
                                }
                            }
                        },
                        [&]() {
                            for (Index_ j = 0; j < otherdim; ++j) {
                                tmp[group[j]] += ptr[j];
                            }
                        }
                    );
 
                    for (size_t g = 0; g < num_groups; ++g) {
                        output[g][i + start] = tmp[g];
                    }
                }
            }, dim, sopt.num_threads);
 
        } else {
            tatami::parallelize([&](int thread, Index_ start, Index_ len) -> void {
                std::vector<sums::RunningDense<Output_, Value_, Index_> > runners;
                runners.reserve(num_groups);
                std::vector<LocalOutputBuffer<Output_> > local_output;
                local_output.reserve(num_groups);
 
                for (size_t g = 0; g < num_groups; ++g) {
                    local_output.emplace_back(thread, start, len, output[g]);
                    runners.emplace_back(len, local_output.back().data(), sopt.skip_nan);
                }
 
                std::vector<Value_> xbuffer(len);
                auto ext = tatami::consecutive_extractor<false>(p, !row, static_cast<Index_>(0), otherdim, start, len);
 
                for (int i = 0; i < otherdim; ++i) {
                    auto ptr = ext->fetch(xbuffer.data());
                    runners[group[i]].add(ptr);
                }
 
                for (size_t g = 0; g < num_groups; ++g) {
                    local_output[g].transfer();
                }
            }, dim, sopt.num_threads);
        }
    }
}
 
template<typename Output_ = double, typename Value_, typename Index_, typename Group_>
std::vector<std::vector<Output_> > by_row(const tatami::Matrix<Value_, Index_>* p, const Group_* group, const Options& sopt) {
    size_t mydim = p->nrow();
    auto ngroup = total_groups(group, p->ncol());
 
    std::vector<std::vector<Output_> > output(ngroup);
    std::vector<Output_*> ptrs;
    ptrs.reserve(output.size());
    for (auto& o : output) {
        o.resize(mydim);
        ptrs.push_back(o.data());
    }
 
    apply(true, p, group, ngroup, ptrs.data(), sopt);
    return output;
}
 
template<typename Output_ = double, typename Value_, typename Index_, typename Group_>
std::vector<std::vector<Output_> > by_row(const tatami::Matrix<Value_, Index_>* p, const Group_* group) {
    return by_row(p, group, Options());
}
 
template<typename Output_ = double, typename Value_, typename Index_, typename Group_>
std::vector<std::vector<Output_> > by_column(const tatami::Matrix<Value_, Index_>* p, const Group_* group, const Options& sopt) {
    size_t mydim = p->ncol();
    auto ngroup = total_groups(group, p->nrow());
 
    std::vector<std::vector<Output_> > output(ngroup);
    std::vector<Output_*> ptrs;
    ptrs.reserve(output.size());
    for (auto& o : output) {
        o.resize(mydim);
        ptrs.push_back(o.data());
    }
 
    apply(false, p, group, ngroup, ptrs.data(), sopt);
    return output;
}
 
template<typename Output_ = double, typename Value_, typename Index_, typename Group_>
std::vector<std::vector<Output_> > by_column(const tatami::Matrix<Value_, Index_>* p, const Group_* group) {
    return by_column(p, group, Options());
}
 
}
 
}
 
#endif