model_gene_variances.hpp

Go to the documentation of this file.

#ifndef SCRAN_MODEL_GENE_VARIANCES_HPP
#define SCRAN_MODEL_GENE_VARIANCES_HPP
 
#include <algorithm>
#include <vector>
#include <limits>
#include <cstddef>
#include <cassert>
#include <optional>
 
#include "tatami/tatami.hpp"
#include "tatami_stats/tatami_stats.hpp"
#include "quickstats/quickstats.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_* mean;
 
    Stat_* variance;
 
    Stat_* fitted;
 
    Stat_* residual;
};
 
template<typename Value_, typename Index_, typename Stat_> 
void model_gene_variances(
    const tatami::Matrix<Value_, Index_>& mat, 
    const ModelGeneVariancesBuffers<Stat_> buffers,
    const ModelGeneVariancesOptions& options
) {
    tatami_stats::VarianceBuffers<Stat_> vbuf;
    vbuf.mean = buffers.mean;    
    vbuf.variance = buffers.variance;    
 
    tatami_stats::VarianceOptions vopt;
    vopt.num_threads = options.num_threads;
    tatami_stats::variance(true, mat, vbuf, vopt);
 
    FitVarianceTrendWorkspace<Stat_> work;
    auto fopt = options.fit_variance_trend_options;
    fopt.num_threads = options.num_threads;
 
    if (buffers.fitted != NULL && buffers.residual != NULL) {
        const auto NR = mat.nrow();
        if (mat.ncol() >= 2) {
            fit_variance_trend(NR, buffers.mean, buffers.variance, buffers.fitted, buffers.residual, work, fopt);
        } else {
            std::fill_n(buffers.fitted, NR, std::numeric_limits<double>::quiet_NaN());
            std::fill_n(buffers.residual, NR, std::numeric_limits<double>::quiet_NaN());
        }
    }
}
 
template<typename Stat_>
struct ModelGeneVariancesResults {
    ModelGeneVariancesResults() = default;
 
    ModelGeneVariancesResults(const std::size_t ngenes, const bool trend) :
        mean(sanisizer::cast<I<decltype(mean.size())> >(ngenes)
#ifdef SCRAN_VARIANCES_TEST_INIT
            , SCRAN_VARIANCES_TEST_INIT
#endif
        ),
        variance(sanisizer::cast<I<decltype(variance.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
        ),
        residual(sanisizer::cast<I<decltype(residual.size())> >(trend ? ngenes : 0)
#ifdef SCRAN_VARIANCES_TEST_INIT
            , SCRAN_VARIANCES_TEST_INIT
#endif
        )
    {}
    std::vector<Stat_> mean;
 
    std::vector<Stat_> variance;
 
    std::vector<Stat_> fitted;
 
    std::vector<Stat_> residual;
};
 
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.mean = output.mean.data();
    buffers.variance = output.variance.data();
 
    if (options.trend) {
        buffers.fitted = output.fitted.data();
        buffers.residual = output.residual.data();
    } else {
        buffers.fitted = NULL;
        buffers.residual = NULL;
    }
 
    model_gene_variances(mat, std::move(buffers), options);
    return output;
}
 
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;
};
 
