model_gene_variances.hpp

Go to the documentation of this file.

#ifndef SCRAN_MODEL_GENE_VARIANCES_H
#define SCRAN_MODEL_GENE_VARIANCES_H
 
#include <algorithm>
#include <vector>
#include <limits>
#include <cstddef>
 
#include "tatami/tatami.hpp"
#include "tatami_stats/tatami_stats.hpp"
#include "scran_blocks/scran_blocks.hpp"
#include "sanisizer/sanisizer.hpp"
 
#include "fit_variance_trend.hpp"
#include "utils.hpp"
 
namespace scran_variances {
 
enum class BlockAveragePolicy : unsigned char { MEAN, QUANTILE, NONE };
 
struct ModelGeneVariancesOptions {
    bool trend = true;
 
    FitVarianceTrendOptions fit_variance_trend_options;
 
    BlockAveragePolicy block_average_policy = BlockAveragePolicy::MEAN;
 
    scran_blocks::WeightPolicy block_weight_policy = scran_blocks::WeightPolicy::VARIABLE;
 
    scran_blocks::VariableWeightParameters variable_block_weight_parameters; 
 
    // Back-compatibility only.
    bool compute_average = true;
    double block_quantile = 0.5;
 
    int num_threads = 1;
};
 
template<typename Stat_>
struct ModelGeneVariancesBuffers {
    Stat_* means;
 
    Stat_* variances;
 
    Stat_* fitted;
 
    Stat_* residuals;
};
 
template<typename Stat_>
struct ModelGeneVariancesResults {
    ModelGeneVariancesResults() = default;
 
    ModelGeneVariancesResults(const std::size_t ngenes, const bool trend) :
        means(sanisizer::cast<I<decltype(means.size())> >(ngenes)
#ifdef SCRAN_VARIANCES_TEST_INIT
            , SCRAN_VARIANCES_TEST_INIT
#endif
        ),
        variances(sanisizer::cast<I<decltype(variances.size())> >(ngenes)
#ifdef SCRAN_VARIANCES_TEST_INIT
            , SCRAN_VARIANCES_TEST_INIT
#endif
        ),
        fitted(sanisizer::cast<I<decltype(fitted.size())> >(trend ? ngenes : 0)
#ifdef SCRAN_VARIANCES_TEST_INIT
            , SCRAN_VARIANCES_TEST_INIT
#endif
        ),
        residuals(sanisizer::cast<I<decltype(residuals.size())> >(trend ? ngenes : 0)
#ifdef SCRAN_VARIANCES_TEST_INIT
            , SCRAN_VARIANCES_TEST_INIT
#endif
        )
    {}
    std::vector<Stat_> means;
 
    std::vector<Stat_> variances;
 
    std::vector<Stat_> fitted;
 
    std::vector<Stat_> residuals;
};
 
template<typename Stat_>
struct ModelGeneVariancesBlockedBuffers {
    std::vector<ModelGeneVariancesBuffers<Stat_> > per_block;
 
    ModelGeneVariancesBuffers<Stat_> average;
};
 
template<typename Stat_>
struct ModelGeneVariancesBlockedResults {
    ModelGeneVariancesBlockedResults() = default;
 
    ModelGeneVariancesBlockedResults(const std::size_t ngenes, const std::size_t nblocks, const bool do_average, const bool do_trend) :
        average(do_average ? ngenes : 0, do_trend)
    {
        per_block.reserve(nblocks);
        for (I<decltype(nblocks)> b = 0; b < nblocks; ++b) {
            per_block.emplace_back(ngenes, do_trend);
        }
    }
    std::vector<ModelGeneVariancesResults<Stat_> > per_block;
 
