group_sum.hpp

Go to the documentation of this file.

#ifndef TATAMI_STATS_GROUPED_SUMS_HPP
#define TATAMI_STATS_GROUPED_SUMS_HPP
 
#include "utils.hpp"
#include "sum.hpp"
 
#include <vector>
#include <algorithm>
#include <cstddef>
 
#include "tatami/tatami.hpp"
#include "sanisizer/sanisizer.hpp"
#include "jiwoo/jiwoo.hpp"
 
namespace tatami_stats {
 
struct GroupSumOptions {
    bool skip_nan = false;
 
    int num_threads = 1;
};
 
template<typename Value_, typename Index_, typename Group_, typename Output_>
void group_sum_direct(
    bool row,
    const tatami::Matrix<Value_, Index_>& mat,
    const Group_* group,
    const Group_ num_groups,
    const std::vector<Output_*>& output,
    const GroupSumOptions& opt
) {
    const Index_ dim = (row ? mat.nrow() : mat.ncol());
    const Index_ otherdim = (row ? mat.ncol() : mat.nrow());
 
    if (mat.sparse()) {
        tatami::parallelize([&](int, Index_ start, Index_ len) -> void {
            auto ext = tatami::consecutive_extractor<true>(mat, row, start, len);
            auto xbuffer = tatami::create_container_of_Index_size<std::vector<Value_> >(otherdim);
            auto ibuffer = tatami::create_container_of_Index_size<std::vector<Index_> >(otherdim);
            auto tmp = sanisizer::create<std::vector<Output_> >(num_groups);
 
            for (Index_ x = 0; x < len; ++x) {
                auto range = ext->fetch(xbuffer.data(), ibuffer.data());
                std::fill(tmp.begin(), tmp.end(), static_cast<Output_>(0));
 
                nanable_ifelse<Value_>(
                    opt.skip_nan,
                    [&]() -> void {
                        for (Index_ j = 0; j < range.number; ++j) {
                            const auto val = range.value[j];
                            if (!std::isnan(val)) {
                                tmp[group[range.index[j]]] += val;
                            }
                        }
                    },
                    [&]() -> void {
                        for (Index_ j = 0; j < range.number; ++j) {
                            tmp[group[range.index[j]]] += range.value[j];
                        }
                    }
                );
 
                for (I<decltype(num_groups)> g = 0; g < num_groups; ++g) {
                    output[g][start + x] = tmp[g];
                }
            }
        }, dim, opt.num_threads);
 
    } else {
        tatami::parallelize([&](int, Index_ start, Index_ len) -> void {
            auto ext = tatami::consecutive_extractor<false>(mat, row, start, len);
            auto xbuffer = tatami::create_container_of_Index_size<std::vector<Value_> >(otherdim);
            auto tmp = sanisizer::create<std::vector<Output_> >(num_groups);
 
            for (Index_ x = 0; x < len; ++x) {
                auto ptr = ext->fetch(xbuffer.data());
                std::fill(tmp.begin(), tmp.end(), static_cast<Output_>(0));
 
                nanable_ifelse<Value_>(
                    opt.skip_nan,
                    [&]() -> void {
                        for (Index_ j = 0; j < otherdim; ++j) {
                            const auto val = ptr[j];
                            if (!std::isnan(val)) {
                                tmp[group[j]] += val;
                            }
                        }
                    },
                    [&]() -> void {
                        for (Index_ j = 0; j < otherdim; ++j) {
                            tmp[group[j]] += ptr[j];
                        }
                    }
                );
 
                for (I<decltype(num_groups)> g = 0; g < num_groups; ++g) {
                    output[g][start + x] = tmp[g];
                }
            }
        }, dim, opt.num_threads);
    }
}
 
template<typename Value_, typename Index_, typename Group_, typename Output_>
void group_sum_running(
    bool row,
    const tatami::Matrix<Value_, Index_>& mat,
    const Group_* group,
    const Group_ num_groups,
    const std::vector<Output_*>& output,
    const GroupSumOptions& opt
) {
    const Index_ dim = (row ? mat.nrow() : mat.ncol());
    const Index_ otherdim = (row ? mat.ncol() : mat.nrow());
    const bool is_sparse = mat.is_sparse();
 
    const auto do_parallel = opt.num_threads > 1;
    std::optional<std::vector<std::optional<jiwoo::EquilengthArrays<Output_> > > > all_partial_sums;
    if (do_parallel) {
        all_partial_sums.emplace(sanisizer::cast<I<decltype(all_partial_sums->size())> >(opt.num_threads - 1));
    }
 
    for (Group_ g = 0; g < num_groups; ++g) {
        std::fill_n(output[g], dim, 0);
    }
 
    const auto nused = tatami::parallelize([&](int thread, Index_ start, Index_ len) -> void {
        // If we can, directly dump the sum to the output pointers, otherwise put it into a temporary.
        std::optional<jiwoo::EquilengthArrays<Output_> > cur_sums;
 
        Output_* const * sum_ptrs;
        if (!do_parallel) {
            sum_ptrs = output.data();
        } else {
            if (thread == 0) {
                sum_ptrs = output.data();
            } else {
                cur_sums.emplace(
                    sanisizer::Cast(num_groups),
                    static_cast<std::size_t>(dim), // cast is safe due to the tatami contract.
                    0
                ); 
                sum_ptrs = cur_sums->get();
            }
        }
 
        if (is_sparse) {
            // 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 topt;
            topt.sparse_ordered_index = false; 
            auto ext = tatami::consecutive_extractor<true>(mat, !row, start, len, topt);
            auto xbuffer = 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 < len; ++x) {
                auto range = ext->fetch(xbuffer.data(), ibuffer.data());
                const auto sum_ptr = sum_ptrs[group[start + x]];
 
                nanable_ifelse<Value_>(
                    opt.skip_nan,
                    [&]() -> void {
                        for (Index_ i = 0; i < range.number; ++i) {
                            const auto val = range.value[i];
                            sum_ptr[range.index[i]] += (std::isnan(val) ? 0 : val); // slightly nicer for vectorization.
                        }
                    },
                    [&]() -> void {
                        for (Index_ i = 0; i < range.number; ++i) {
                            sum_ptr[range.index[i]] += range.value[i];
                        }
                    }
                );
            }
 
        } else {
            auto ext = tatami::consecutive_extractor<false>(mat, !row, start, len);
            auto buffer = tatami::create_container_of_Index_size<std::vector<Value_> >(dim);
 
            for (Index_ x = 0; x < len; ++x) {
                auto ptr = ext->fetch(buffer.data());
                const auto sum_ptr = sum_ptrs[group[start + x]];
 
                nanable_ifelse<Value_>(
                    opt.skip_nan,
                    [&]() -> void {
                        for (Index_ d = 0; d < dim; ++d) {
                            const auto val = ptr[d];
                            sum_ptr[d] += (std::isnan(val) ? 0 : val); // slightly nicer for vectorization.
                        }
                    },
                    [&]() -> void {
                        for (Index_ d = 0; d < dim; ++d) {
                            sum_ptr[d] += ptr[d];
                        }
                    }
                );
            }
        }
 
        if (do_parallel) {
            if (thread > 0) {
                (*all_partial_sums)[thread - 1] = std::move(cur_sums);
            }
        }
    }, otherdim, opt.num_threads);
 
