sums.hpp

Go to the documentation of this file.

#ifndef TATAMI_STATS__SUMS_HPP
#define TATAMI_STATS__SUMS_HPP
 
#include "tatami/tatami.hpp"
#include "utils.hpp"
 
#include <vector>
#include <numeric>
#include <algorithm>
 
namespace tatami_stats {
 
namespace sums {
 
struct Options {
    bool skip_nan = false;
 
    int num_threads = 1;
};
 
template<typename Output_ = double, typename Value_, typename Index_>
Output_ direct(const Value_* ptr, Index_ num, bool skip_nan) {
    return internal::nanable_ifelse_with_value<Value_>(
        skip_nan,
        [&]() -> Output_ {
            Output_ sum = 0;
            for (Index_ i = 0; i < num; ++i) {
                auto val = ptr[i];
                if (!std::isnan(val)) {
                    sum += val;
                }
            }
            return sum;
        },
        [&]() -> Output_ {
            return std::accumulate(ptr, ptr + num, static_cast<Output_>(0));
        }
    );
}
 
template<typename Output_, typename Value_, typename Index_>
class RunningDense {
public:
    RunningDense(Index_ num, Output_* sum, bool skip_nan) : my_num(num), my_sum(sum), my_skip_nan(skip_nan) {}
 
    void add(const Value_* ptr) {
        internal::nanable_ifelse<Value_>(
            my_skip_nan,
            [&]() {
                for (Index_ i = 0; i < my_num; ++i) {
                    auto val = ptr[i];
                    if (!std::isnan(val)) {
                        my_sum[i] += val;
                    }
                }
            },
            [&]() {
                for (Index_ i = 0; i < my_num; ++i) {
                    my_sum[i] += ptr[i];
                }
            }
        );
    }
 
private:
    Index_ my_num;
    Output_* my_sum;
    bool my_skip_nan;
};
 
template<typename Output_, typename Value_, typename Index_>
class RunningSparse {
public:
    RunningSparse(Output_* sum, bool skip_nan, Index_ subtract = 0) : 
        my_sum(sum), my_skip_nan(skip_nan), my_subtract(subtract) {}
 
    void add(const Value_* value, const Index_* index, Index_ number) {
        internal::nanable_ifelse<Value_>(
            my_skip_nan,
            [&]() {
                for (Index_ i = 0; i < number; ++i) {
                    auto val = value[i];
                    if (!std::isnan(val)) {
                        my_sum[index[i] - my_subtract] += val;
                    }
                }
            },
            [&]() {
                for (Index_ i = 0; i < number; ++i) {
                    my_sum[index[i] - my_subtract] += value[i];
                }
            }
        );
    }
 
private:
    Output_* my_sum;
    bool my_skip_nan;
    Index_ my_subtract;
};
 
template<typename Value_, typename Index_, typename Output_>
void apply(bool row, const tatami::Matrix<Value_, Index_>* p, Output_* output, const Options& sopt) {
    auto dim = (row ? p->nrow() : p->ncol());
    auto otherdim = (row ? p->ncol() : p->nrow());
    const bool direct = p->prefer_rows() == row;
 
    if (p->sparse()) {
        if (direct) {
            tatami::Options opt;
            opt.sparse_extract_index = false;
 
            tatami::parallelize([&](int, Index_ s, Index_ l) {
                auto ext = tatami::consecutive_extractor<true>(p, row, s, l, opt);
                std::vector<Value_> vbuffer(otherdim);
                for (Index_ x = 0; x < l; ++x) {
                    auto out = ext->fetch(vbuffer.data(), NULL);
                    output[x + s] = sums::direct(out.value, out.number, sopt.skip_nan);
                }
            }, dim, sopt.num_threads);
 
        } else {
            tatami::Options opt;
            opt.sparse_ordered_index = false;
 
            tatami::parallelize([&](int thread, Index_ s, Index_ l) {
                auto ext = tatami::consecutive_extractor<true>(p, !row, static_cast<Index_>(0), otherdim, s, l, opt);
                std::vector<Value_> vbuffer(l);
                std::vector<Index_> ibuffer(l);
 
                LocalOutputBuffer<Output_> local_output(thread, s, l, output);
                sums::RunningSparse<Output_, Value_, Index_> runner(local_output.data(), sopt.skip_nan, s);
 
                for (Index_ x = 0; x < otherdim; ++x) {
                    auto out = ext->fetch(vbuffer.data(), ibuffer.data());
                    runner.add(out.value, out.index, out.number);
                }
 
                local_output.transfer();
            }, dim, sopt.num_threads);
        }
 
    } else {
        if (direct) {
            tatami::parallelize([&](int, Index_ s, Index_ l) {
                auto ext = tatami::consecutive_extractor<false>(p, row, s, l);
                std::vector<Value_> buffer(otherdim);
                for (Index_ x = 0; x < l; ++x) {
                    auto out = ext->fetch(buffer.data());
                    output[x + s] = sums::direct(out, otherdim, sopt.skip_nan);
                }
            }, dim, sopt.num_threads);
 
        } else {
            tatami::parallelize([&](int thread, Index_ s, Index_ l) {
                auto ext = tatami::consecutive_extractor<false>(p, !row, static_cast<Index_>(0), otherdim, s, l);
                std::vector<Value_> buffer(l);
 
                LocalOutputBuffer<Output_> local_output(thread, s, l, output);
                sums::RunningDense<Output_, Value_, Index_> runner(l, local_output.data(), sopt.skip_nan);
 
                for (Index_ x = 0; x < otherdim; ++x) {
                    auto out = ext->fetch(buffer.data());
                    runner.add(out);
                }
 
                local_output.transfer();
            }, dim, sopt.num_threads);
        }
    }
 
    return;
}
 
template<typename Output_ = double, typename Value_, typename Index_>
std::vector<Output_> by_column(const tatami::Matrix<Value_, Index_>* p, const Options& sopt) {
    std::vector<Output_> output(p->ncol());
    apply(false, p, output.data(), sopt);
    return output;
}
 
template<typename Output_ = double, typename Value_, typename Index_>
std::vector<Output_> by_column(const tatami::Matrix<Value_, Index_>* p) {
    return by_column(p, Options());
}
 
template<typename Output_ = double, typename Value_, typename Index_>
std::vector<Output_> by_row(const tatami::Matrix<Value_, Index_>* p, const Options& sopt) {
    std::vector<Output_> output(p->nrow());
    apply(true, p, output.data(), sopt);
    return output;
}
 
template<typename Output_ = double, typename Value_, typename Index_>
std::vector<Output_> by_row(const tatami::Matrix<Value_, Index_>* p) {
    return by_row(p, Options());
}
 
}
 
}
 
#endif