    ModelGeneVariancesResults<Stat_> average;
};
 
namespace internal {
 
template<typename Value_, typename Index_, typename Stat_, typename Block_> 
void compute_variances_dense_row(
    const tatami::Matrix<Value_, Index_>& mat,
    const std::vector<ModelGeneVariancesBuffers<Stat_> >& buffers,
    const Block_* const block,
    const std::vector<Index_>& block_size,
    const int num_threads)
{
    const bool blocked = (block != NULL);
    const auto nblocks = block_size.size();
    const auto NR = mat.nrow(), NC = mat.ncol();
 
    tatami::parallelize([&](const int, const Index_ start, const Index_ length) -> void {
        auto tmp_means = sanisizer::create<std::vector<Stat_> >(blocked ? nblocks : 0);
        auto tmp_vars = sanisizer::create<std::vector<Stat_> >(blocked ? nblocks : 0);
 
        auto buffer = tatami::create_container_of_Index_size<std::vector<Value_> >(NC);
        auto ext = tatami::consecutive_extractor<false>(mat, true, start, length);
        for (Index_ r = start, end = start + length; r < end; ++r) {
            auto ptr = ext->fetch(buffer.data());
 
            if (blocked) {
                tatami_stats::grouped_variances::direct(
                    ptr,
                    NC,
                    block,
                    nblocks,
                    block_size.data(),
                    tmp_means.data(),
                    tmp_vars.data(),
                    false,
                    static_cast<Index_*>(NULL)
                );
                for (I<decltype(nblocks)> b = 0; b < nblocks; ++b) {
                    buffers[b].means[r] = tmp_means[b];
                    buffers[b].variances[r] = tmp_vars[b];
                }
            } else {
                const auto stat = tatami_stats::variances::direct(ptr, NC, false);
                buffers[0].means[r] = stat.first;
                buffers[0].variances[r] = stat.second;
            }
        }
    }, NR, num_threads);
}
 
template<typename Value_, typename Index_, typename Stat_, typename Block_> 
void compute_variances_sparse_row(
    const tatami::Matrix<Value_, Index_>& mat,
    const std::vector<ModelGeneVariancesBuffers<Stat_> >& buffers,
    const Block_* const block,
    const std::vector<Index_>& block_size,
    const int num_threads)
{
    const bool blocked = (block != NULL);
    const auto nblocks = block_size.size();
    const auto NR = mat.nrow(), NC = mat.ncol();
 
    tatami::parallelize([&](const int, const Index_ start, const Index_ length) -> void {
        auto tmp_means = sanisizer::create<std::vector<Stat_> >(nblocks);
        auto tmp_vars = sanisizer::create<std::vector<Stat_> >(nblocks);
        auto tmp_nzero = sanisizer::create<std::vector<Index_> >(nblocks);
 
        auto vbuffer = tatami::create_container_of_Index_size<std::vector<Value_> >(NC);
        auto ibuffer = tatami::create_container_of_Index_size<std::vector<Index_> >(NC);
        auto ext = tatami::consecutive_extractor<true>(mat, true, start, length, [&]{
            tatami::Options opt;
            opt.sparse_ordered_index = false;
            return opt;
        }());
 
        for (Index_ r = start, end = start + length; r < end; ++r) {
            auto range = ext->fetch(vbuffer.data(), ibuffer.data());
 
            if (blocked) {
                tatami_stats::grouped_variances::direct(
                    range.value,
                    range.index,
                    range.number,
                    block,
                    nblocks,
                    block_size.data(),
                    tmp_means.data(),
                    tmp_vars.data(),
                    tmp_nzero.data(),
                    false,
                    static_cast<Index_*>(NULL)
                );
                for (I<decltype(nblocks)> b = 0; b < nblocks; ++b) {
                    buffers[b].means[r] = tmp_means[b];
                    buffers[b].variances[r] = tmp_vars[b];
                }
            } else {
                const auto stat = tatami_stats::variances::direct(range.value, range.number, NC, false);
                buffers[0].means[r] = stat.first;
                buffers[0].variances[r] = stat.second;
            }
        }
    }, NR, num_threads);
}
 
template<typename Value_, typename Index_, typename Stat_, typename Block_> 
void compute_variances_dense_column(
    const tatami::Matrix<Value_, Index_>& mat,
    const std::vector<ModelGeneVariancesBuffers<Stat_> >& buffers,
    const Block_* const block,
    const std::vector<Index_>& block_size,
    const int num_threads)
{
    const bool blocked = (block != NULL);
    const auto nblocks = block_size.size();
    const auto NR = mat.nrow(), NC = mat.ncol();
 
    tatami::parallelize([&](const int thread, const Index_ start, const Index_ length) -> void {
        auto buffer = tatami::create_container_of_Index_size<std::vector<Value_> >(length);
        auto ext = tatami::consecutive_extractor<false>(mat, false, static_cast<Index_>(0), NC, start, length);
 
        auto get_var = [&](Index_ b) -> Stat_* { return buffers[b].variances; };
        tatami_stats::LocalOutputBuffers<Stat_, decltype(get_var)> local_vars(thread, nblocks, start, length, std::move(get_var));
        auto get_mean = [&](Index_ b) -> Stat_* { return buffers[b].means; };
        tatami_stats::LocalOutputBuffers<Stat_, decltype(get_mean)> local_means(thread, nblocks, start, length, std::move(get_mean));
 
        std::vector<tatami_stats::variances::RunningDense<Stat_, Value_, Index_> > runners;
        runners.reserve(nblocks);
        for (I<decltype(nblocks)> b = 0; b < nblocks; ++b) {
            runners.emplace_back(length, local_means.data(b), local_vars.data(b), false);
        }
 
        if (blocked) {
            for (I<decltype(NC)> c = 0; c < NC; ++c) {
                auto ptr = ext->fetch(buffer.data());
                runners[block[c]].add(ptr);
            }
        } else {
            for (I<decltype(NC)> c = 0; c < NC; ++c) {
                auto ptr = ext->fetch(buffer.data());
                runners[0].add(ptr);
            }
        }
 
        for (I<decltype(nblocks)> b = 0; b < nblocks; ++b) {
            runners[b].finish();
        }
        local_vars.transfer();
        local_means.transfer();
    }, NR, num_threads);
}
 
template<typename Value_, typename Index_, typename Stat_, typename Block_> 
void compute_variances_sparse_column(
    const tatami::Matrix<Value_, Index_>& mat,
    const std::vector<ModelGeneVariancesBuffers<Stat_> >& buffers,
    const Block_* const block,
    const std::vector<Index_>& block_size,
    const int num_threads) 
{
    const bool blocked = (block != NULL);
    const auto nblocks = block_size.size();
    const auto NR = mat.nrow(), NC = mat.ncol();
    auto nonzeros = sanisizer::create<std::vector<std::vector<Index_> > >(
        nblocks,
        tatami::create_container_of_Index_size<std::vector<Index_> >(NR)
    );
 
    tatami::parallelize([&](const int thread, const Index_ start, const Index_ length) -> void {
        auto vbuffer = tatami::create_container_of_Index_size<std::vector<Value_> >(length);
        auto ibuffer = tatami::create_container_of_Index_size<std::vector<Index_> >(length);
        auto ext = tatami::consecutive_extractor<true>(mat, false, static_cast<Index_>(0), NC, start, length, [&]{
            tatami::Options opt;
            opt.sparse_ordered_index = false;
            return opt;
        }());
 
        auto get_var = [&](Index_ b) -> Stat_* { return buffers[b].variances; };
        tatami_stats::LocalOutputBuffers<Stat_, decltype(get_var)> local_vars(thread, nblocks, start, length, std::move(get_var));
        auto get_mean = [&](Index_ b) -> Stat_* { return buffers[b].means; };
        tatami_stats::LocalOutputBuffers<Stat_, decltype(get_mean)> local_means(thread, nblocks, start, length, std::move(get_mean));
 
        std::vector<tatami_stats::variances::RunningSparse<Stat_, Value_, Index_> > runners;
        runners.reserve(nblocks);
        for (I<decltype(nblocks)> b = 0; b < nblocks; ++b) {
            runners.emplace_back(length, local_means.data(b), local_vars.data(b), false, start);
        }
 
        if (blocked) {
            for (I<decltype(NC)> c = 0; c < NC; ++c) {
                auto range = ext->fetch(vbuffer.data(), ibuffer.data());
                runners[block[c]].add(range.value, range.index, range.number);
            }
        } else {
            for (I<decltype(NC)> c = 0; c < NC; ++c) {
                auto range = ext->fetch(vbuffer.data(), ibuffer.data());
                runners[0].add(range.value, range.index, range.number);
            }
        }
 
        for (I<decltype(nblocks)> b = 0; b < nblocks; ++b) {
            runners[b].finish();
        }
        local_vars.transfer();
        local_means.transfer();
    }, NR, num_threads);
}
 
template<typename Value_, typename Index_, typename Stat_, typename Block_> 
void compute_variances(
    const tatami::Matrix<Value_, Index_>& mat,
    const std::vector<ModelGeneVariancesBuffers<Stat_> >& buffers,
    const Block_* const block,
    const std::vector<Index_>& block_size,
    const int num_threads) 
{
    if (mat.prefer_rows()) {
        if (mat.sparse()) {
            compute_variances_sparse_row(mat, buffers, block, block_size, num_threads);
        } else {
            compute_variances_dense_row(mat, buffers, block, block_size, num_threads);
        }
    } else {
        if (mat.sparse()) {
            compute_variances_sparse_column(mat, buffers, block, block_size, num_threads);
        } else {
            compute_variances_dense_column(mat, buffers, block, block_size, num_threads);
        }
    }
}
 
template<typename Stat_, typename Index_>
void extract_weights(
    const std::vector<Stat_>& block_weights,
    const std::vector<Index_>& block_size,
    const Index_ min_size,
    std::vector<Stat_>& tmp_weights
) {
    const auto nblocks = block_weights.size();
    tmp_weights.clear();
    for (I<decltype(nblocks)> b = 0; b < nblocks; ++b) {
        if (block_size[b] < min_size) { // skip blocks with insufficient cells.
            continue;
        }
        tmp_weights.push_back(block_weights[b]);
    }
}
 
template<typename Stat_, typename Index_, class Function_>
void extract_pointers(
    const std::vector<ModelGeneVariancesBuffers<Stat_> >& per_block, 
    const std::vector<Index_>& block_size,
    const Index_ min_size,
    const Function_ fun,
    std::vector<Stat_*>& tmp_pointers
) {
    const auto nblocks = per_block.size();
    tmp_pointers.clear();
    for (I<decltype(nblocks)> b = 0; b < nblocks; ++b) {
        if (block_size[b] < min_size) { // skip blocks with insufficient cells.
            continue;
        }
        tmp_pointers.push_back(fun(per_block[b]));
    }
}
 
}
template<typename Value_, typename Index_, typename Block_, typename Stat_>
void model_gene_variances_blocked(
    const tatami::Matrix<Value_, Index_>& mat, 
    const Block_* const block, 
    const ModelGeneVariancesBlockedBuffers<Stat_>& buffers,
    const ModelGeneVariancesOptions& options
) {
    const Index_ NR = mat.nrow(), NC = mat.ncol();
    std::vector<Index_> block_size;
 
    if (block) {
        block_size = tatami_stats::tabulate_groups(block, NC);
        internal::compute_variances(mat, buffers.per_block, block, block_size, options.num_threads);
    } else {
        block_size.push_back(NC); // everything is one big block.
        internal::compute_variances(mat, buffers.per_block, block, block_size, options.num_threads);
    }
    const auto nblocks = block_size.size();
 
    FitVarianceTrendWorkspace<Stat_> work;
    auto fopt = options.fit_variance_trend_options;
    fopt.num_threads = options.num_threads;
    bool all_trends_fitted = true;
 
    for (I<decltype(nblocks)> b = 0; b < nblocks; ++b) {
        const auto& current = buffers.per_block[b];
        if (current.fitted == NULL || current.residuals == NULL) {
            all_trends_fitted = false;
            continue;
        }
        if (block_size[b] >= 2) {
            fit_variance_trend(NR, current.means, current.variances, current.fitted, current.residuals, work, fopt);
        } else {
            std::fill_n(current.fitted, NR, std::numeric_limits<double>::quiet_NaN());
            std::fill_n(current.residuals, NR, std::numeric_limits<double>::quiet_NaN());
        }
    }
 
    const auto ave_means = buffers.average.means;
    const auto ave_variances = buffers.average.variances;
    const auto ave_fitted = buffers.average.fitted;
    const auto ave_residuals = buffers.average.residuals;
 
    if ((ave_fitted || ave_residuals) && !all_trends_fitted) {
        throw std::runtime_error("cannot compute average fitted values/residuals without per-block trend fits");
    }
 
    std::vector<Stat_*> tmp_pointers;
    tmp_pointers.reserve(nblocks);
 
    if (options.block_average_policy == BlockAveragePolicy::MEAN) {
        const auto block_weight = scran_blocks::compute_weights<Stat_>(block_size, options.block_weight_policy, options.variable_block_weight_parameters);
        std::vector<Stat_> tmp_weights;
        tmp_weights.reserve(nblocks);
 
        if (ave_means) {
            internal::extract_weights(block_weight, block_size, static_cast<Index_>(1), tmp_weights);
            internal::extract_pointers(buffers.per_block, block_size, static_cast<Index_>(1), [](const auto& x) -> Stat_* { return x.means; }, tmp_pointers);
            scran_blocks::parallel_weighted_means(NR, tmp_pointers, tmp_weights.data(), ave_means, /* skip_nan = */ false);
        }
 
        // Skip blocks without enough cells to compute the variance.
        internal::extract_weights(block_weight, block_size, static_cast<Index_>(2), tmp_weights);
 
        if (ave_variances) {
            internal::extract_pointers(buffers.per_block, block_size, static_cast<Index_>(2), [](const auto& x) -> Stat_* { return x.variances; }, tmp_pointers);
            scran_blocks::parallel_weighted_means(NR, tmp_pointers, tmp_weights.data(), ave_variances, /* skip_nan = */ false);
        }
 
        if (ave_fitted) {
            internal::extract_pointers(buffers.per_block, block_size, static_cast<Index_>(2), [](const auto& x) -> Stat_* { return x.fitted; }, tmp_pointers);
            scran_blocks::parallel_weighted_means(NR, tmp_pointers, tmp_weights.data(), ave_fitted, /* skip_nan = */ false);
        }
 
        if (ave_residuals) {
            internal::extract_pointers(buffers.per_block, block_size, static_cast<Index_>(2), [](const auto& x) -> Stat_* { return x.residuals; }, tmp_pointers);
            scran_blocks::parallel_weighted_means(NR, tmp_pointers, tmp_weights.data(), ave_residuals, /* skip_nan = */ false);
        }
 
    } else if (options.block_average_policy == BlockAveragePolicy::QUANTILE) {
        if (ave_means) {
            internal::extract_pointers(buffers.per_block, block_size, static_cast<Index_>(1), [](const auto& x) -> Stat_* { return x.means; }, tmp_pointers);
            scran_blocks::parallel_quantiles(NR, tmp_pointers, options.block_quantile, ave_means, /* skip_nan = */ false);
        }
 
        // Skip blocks without enough cells to compute the variance.
 
        if (ave_variances) {
            internal::extract_pointers(buffers.per_block, block_size, static_cast<Index_>(2), [](const auto& x) -> Stat_* { return x.variances; }, tmp_pointers);
            scran_blocks::parallel_quantiles(NR, tmp_pointers, options.block_quantile, ave_variances, /* skip_nan = */ false);
        }
 
        if (ave_fitted) {
            internal::extract_pointers(buffers.per_block, block_size, static_cast<Index_>(2), [](const auto& x) -> Stat_* { return x.fitted; }, tmp_pointers);
            scran_blocks::parallel_quantiles(NR, tmp_pointers, options.block_quantile, ave_fitted, /* skip_nan = */ false);
        }
 
        if (ave_residuals) {
            internal::extract_pointers(buffers.per_block, block_size, static_cast<Index_>(2), [](const auto& x) -> Stat_* { return x.residuals; }, tmp_pointers);
            scran_blocks::parallel_quantiles(NR, tmp_pointers, options.block_quantile, ave_residuals, /* skip_nan = */ false);
        }
    }
}
 
template<typename Value_, typename Index_, typename Stat_> 
void model_gene_variances(
    const tatami::Matrix<Value_, Index_>& mat, 
    ModelGeneVariancesBuffers<Stat_> buffers, // yes, the lack of a const ref here is deliberate, we need to move it into bbuffers anyway.
    const ModelGeneVariancesOptions& options)
{
    ModelGeneVariancesBlockedBuffers<Stat_> bbuffers;
    bbuffers.per_block.emplace_back(std::move(buffers));
 
    bbuffers.average.means = NULL;
    bbuffers.average.variances = NULL;
    bbuffers.average.fitted = NULL;
    bbuffers.average.residuals = NULL;
 
    model_gene_variances_blocked(mat, static_cast<Index_*>(NULL), bbuffers, options);
}
 
template<typename Stat_ = double, typename Value_, typename Index_>
ModelGeneVariancesResults<Stat_> model_gene_variances(const tatami::Matrix<Value_, Index_>& mat, const ModelGeneVariancesOptions& options) {
    ModelGeneVariancesResults<Stat_> output(mat.nrow(), options.trend); // cast is safe, as any tatami Index_ can always fit into a size_t.
 
    ModelGeneVariancesBuffers<Stat_> buffers;
    buffers.means = output.means.data();
    buffers.variances = output.variances.data();
 
    if (options.trend) {
        buffers.fitted = output.fitted.data();
        buffers.residuals = output.residuals.data();
    } else {
        buffers.fitted = NULL;
        buffers.residuals = NULL;
    }
 
    model_gene_variances(mat, std::move(buffers), options);
    return output;
}
 
template<typename Stat_ = double, typename Value_, typename Index_, typename Block_>
ModelGeneVariancesBlockedResults<Stat_> model_gene_variances_blocked(const tatami::Matrix<Value_, Index_>& mat, const Block_* const block, const ModelGeneVariancesOptions& options) {
    const auto nblocks = (block ? tatami_stats::total_groups(block, mat.ncol()) : 1);
 
    const bool do_average = options.compute_average /* for back-compatibility */ && options.block_average_policy != BlockAveragePolicy::NONE;
    ModelGeneVariancesBlockedResults<Stat_> output(
        mat.nrow(), // cast is safe, any tatami Index_ can always fit into a size_t.
        nblocks,
        do_average,
        options.trend
    );
 
    ModelGeneVariancesBlockedBuffers<Stat_> buffers;
    sanisizer::resize(buffers.per_block, nblocks);
    for (I<decltype(nblocks)> b = 0; b < nblocks; ++b) {
        auto& current = buffers.per_block[b];
        current.means = output.per_block[b].means.data();
        current.variances = output.per_block[b].variances.data();
 
        if (options.trend) {
            current.fitted = output.per_block[b].fitted.data();
            current.residuals = output.per_block[b].residuals.data();
        } else {
            current.fitted = NULL;
            current.residuals = NULL;
        }
    }
 
    if (!do_average) {
        buffers.average.means = NULL;
        buffers.average.variances = NULL;
        buffers.average.fitted = NULL;
        buffers.average.residuals = NULL;
    } else {
        buffers.average.means = output.average.means.data();
        buffers.average.variances = output.average.variances.data();
 
        if (options.trend) {
            buffers.average.fitted = output.average.fitted.data();
            buffers.average.residuals = output.average.residuals.data();
        } else {
            buffers.average.fitted = NULL;
            buffers.average.residuals = NULL;
        }
    }
 
    model_gene_variances_blocked(mat, block, buffers, options);
    return output;
}
 
}
 
#endif