rss.hpp

Go to the documentation of this file.

#ifndef TATAMI_STATS_SKIP_NAN_RSS_HPP
#define TATAMI_STATS_SKIP_NAN_RSS_HPP
 
#include "../utils.hpp"
 
#include <vector>
#include <cmath>
#include <numeric>
#include <limits>
#include <algorithm>
#include <cstddef>
#include <optional>
 
#include "tatami/tatami.hpp"
#include "sanisizer/sanisizer.hpp"
#include "quickstats/quickstats.hpp"
 
namespace tatami_stats {
 
namespace skip_nan {
 
struct RssOptions {
    int num_threads = 1;
};
 
template<typename Output_, typename Count_>
struct RssBuffers {
    Output_* mean;
 
    Output_* rss;
 
    Count_* count;
};
 
template<typename Value_, typename Index_, typename Output_, typename Count_>
void rss_direct(bool row, const tatami::Matrix<Value_, Index_>& mat, RssBuffers<Output_, Count_>& output, const RssOptions& opt) {
    const auto dim = (row ? mat.nrow() : mat.ncol());
    const auto otherdim = (row ? mat.ncol() : mat.nrow());
 
    if (mat.sparse()) {
        tatami::Options topt;
        topt.sparse_extract_index = false;
 
        tatami::parallelize([&](int, Index_ s, Index_ l) -> void {
            auto ext = tatami::consecutive_extractor<true>(mat, row, s, l, topt);
            auto vbuffer = tatami::create_container_of_Index_size<std::vector<Value_> >(otherdim);
            quickstats::RssWorkspace<Output_> work;
            for (Index_ x = 0; x < l; ++x) {
                auto out = ext->fetch(vbuffer.data(), NULL);
                tatami::copy_n(out.value, out.number, vbuffer.data());
                const auto new_number = shift_nans(vbuffer.data(), out.number);
                const Index_ new_total = otherdim - (out.number - new_number);
                const auto res = quickstats::rss(new_total, new_number, vbuffer.data(), work);
                output.mean[x + s] = res.mean;
                output.rss[x + s] = res.rss;
                output.count[x + s] = new_total;
            }
        }, 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);
            quickstats::RssWorkspace<Output_> work;
            for (Index_ x = 0; x < l; ++x) {
                auto out = ext->fetch(buffer.data());
                tatami::copy_n(out, otherdim, buffer.data());
                const auto new_total = shift_nans(buffer.data(), otherdim);
                const auto res = quickstats::rss(new_total, buffer.data(), work);
                output.mean[x + s] = res.mean;
                output.rss[x + s] = res.rss;
                output.count[x + s] = new_total;
            }
        }, dim, opt.num_threads);
    }
}
 