    if (do_parallel) {
        for (Group_ g = 0; g < num_groups; ++g) {
            const auto cur_out = output[g];
            for (int u = 1; u < nused; ++u) {
                const auto cur_sum = (*((*all_partial_sums)[u - 1]))[g];
                for (Index_ d = 0; d < dim; ++d) {
                    cur_out[d] += cur_sum[d];
                }
            }
        }
    }
}
template<typename Value_, typename Index_, typename Group_, typename Output_>
void group_sum(
    bool row,
    const tatami::Matrix<Value_, Index_>& mat,
    const Group_* group,
    const Group_ num_groups,
    const std::vector<Output_*>& output,
    const GroupSumOptions& opt
) {
    if (mat.prefer_rows() == row) {
        group_sum_direct(row, mat, group, num_groups, output, opt);
    } else {
        group_sum_running(row, mat, group, num_groups, output, opt);
    }
}
 
template<typename Output_ = double, typename Value_, typename Index_, typename Group_>
std::vector<std::vector<Output_> > group_sum(
    bool row,
    const tatami::Matrix<Value_, Index_>& mat,
    const Group_* group,
    const Group_ num_groups,
    const GroupSumOptions& opt
) {
    auto output = sanisizer::create<std::vector<std::vector<Output_> > >(num_groups);
    auto ptrs = sanisizer::create<std::vector<Output_*> >(num_groups);
    const Index_ dim = (row ? mat.nrow() : mat.ncol());
    for (Group_ g = 0; g < num_groups; ++g) {
        tatami::resize_container_to_Index_size(output[g], dim
#ifdef TATAMI_STATS_TEST_DIRTY
            , -1
#endif
        );
        ptrs[g] = output[g].data();
    }
    group_sum(row, mat, group, num_groups, ptrs, opt);
    return output;
}
 
}
 
#endif