model_gene_variances.hpp

Go to the documentation of this file.

#ifndef SCRAN_MODEL_GENE_VARIANCES_H
#define SCRAN_MODEL_GENE_VARIANCES_H
 
#include "tatami/tatami.hpp"
#include "tatami_stats/tatami_stats.hpp"
#include "scran_blocks/scran_blocks.hpp"
 
#include "fit_variance_trend.hpp"
 
#include <algorithm>
#include <vector>
#include <limits>
 
namespace scran_variances {
 
struct ModelGeneVariancesOptions {
    FitVarianceTrendOptions fit_variance_trend_options;
 
    scran_blocks::WeightPolicy block_weight_policy = scran_blocks::WeightPolicy::VARIABLE;
 
    scran_blocks::VariableWeightParameters variable_block_weight_parameters; 
 
    bool compute_average = true;
 
    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(size_t ngenes) : means(ngenes), variances(ngenes), fitted(ngenes), residuals(ngenes) {}
    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(size_t ngenes, size_t nblocks, bool compute_average) : average(compute_average ? ngenes : 0) {
        per_block.reserve(nblocks);
        for (size_t b = 0; b < nblocks; ++b) {
            per_block.emplace_back(ngenes);
        }
    }
    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_* block,
    const std::vector<Index_>& block_size,
    int num_threads)
{
    bool blocked = (block != NULL);
    auto nblocks = block_size.size();
    auto NR = mat.nrow(), NC = mat.ncol();
 
    tatami::parallelize([&](int, Index_ start, Index_ length) -> void {
        std::vector<Stat_> tmp_means(blocked ? nblocks : 0);
        std::vector<Stat_> tmp_vars(blocked ? nblocks : 0);
 
        std::vector<Value_> buffer(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 (size_t b = 0; b < nblocks; ++b) {
                    buffers[b].means[r] = tmp_means[b];
                    buffers[b].variances[r] = tmp_vars[b];
                }
            } else {
                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_* block,
    const std::vector<Index_>& block_size,
    int num_threads)
{
    bool blocked = (block != NULL);
    auto nblocks = block_size.size();
    auto NR = mat.nrow(), NC = mat.ncol();
 
    tatami::parallelize([&](int, Index_ start, Index_ length) -> void {
        std::vector<Stat_> tmp_means(nblocks);
        std::vector<Stat_> tmp_vars(nblocks);
        std::vector<Index_> tmp_nzero(nblocks);
 
        std::vector<Value_> vbuffer(NC);
        std::vector<Index_> ibuffer(NC);
        tatami::Options opt;
        opt.sparse_ordered_index = false;
        auto ext = tatami::consecutive_extractor<true>(&mat, true, start, length, 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 (size_t b = 0; b < nblocks; ++b) {
                    buffers[b].means[r] = tmp_means[b];
                    buffers[b].variances[r] = tmp_vars[b];
                }
            } else {
                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_* block,
    const std::vector<Index_>& block_size,
    int num_threads)
{
    bool blocked = (block != NULL);
    auto nblocks = block_size.size();
    auto NR = mat.nrow(), NC = mat.ncol();
 
    tatami::parallelize([&](int thread, Index_ start, Index_ length) -> void {
        std::vector<Value_> buffer(length);
        auto ext = tatami::consecutive_extractor<false>(&mat, false, 0, NC, start, length);
 
        std::vector<tatami_stats::LocalOutputBuffer<Stat_> > local_var_output;
        local_var_output.reserve(nblocks);
        std::vector<tatami_stats::LocalOutputBuffer<Stat_> > local_mean_output;
        local_mean_output.reserve(nblocks);
        std::vector<tatami_stats::variances::RunningDense<Stat_, Value_, Index_> > runners;
        runners.reserve(nblocks);
 
        for (size_t b = 0; b < nblocks; ++b) {
            local_var_output.emplace_back(thread, start, length, buffers[b].variances);
            local_mean_output.emplace_back(thread, start, length, buffers[b].means);
            runners.emplace_back(length, local_mean_output.back().data(), local_var_output.back().data(), false);
        }
 
        if (blocked) {
            for (Index_ c = 0; c < NC; ++c) {
                auto ptr = ext->fetch(buffer.data());
                runners[block[c]].add(ptr);
            }
        } else {
            for (Index_ c = 0; c < NC; ++c) {
                auto ptr = ext->fetch(buffer.data());
                runners[0].add(ptr);
            }
        }
 
        for (size_t b = 0; b < nblocks; ++b) {
            runners[b].finish();
            local_var_output[b].transfer();
            local_mean_output[b].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_* block,
    const std::vector<Index_>& block_size,
    int num_threads) 
{
    bool blocked = (block != NULL);
    auto nblocks = block_size.size();
    auto NR = mat.nrow(), NC = mat.ncol();
    std::vector<std::vector<Index_> > nonzeros(nblocks, std::vector<Index_>(NR));
 
    tatami::parallelize([&](int thread, Index_ start, Index_ length) -> void {
        std::vector<Value_> vbuffer(length);
        std::vector<Index_> ibuffer(length);
        tatami::Options opt;
        opt.sparse_ordered_index = false;
        auto ext = tatami::consecutive_extractor<true>(&mat, false, 0, NC, start, length, opt);
 
        std::vector<tatami_stats::LocalOutputBuffer<Stat_> > local_var_output;
        local_var_output.reserve(nblocks);
        std::vector<tatami_stats::LocalOutputBuffer<Stat_> > local_mean_output;
        local_mean_output.reserve(nblocks);
        std::vector<tatami_stats::variances::RunningSparse<Stat_, Value_, Index_> > runners;
        runners.reserve(nblocks);
 
        for (size_t b = 0; b < nblocks; ++b) {
            local_var_output.emplace_back(thread, start, length, buffers[b].variances);
            local_mean_output.emplace_back(thread, start, length, buffers[b].means);
            runners.emplace_back(length, local_mean_output.back().data(), local_var_output.back().data(), false, start);
        }
 
        if (blocked) {
            for (Index_ 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 (Index_ c = 0; c < NC; ++c) {
                auto range = ext->fetch(vbuffer.data(), ibuffer.data());
                runners[0].add(range.value, range.index, range.number);
            }
        }
 
        for (size_t b = 0; b < nblocks; ++b) {
            runners[b].finish();
            local_var_output[b].transfer();
            local_mean_output[b].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_* block,
    const std::vector<Index_>& block_size,
    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 Index_, typename Stat_, class Function_>
void compute_average(
    Index_ ngenes, 
    const std::vector<ModelGeneVariancesBuffers<Stat_> >& per_block, 
    const std::vector<Index_>& block_size,
    const std::vector<Stat_>& block_weights,
    Index_ min_size,
    Function_ fun,
    std::vector<Stat_*>& tmp_pointers,
    std::vector<Stat_>& tmp_weights,
    Stat_* output) 
{
    if (!output) {
        return;
    }
 
    tmp_pointers.clear();
    tmp_weights.clear();
    for (size_t b = 0, nblocks = per_block.size(); b < nblocks; ++b) {
        if (block_size[b] < min_size) { // skip blocks with insufficient cells.
            continue;
        }
        tmp_weights.push_back(block_weights[b]);
        tmp_pointers.push_back(fun(per_block[b]));
    }
 
    scran_blocks::average_vectors_weighted(ngenes, tmp_pointers, tmp_weights.data(), output, /* skip_nan = */ false);
}
 
}
template<typename Value_, typename Index_, typename Block_, typename Stat_>
void model_gene_variances_blocked(
    const tatami::Matrix<Value_, Index_>& mat, 
    const Block_* block, 
    const ModelGeneVariancesBlockedBuffers<Stat_>& buffers,
    const ModelGeneVariancesOptions& options)
{
    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);
    }
    size_t nblocks = block_size.size();
 
    FitVarianceTrendWorkspace<Stat_> work;
    auto fopt = options.fit_variance_trend_options;
    fopt.num_threads = options.num_threads;
    for (size_t b = 0; b < nblocks; ++b) {
        const auto& current = buffers.per_block[b];
        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());
        }
    }
 
    auto ave_means = buffers.average.means;
    auto ave_variances = buffers.average.variances;
    auto ave_fitted = buffers.average.fitted;
    auto ave_residuals = buffers.average.residuals;
 
    if (ave_means || ave_variances || ave_fitted || ave_residuals) {
        auto block_weight = scran_blocks::compute_weights<Stat_>(block_size, options.block_weight_policy, options.variable_block_weight_parameters);
 
        std::vector<Stat_*> tmp_pointers;
        std::vector<Stat_> tmp_weights;
        tmp_pointers.reserve(nblocks);
        tmp_weights.reserve(nblocks);
 
        internal::compute_average(NR, buffers.per_block, block_size, block_weight,
            /* min_size = */ 1,  // skip blocks without enough cells to compute the mean.
            [](const auto& x) -> Stat_* { return x.means; }, tmp_pointers, tmp_weights, ave_means);
 
        internal::compute_average(NR, buffers.per_block, block_size, block_weight,
            /* min_size = */ 2, // skip blocks without enough cells to compute the variance.
            [](const auto& x) -> Stat_* { return x.variances; }, tmp_pointers, tmp_weights, ave_variances);
 
        internal::compute_average(NR, buffers.per_block, block_size, block_weight, 
            /* min_size = */ 2, // ditto.
            [](const auto& x) -> Stat_* { return x.fitted; }, tmp_pointers, tmp_weights, ave_fitted);
 
        internal::compute_average(NR, buffers.per_block, block_size, block_weight, 
            /* min_size = */ 2, // ditto.
            [](const auto& x) -> Stat_* { return x.residuals; }, tmp_pointers, tmp_weights, ave_residuals);
    }
}
 
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());
 
    ModelGeneVariancesBuffers<Stat_> buffers;
    buffers.means = output.means.data();
    buffers.variances = output.variances.data();
    buffers.fitted = output.fitted.data();
    buffers.residuals = output.residuals.data();
 
    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_* block, const ModelGeneVariancesOptions& options) {
    size_t nblocks = (block ? tatami_stats::total_groups(block, mat.ncol()) : 1);
    ModelGeneVariancesBlockedResults<Stat_> output(mat.nrow(), nblocks, options.compute_average);
 
    ModelGeneVariancesBlockedBuffers<Stat_> buffers;
    buffers.per_block.resize(nblocks);
    for (size_t 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();
        current.fitted = output.per_block[b].fitted.data();
        current.residuals = output.per_block[b].residuals.data();
    }
 
    if (!options.compute_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();
        buffers.average.fitted = output.average.fitted.data();
        buffers.average.residuals = output.average.residuals.data();
    }
 
    model_gene_variances_blocked(mat, block, buffers, options);
    return output;
}
 
}
 
#endif