template<typename Value_, typename Index_, typename Output_, typename Count_>
void rss_running(bool row, const tatami::Matrix<Value_, Index_>& mat, RssBuffers<Output_, Count_>& output, const RssOptions& opt) {
    const auto dim = (row ? mat.nrow() : mat.ncol());
    const auto otherdim = (row ? mat.ncol() : mat.nrow());
    const bool is_sparse = mat.is_sparse();
 
    if (otherdim == 0) {
        std::fill_n(output.mean, dim, std::numeric_limits<Output_>::quiet_NaN());
        std::fill_n(output.rss, dim, 0);
        std::fill_n(output.count, dim, 0);
        return;
    }
 
    assert(opt.num_threads > 0);
    const bool do_parallel = opt.num_threads > 1;
    std::optional<std::vector<std::optional<std::vector<Output_> > > > all_partial_mean, all_partial_rss;
    std::optional<std::vector<std::optional<std::vector<Count_> > > > all_partial_count;
    if (do_parallel) {
        // -1, as we'll repurpose the output buffers to store the output of the first thread.
        all_partial_rss.emplace(sanisizer::cast<I<decltype(all_partial_rss->size())> >(opt.num_threads - 1));
        all_partial_mean.emplace(sanisizer::cast<I<decltype(all_partial_mean->size())> >(opt.num_threads));
        all_partial_count.emplace(sanisizer::cast<I<decltype(all_partial_mean->size())> >(opt.num_threads));
    }
 
    std::fill_n(output.rss, dim, 0);
    std::fill_n(output.mean, dim, 0);
    std::fill_n(output.count, dim, 0);
 
    const int nused = tatami::parallelize([&](int thread, Index_ s, Index_ l) -> void {
        Output_* rss_ptr;
        Output_* mean_ptr;
        Count_* count_ptr;
        std::optional<std::vector<Output_> > cur_rss, cur_mean;
        std::optional<std::vector<Count_> > cur_count;
 
        if (!do_parallel) {
            // Storing mean and RSS directly in the output vector to cut down two allocations if we're not working in parallel.
            rss_ptr = output.rss;
            mean_ptr = output.mean;
            count_ptr = output.count;
        } else {
            // Storing the partial RSS directly in the output vector to save ourselves an allocation if we're in the first thread.
            // We can't do the same for the mean, though, as we need to keep the partial mean and the global mean separate for the reduction.
            cur_mean.emplace(tatami::cast_Index_to_container_size<std::vector<Output_> >(dim));
            mean_ptr = cur_mean->data();
            cur_count.emplace(tatami::cast_Index_to_container_size<std::vector<Count_> >(dim));
            count_ptr = cur_count->data();
            if (thread == 0) {
                rss_ptr = output.rss;
            } else {
                cur_rss.emplace(tatami::cast_Index_to_container_size<std::vector<Output_> >(dim));
                rss_ptr = cur_rss->data();
            }
        }
 
        if (is_sparse) {
            tatami::Options topt;
            topt.sparse_ordered_index = false;
            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);
            auto nonzeros = tatami::create_container_of_Index_size<std::vector<Count_> >(dim);
 
            quickstats::RssRunningSparseSkip<Count_, Value_, Output_> runner(dim, mean_ptr, rss_ptr, nonzeros.data(), count_ptr);
            for (Index_ x = 0; x < l; ++x) {
                auto out = ext->fetch(vbuffer.data(), ibuffer.data());
                runner.add(
                    out.number,
                    out.value,
                    out.index,
                    [](std::size_t, const Value_ val) -> bool {
                        return std::isnan(val);
                    }
                );
            }
            runner.finish();
 
        } else {
            auto ext = tatami::consecutive_extractor<false>(mat, !row, s, l);
            auto buffer = tatami::create_container_of_Index_size<std::vector<Value_> >(dim);
 
            quickstats::RssRunningDenseSkip<Count_, Value_, Output_> runner(dim, mean_ptr, rss_ptr, count_ptr);
            for (Index_ x = 0; x < l; ++x) {
                auto out = ext->fetch(buffer.data());
                runner.add(
                    out,
                    [](std::size_t, const Value_ val) -> bool {
                        return std::isnan(val);
                    }
                );
            }
            runner.finish();
        }
 
        // Moving results to the main containers.
        if (do_parallel) {
            (*all_partial_mean)[thread] = std::move(cur_mean);
            (*all_partial_count)[thread] = std::move(cur_count);
            if (thread > 0) {
                (*all_partial_rss)[thread - 1] = std::move(cur_rss);
            }
        }
    }, otherdim, opt.num_threads);
    assert(nused > 0);
 
    // Don't check nused > 1, as it's possible for do_parallel = true with nused = 1 if not all threads are used.
    // This would cause us to leave output.mean and output.rss empty.
    if (do_parallel) {
        const auto& ap_mean = *all_partial_mean;
        const auto& ap_rss = *all_partial_rss;
        const auto& ap_count = *all_partial_count;
 
        // Computing the global total.
        for (int u = 0; u < nused; ++u) {
            const auto& cur_count = *(ap_count[u]);
            for (Index_ d = 0; d < dim; ++d) {
                output.count[d] += cur_count[d];
            }
        }
 
        // Computing the global mean from its components.
        for (int u = 0; u < nused; ++u) {
            const auto& cur_count = *(ap_count[u]);
            const auto& cur_mean = *(ap_mean[u]);
            for (Index_ d = 0; d < dim; ++d) {
                if (cur_count[d] > 0) { // protect against NaN means at a count of 0.
                    const auto mult = static_cast<Output_>(cur_count[d]) / static_cast<Output_>(output.count[d]);
                    output.mean[d] += cur_mean[d] * mult;
                }
            }
        }
 
        for (Index_ d = 0; d < dim; ++ d) {
            if (output.count[d] == 0) { 
                output.mean[d] = std::numeric_limits<Output_>::quiet_NaN();
            }
        }
 
        // Combining the RSS. This time, we need to use the safe version as we don't know whether all elements were skipped in a thread.
        for (int u = 0; u < nused; ++u) {
            const auto& cur_count = *(ap_count[u]);
            const auto& cur_mean = *(ap_mean[u]);
            if (u == 0) {
                for (Index_ d = 0; d < dim; ++d) {
                    output.rss[d] = quickstats::recenter_rss(cur_count[d], output.rss[d], cur_mean[d], output.mean[d]); 
                }
            } else {
                const auto& cur_rss = *(ap_rss[u - 1]);
                for (Index_ d = 0; d < dim; ++d) {
                    output.rss[d] += quickstats::recenter_rss(cur_count[d], cur_rss[d], cur_mean[d], output.mean[d]); 
                }
            }
        }
    }
}
template<typename Value_, typename Index_, typename Output_, typename Count_>
void rss(bool row, const tatami::Matrix<Value_, Index_>& mat, RssBuffers<Output_, Count_>& output, const RssOptions& opt) {
    if (mat.prefer_rows() == row) {
        rss_direct(row, mat, output, opt);
    } else {
        rss_running(row, mat, output, opt);
    }
}
 
template<typename Output_, typename Count_>
struct RssResult {
    std::vector<Output_> mean;
 
    std::vector<Output_> rss;
 
    std::vector<Count_> count;
};
 
template<typename Output_ = double, typename Count_, typename Value_, typename Index_>
RssResult<Output_, Count_> rss(bool row, const tatami::Matrix<Value_, Index_>& mat, const RssOptions& opt) {
    RssResult<Output_, Count_> output;
    const auto dim = (row ? mat.nrow() : mat.ncol());
 
    tatami::resize_container_to_Index_size(output.mean, dim
#ifdef TATAMI_STATS_TEST_DIRTY
        , -1
#endif
    );
    tatami::resize_container_to_Index_size(output.rss, dim
#ifdef TATAMI_STATS_TEST_DIRTY
        , -1
#endif
    );
    tatami::resize_container_to_Index_size(output.count, dim
#ifdef TATAMI_STATS_TEST_DIRTY
        , -1
#endif
    );
 
    RssBuffers<Output_, Count_> buffers;
    buffers.mean = output.mean.data();
    buffers.rss = output.rss.data();
    buffers.count = output.count.data();
 
    rss(row, mat, buffers, opt);
    return output;
}
 
}
 
}
 
#endif