template<typename Stat_, typename Index_>
void extract_blocked_weights(
    const std::size_t num_blocks,
    const std::vector<Stat_>& block_weights,
    const std::vector<Index_>& block_sizes,
    const Index_ min_size,
    std::vector<Stat_>& tmp_weights
) {
    assert(sanisizer::is_equal(num_blocks, block_weights.size()));
    assert(sanisizer::is_equal(num_blocks, block_sizes.size()));
    tmp_weights.clear();
    for (std::size_t b = 0; b < num_blocks; ++b) {
        if (block_sizes[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_blocked_pointers(
    const std::size_t num_blocks,
    const std::vector<ModelGeneVariancesBuffers<Stat_> >& per_block, 
    const std::vector<Index_>& block_sizes,
    const Index_ min_size,
    const Function_ fun,
    std::vector<Stat_*>& tmp_pointers
) {
    assert(sanisizer::is_equal(num_blocks, per_block.size()));
    assert(sanisizer::is_equal(num_blocks, block_sizes.size()));
    tmp_pointers.clear();
    for (std::size_t b = 0; b < num_blocks; ++b) {
        if (block_sizes[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 std::size_t num_blocks,
    const ModelGeneVariancesBlockedBuffers<Stat_>& buffers,
    const ModelGeneVariancesOptions& options
) {
    if (!sanisizer::is_equal(num_blocks, buffers.per_block.size())) {
        throw std::runtime_error("length of 'buffers.per_block' is not equal to 'num_blocks'");
    }
    assert(mat.ncol() == 0 || sanisizer::is_less_than(*std::max_element(block, block + mat.ncol()), num_blocks));
 
    tatami_stats::GroupRssBuffers<Stat_, Index_> vbuf;
    vbuf.mean.reserve(num_blocks);
    vbuf.rss.reserve(num_blocks);
    for (std::size_t b = 0; b < num_blocks; ++b) {
        vbuf.mean.push_back(buffers.per_block[b].mean);
        vbuf.rss.push_back(buffers.per_block[b].variance);
    }
 
    auto block_sizes = sanisizer::create<std::vector<Index_> >(num_blocks);
    vbuf.count = block_sizes.data();
 
    // Using group_rss() instead of group_variance(), as the latter doesn't pass back the block sizes yet.
    tatami_stats::GroupRssOptions vopt;
    vopt.num_threads = options.num_threads;
    tatami_stats::group_rss(true, mat, block, num_blocks, vbuf, vopt);
 
    const auto NR = mat.nrow();
    for (std::size_t b = 0; b < num_blocks; ++b) {
        quickstats::rss_to_variance(NR, block_sizes[b], vbuf.rss[b]); 
    }
 
    FitVarianceTrendWorkspace<Stat_> work;
    auto fopt = options.fit_variance_trend_options;
    fopt.num_threads = options.num_threads;
    bool all_trends_fitted = true;
 
    for (std::size_t b = 0; b < num_blocks; ++b) {
        const auto& current = buffers.per_block[b];
        if (current.fitted == NULL || current.residual == NULL) {
            all_trends_fitted = false;
            continue;
        }
        if (block_sizes[b] >= 2) {
            fit_variance_trend(NR, current.mean, current.variance, current.fitted, current.residual, work, fopt);
        } else {
            std::fill_n(current.fitted, NR, std::numeric_limits<double>::quiet_NaN());
            std::fill_n(current.residual, NR, std::numeric_limits<double>::quiet_NaN());
        }
    }
 
    const auto ave_means = buffers.average.mean;
    const auto ave_variances = buffers.average.variance;
    const auto ave_fitted = buffers.average.fitted;
    const auto ave_residuals = buffers.average.residual;
 
    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(num_blocks);
 
    if (options.block_average_policy == BlockAveragePolicy::MEAN) {
        const auto block_weight = scran_blocks::compute_weights<Stat_>(block_sizes, options.block_weight_policy, options.variable_block_weight_parameters);
        std::vector<Stat_> tmp_weights;
        tmp_weights.reserve(num_blocks);
 
        if (ave_means) {
            extract_blocked_weights(num_blocks, block_weight, block_sizes, static_cast<Index_>(1), tmp_weights);
            extract_blocked_pointers(num_blocks, buffers.per_block, block_sizes, static_cast<Index_>(1), [](const auto& x) -> Stat_* { return x.mean; }, 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.
        extract_blocked_weights(num_blocks, block_weight, block_sizes, static_cast<Index_>(2), tmp_weights);
 
        if (ave_variances) {
            extract_blocked_pointers(num_blocks, buffers.per_block, block_sizes, static_cast<Index_>(2), [](const auto& x) -> Stat_* { return x.variance; }, tmp_pointers);
            scran_blocks::parallel_weighted_means(NR, tmp_pointers, tmp_weights.data(), ave_variances, /* skip_nan = */ false);
        }
 
        if (ave_fitted) {
            extract_blocked_pointers(num_blocks, buffers.per_block, block_sizes, 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) {
            extract_blocked_pointers(num_blocks, buffers.per_block, block_sizes, static_cast<Index_>(2), [](const auto& x) -> Stat_* { return x.residual; }, 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) {
            extract_blocked_pointers(num_blocks, buffers.per_block, block_sizes, static_cast<Index_>(1), [](const auto& x) -> Stat_* { return x.mean; }, tmp_pointers);
            scran_blocks::parallel_quantiles(NR, tmp_pointers, options.block_quantile, ave_means, /* skip_nan = */ false);
        }
 
        // Again, skip blocks without enough cells to compute the variance.
 
        if (ave_variances) {
            extract_blocked_pointers(num_blocks, buffers.per_block, block_sizes, static_cast<Index_>(2), [](const auto& x) -> Stat_* { return x.variance; }, tmp_pointers);
            scran_blocks::parallel_quantiles(NR, tmp_pointers, options.block_quantile, ave_variances, /* skip_nan = */ false);
        }
 
        if (ave_fitted) {
            extract_blocked_pointers(num_blocks, buffers.per_block, block_sizes, 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) {
            extract_blocked_pointers(num_blocks, buffers.per_block, block_sizes, static_cast<Index_>(2), [](const auto& x) -> Stat_* { return x.residual; }, tmp_pointers);
            scran_blocks::parallel_quantiles(NR, tmp_pointers, options.block_quantile, ave_residuals, /* skip_nan = */ false);
        }
    }
}
 
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 std::size_t num_blocks,
    const ModelGeneVariancesOptions& options
) {
    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.
        num_blocks,
        do_average,
        options.trend
    );
 
    ModelGeneVariancesBlockedBuffers<Stat_> buffers;
    sanisizer::resize(buffers.per_block, num_blocks);
    for (std::size_t b = 0; b < num_blocks; ++b) {
        auto& current = buffers.per_block[b];
        current.mean = output.per_block[b].mean.data();
        current.variance = output.per_block[b].variance.data();
 
        if (options.trend) {
            current.fitted = output.per_block[b].fitted.data();
            current.residual = output.per_block[b].residual.data();
        } else {
            current.fitted = NULL;
            current.residual = NULL;
        }
    }
 
    if (!do_average) {
        buffers.average.mean = NULL;
        buffers.average.variance = NULL;
        buffers.average.fitted = NULL;
        buffers.average.residual = NULL;
    } else {
        buffers.average.mean = output.average.mean.data();
        buffers.average.variance = output.average.variance.data();
 
        if (options.trend) {
            buffers.average.fitted = output.average.fitted.data();
            buffers.average.residual = output.average.residual.data();
        } else {
            buffers.average.fitted = NULL;
            buffers.average.residual = NULL;
        }
    }
 
    model_gene_variances_blocked(mat, block, num_blocks, buffers, options);
    return output;
}
 
}
 
#endif