sum.hpp

Go to the documentation of this file.

#ifndef TATAMI_STATS_SUM_HPP
#define TATAMI_STATS_SUM_HPP
 
#include "utils.hpp"
 
#include <vector>
#include <numeric>
#include <algorithm>
#include <optional>
#include <cmath>
 
#include "tatami/tatami.hpp"
 
namespace tatami_stats {
 
struct SumOptions {
    bool skip_nan = false;
 
    int num_threads = 1;
};
 
template<typename Value_, typename Index_, typename Output_>
void sum_direct(bool row, const tatami::Matrix<Value_, Index_>& mat, Output_* output, const SumOptions& opt) {
    const auto dim = (row ? mat.nrow() : mat.ncol());
    const auto otherdim = (row ? mat.ncol() : mat.nrow());
 
    if (mat.is_sparse()) {
        tatami::parallelize([&](int, Index_ s, Index_ l) -> void {
            tatami::Options topt;
            topt.sparse_extract_index = false;
            auto ext = tatami::consecutive_extractor<true>(mat, row, s, l, topt);
            auto vbuffer = tatami::create_container_of_Index_size<std::vector<Value_> >(otherdim);
 
            nanable_ifelse<Value_>(
                opt.skip_nan,
                [&]() -> void {
                    for (Index_ x = 0; x < l; ++x) {
                        const auto out = ext->fetch(vbuffer.data(), NULL);
                        Output_ sum = 0;
                        for (Index_ i = 0; i < out.number; ++i) {
                            const auto val = out.value[i];
                            if (!std::isnan(val)) {
                                sum += val;
                            }
                        }
                        output[x + s] = sum;
                    }
                },
                [&]() -> void {
                    quickstats::PairwiseSumWorkspace<Output_> work;
                    for (Index_ x = 0; x < l; ++x) {
                        const auto out = ext->fetch(vbuffer.data(), NULL);
                        output[x + s] = quickstats::pairwise_sum(out.number, out.value, work); // Index_ -> size_t conversion is safe, as per tatami's contract.
                    }
                }
            );
        }, dim, opt.num_threads);
 
    } else {
        tatami::parallelize([&](int, Index_ s, Index_ l) -> void {
            auto ext = tatami::consecutive_extractor<false>(mat, row, s, l);
            auto buffer = tatami::create_container_of_Index_size<std::vector<Value_> >(otherdim);
 
            nanable_ifelse<Value_>(
                opt.skip_nan,
                [&]() -> void {
                    for (Index_ x = 0; x < l; ++x) {
                        const auto ptr = ext->fetch(buffer.data());
                        Output_ sum = 0;
                        for (Index_ i = 0; i < otherdim; ++i) {
                            const auto val = ptr[i];
                            if (!std::isnan(val)) {
                                sum += val;
                            }
                        }
                        output[x + s] = sum;
                    }
                },
                [&]() -> void {
                    quickstats::PairwiseSumWorkspace<Output_> work;
                    for (Index_ x = 0; x < l; ++x) {
                        const auto ptr = ext->fetch(buffer.data());
                        output[x + s] = quickstats::pairwise_sum(otherdim, ptr, work); // Index_ -> size_t conversion is safe, as per tatami's contract.
                    }
                }
            );
        }, dim, opt.num_threads);
    }
}
 
template<typename Value_, typename Index_, typename Output_>
void sum_running(bool row, const tatami::Matrix<Value_, Index_>& mat, Output_* output, const SumOptions& opt) {
    const auto dim = (row ? mat.nrow() : mat.ncol());
    const auto otherdim = (row ? mat.ncol() : mat.nrow());
 
    const bool do_parallel = (opt.num_threads > 1);
    std::optional<std::vector<std::optional<std::vector<Output_> > > > all_partial_sum;
    if (do_parallel) {
        all_partial_sum.emplace(sanisizer::cast<I<decltype(all_partial_sum->size())> >(opt.num_threads - 1));
    }
 
    std::fill_n(output, dim, 0);
 
    const auto nused = tatami::parallelize([&](int thread, Index_ s, Index_ l) -> void {
        Output_* sum_ptr;
        std::optional<std::vector<Output_> > cur_sum;
        if (!do_parallel) {
            sum_ptr = output;
        } else {
            if (thread == 0) {
                sum_ptr = output;
            } else {
                cur_sum.emplace(tatami::cast_Index_to_container_size<std::vector<Output_> >(dim));
                sum_ptr = cur_sum->data();
            }
        }
 
        if (mat.is_sparse()) {
            tatami::Options topt;
            topt.sparse_ordered_index = false; // ordering doesn't matter.
            auto ext = tatami::consecutive_extractor<true>(mat, !row, s, l, topt);
            auto vbuffer = tatami::create_container_of_Index_size<std::vector<Value_> >(dim);
            auto ibuffer = tatami::create_container_of_Index_size<std::vector<Index_> >(dim);
 
            for (Index_ x = 0; x < l; ++x) {
                const auto out = ext->fetch(vbuffer.data(), ibuffer.data());
                nanable_ifelse<Value_>(
                    opt.skip_nan,
                    [&]() -> void {
                        for (Index_ i = 0; i < out.number; ++i) {
                            const auto val = out.value[i];
                            sum_ptr[out.index[i]] += (std::isnan(val) ? 0 : val); // slightly nicer for vectorization.
                        }
                    },
                    [&]() -> void {
                        for (Index_ i = 0; i < out.number; ++i) {
                            sum_ptr[out.index[i]] += out.value[i];
                        }
                    }
                );
            }
 
        } else {
            auto ext = tatami::consecutive_extractor<false>(mat, !row, s, l);
            auto buffer = tatami::create_container_of_Index_size<std::vector<Value_> >(dim);
 
            for (Index_ x = 0; x < l; ++x) {
                const auto ptr = ext->fetch(buffer.data());
                nanable_ifelse<Value_>(
                    opt.skip_nan,
                    [&]() -> void {
                        for (Index_ i = 0; i < dim; ++i) {
                            const auto val = ptr[i];
                            sum_ptr[i] += (std::isnan(val) ? 0 : val); // slightly nicer for vectorization.
                        }
                    },
                    [&]() -> void {
                        for (Index_ i = 0; i < dim; ++i) {
                            sum_ptr[i] += ptr[i];
                        }
                    }
                );
            }
        }
 
        if (do_parallel) {
            if (thread > 0) {
                (*all_partial_sum)[thread - 1] = std::move(cur_sum);
            }
        }
    }, otherdim, opt.num_threads);
 
    if (do_parallel) {
        for (int u = 1; u < nused; ++u) {
            const auto& cur_sum = *((*all_partial_sum)[u - 1]);
            for (Index_ d = 0; d < dim; ++d) {
                output[d] += cur_sum[d];
            }
        }
    }
}
template<typename Value_, typename Index_, typename Output_>
void sum(bool row, const tatami::Matrix<Value_, Index_>& mat, Output_* output, const SumOptions& opt) {
    if (mat.prefer_rows() == row) {
        sum_direct(row, mat, output, opt);
    } else {
        sum_running(row, mat, output, opt);
    }
}
 
template<typename Output_ = double, typename Value_, typename Index_>
std::vector<Output_> sum(bool row, const tatami::Matrix<Value_, Index_>& mat, const SumOptions& opt) {
    const auto dim = (row ? mat.nrow() : mat.ncol());
    auto output = sanisizer::create<std::vector<Output_> >(dim
#ifdef TATAMI_STATS_TEST_DIRTY
        , -1
#endif
    );
    sum(row, mat, output.data(), opt);
    return output;
}
 
}
 
#endif