group_variance.hpp

Go to the documentation of this file.

#ifndef TATAMI_STATS_GROUP_VARIANCE_HPP
#define TATAMI_STATS_GROUP_VARIANCE_HPP
 
#include <vector>
#include <algorithm>
#include <cstddef>
#include <optional>
#include <cassert>
#include <limits>
#include <cmath>
 
#include "tatami/tatami.hpp"
#include "sanisizer/sanisizer.hpp"
 
#include "group_rss.hpp"
#include "skip_nan/group_rss.hpp"
#include "utils.hpp"
 
namespace tatami_stats {
 
template<typename Output_ = double>
struct GroupVarianceOptions {
    bool skip_nan = false;
 
    int num_threads = 1;
 
    Output_ mean_placeholder = quickstats::nan_if_available_else_zero<Output_>();
 
    Output_ variance_placeholder = quickstats::nan_if_available_else_zero<Output_>();
};
 
template<typename Output_>
struct GroupVarianceBuffers {
    std::vector<Output_*> mean;
 
    std::vector<Output_*> variance;
};
 
template<typename Value_, typename Index_, typename Group_, typename Count_, typename Output_>
void group_variance(
    bool row,
    const tatami::Matrix<Value_, Index_>& mat,
    const Group_* const group,
    const Group_ num_groups,
    const Count_* const group_size,
    const GroupVarianceBuffers<Output_>& output,
    const GroupVarianceOptions<Output_>& opt
) {
    assert(sanisizer::is_equal(num_groups, output.mean.size()));
    assert(sanisizer::is_equal(num_groups, output.variance.size()));
    const auto dim = (row ? mat.nrow() : mat.ncol());
 
    nanable_ifelse<Value_>(
        opt.skip_nan,
 
        [&]() -> void {
            skip_nan::GroupRssBuffers<Output_, Index_> tmp;
            tmp.mean = output.mean;
            tmp.rss = output.variance;
 
            auto count = sanisizer::create<std::vector<std::vector<Index_> > >(num_groups);
            tmp.count.reserve(num_groups);
            for (Group_ g = 0; g < num_groups; ++g) {
                tatami::resize_container_to_Index_size(count[g], dim);
                tmp.count.push_back(count[g].data());
            }
 
            skip_nan::GroupRssOptions ropt;
            ropt.num_threads = opt.num_threads;
            skip_nan::group_rss(row, mat, group, num_groups, tmp, ropt);
            for (Group_ g = 0; g < num_groups; ++g) {
                const auto outvar = output.variance[g];
                const auto curcounts = count[g];
                for (Index_ d = 0; d < dim; ++d) {
                    if (curcounts[d] <= 1) {
                        outvar[d] = opt.variance_placeholder;
                    } else {
                        outvar[d] /= curcounts[d] - 1;
                    }
                }
            }
        },
 
        [&]() -> void {
            GroupRssBuffers<Output_> tmp;
            tmp.mean = output.mean;
            tmp.rss = output.variance;
 
            GroupRssOptions ropt;
            ropt.num_threads = opt.num_threads;
            group_rss(row, mat, group, num_groups, group_size, tmp, ropt);
 
            for (Group_ g = 0; g < num_groups; ++g) {
                const auto outvar = output.variance[g];
                const auto gsize = group_size[g];
                if (gsize <= 1) {
                    std::fill_n(outvar, dim, opt.variance_placeholder);
                } else {
                    for (Index_ d = 0; d < dim; ++d) {
                        outvar[d] /= gsize - 1;
                    }
                }
            }
        }
    );
}
 
template<typename Value_, typename Index_, typename Group_, typename Output_>
void group_variance(
    bool row,
    const tatami::Matrix<Value_, Index_>& mat,
    const Group_* const group,
    const Group_ num_groups,
    const GroupVarianceBuffers<Output_>& output,
    const GroupVarianceOptions<Output_>& opt
) {
    auto group_size = sanisizer::create<std::vector<Index_> >(num_groups);
    const auto otherdim = (row ? mat.ncol() : mat.nrow());
    for (Index_ o = 0; o < otherdim; ++o) {
        group_size[group[o]] += 1;
    }
    group_variance(row, mat, group, num_groups, group_size.data(), output, opt);
}
 
template<typename Output_>
struct GroupVarianceResult {
    std::vector<std::vector<Output_> > mean;
 
    std::vector<std::vector<Output_> > variance;
};
 
template<typename Output_ = double, typename Value_, typename Index_, typename Group_> 
GroupVarianceResult<Output_> group_variance(
    bool row,
    const tatami::Matrix<Value_, Index_>& mat,
    const Group_* const group,
    const Group_ num_groups,
    const GroupVarianceOptions<Output_>& opt
) {
    GroupVarianceResult<Output_> output;
    sanisizer::resize(output.mean, num_groups);
    sanisizer::resize(output.variance, num_groups);
 
    GroupVarianceBuffers<Output_> buffers;
    sanisizer::resize(buffers.mean, num_groups);
    sanisizer::resize(buffers.variance, num_groups);
    const auto dim = (row ? mat.nrow() : mat.ncol());
 
    for (Group_ g = 0; g < num_groups; ++g) {
        tatami::resize_container_to_Index_size(output.mean[g], dim
#ifdef TATAMI_STATS_TEST_DIRTY
            , -1
#endif
        );
        buffers.mean[g] = output.mean[g].data();
        tatami::resize_container_to_Index_size(output.variance[g], dim
#ifdef TATAMI_STATS_TEST_DIRTY
            , -1
#endif
        );
        buffers.variance[g] = output.variance[g].data();
    }
 
    group_variance(row, mat, group, num_groups, buffers, opt);
    return output;
}
 
}
 
#endif