GroupedSizeFactors.hpp

Go to the documentation of this file.

#ifndef SCRAN_CLUSTERED_SIZE_FACTORS_HPP
#define SCRAN_CLUSTERED_SIZE_FACTORS_HPP
 
#include "../utils/macros.hpp"
 
#include "MedianSizeFactors.hpp"
#include "SanitizeSizeFactors.hpp"
#include "CenterSizeFactors.hpp"
#include "../aggregation/AggregateAcrossCells.hpp"
 
#include "tatami/tatami.hpp"
 
#include "../utils/blocking.hpp"
 
#include <memory>
#include <vector>
#include <algorithm>
 
namespace scran {
 
class GroupedSizeFactors {
public:
    struct Defaults {
        static constexpr bool center = true;
 
        static constexpr bool handle_zeros = false;
 
        static constexpr bool handle_non_finite = false;
 
        static constexpr int num_threads = 1;
    };
 
public:
    GroupedSizeFactors& set_center(bool c = Defaults::center) {
        center = c;
        return *this;
    }
 
    GroupedSizeFactors& set_prior_count(double p = MedianSizeFactors::Defaults::prior_count) {
        prior_count = p;
        return *this;
    }
 
    GroupedSizeFactors& set_handle_zeros(bool z = Defaults::handle_zeros) {
        handle_zeros = z;
        return *this;
    }
 
    GroupedSizeFactors& set_handle_non_finite(bool n = Defaults::handle_non_finite) {
        handle_non_finite = n;
        return *this;
    }
 
    GroupedSizeFactors& set_num_threads(int n = Defaults::num_threads) {
        num_threads = n;
        return *this;
    }
 
private:
    bool center = Defaults::center;
    double prior_count = MedianSizeFactors::Defaults::prior_count;
    int num_threads = Defaults::num_threads;
    bool handle_zeros = Defaults::handle_zeros;
    bool handle_non_finite = Defaults::handle_non_finite;
 
public:
    template<typename T, typename IDX, typename Group, typename Out>
    void run(const tatami::Matrix<T, IDX>* mat, const Group* group, Out* output) const {
        run_internal(mat, group, false, output);
    }
 
    template<typename T, typename IDX, typename Group, typename Out>
    void run(const tatami::Matrix<T, IDX>* mat, const Group* group, size_t reference, Out* output) const {
        run_internal(mat, group, reference, output);
    }
 
private:
    template<typename T, typename IDX, typename Group, typename Ref, typename Out>
    void run_internal(const tatami::Matrix<T, IDX>* mat, const Group* group, Ref reference, Out* output) const {
        size_t NR = mat->nrow(), NC = mat->ncol();
        if (!NC) {
            return;
        }
 
        // Aggregating each group to get a pseudo-bulk sample.
        auto ngroups = count_ids(NC, group);
        std::vector<double> combined(ngroups * NR);
        {
            std::vector<double*> sums(ngroups);
            for (size_t i = 0; i < ngroups; ++i) {
                sums[i] = combined.data() + i * NR;
            }
            AggregateAcrossCells aggregator;
            aggregator.set_num_threads(num_threads).run(mat, group, std::move(sums), std::vector<int*>());
        }
 
        size_t ref = 0;
        if constexpr(std::is_same<Ref, size_t>::value) {
            ref = reference;
        } else {
            // Choosing one of them to be the reference. Here, we borrow some logic
            // from edgeR and use the one with the largest sum of square roots,
            // which provides a compromise between transcriptome coverage and
            // complexity. The root ensures that we don't pick a sample that just
            // has very high expression in a small subset of genes, while still
            // remaining responsive to the overall coverage level.
            double best = 0;
 
            for (size_t i = 0; i < ngroups; ++i) {
                auto start = combined.data() + i * NR;
                double current = 0;
                for (size_t j = 0; j < NR; ++j) {
                    current += std::sqrt(start[j]);
                }
                if (current > best) {
                    ref = i;
                    best = current;
                }
            }
        }
 
        // Computing median-based size factors. No need to center
        // here as we'll be recentering afterwards anyway.
        tatami::ArrayView view(combined.data(), combined.size());
        tatami::DenseColumnMatrix<T, IDX, decltype(view)> aggmat(NR, ngroups, std::move(view));
 
        MedianSizeFactors med;
        med.set_num_threads(num_threads).set_center(false).set_prior_count(prior_count);
        auto mres = med.run(&aggmat, combined.data() + ref * NR);
 
        // We need to rescue odd size factors here; if they're zero or infinite,
        // we lose information about the relative scaling of cells within each group.
        auto cres = validate_size_factors(mres.factors.size(), mres.factors.data());
        SanitizeSizeFactors sanitizer;
        sanitizer.set_handle_zero(handle_zeros ? SanitizeSizeFactors::HandlerAction::SANITIZE : SanitizeSizeFactors::HandlerAction::ERROR);
        sanitizer.set_handle_non_finite(handle_non_finite ? SanitizeSizeFactors::HandlerAction::SANITIZE : SanitizeSizeFactors::HandlerAction::ERROR);
        sanitizer.run(mres.factors.size(), mres.factors.data(), cres);
 
        // Propagating to each cell via library size-based normalization.
        auto aggcolsums = tatami::column_sums(&aggmat, num_threads);
        auto colsums = tatami::column_sums(mat, num_threads);
        for (size_t i = 0; i < NC; ++i) {
            auto curgroup = group[i];
            if (aggcolsums[curgroup]) { 
                auto scale = static_cast<double>(colsums[i])/static_cast<double>(aggcolsums[curgroup]);
                output[i] = scale * mres.factors[curgroup];
            } else {
                // implies colsums[i] == 0, so instead of getting NaN size factors,
                // we just set the factors to zero, given that everything in the group is also zero.
                output[i] = 0; 
            }
        }
 
        if (center) {
            CenterSizeFactors centerer;
            centerer.run(NC, output);
        }
    }
 
public:
    template<typename Out>
    struct Results {
        Results(size_t NC) : factors(NC) {}
        std::vector<Out> factors;
    };
 
    template<typename Out = double, typename T, typename IDX, typename Group>
    Results<Out> run(const tatami::Matrix<T, IDX>* mat, const Group* group) const {
        Results<Out> output(mat->ncol());
        run(mat, group, output.factors.data());
        return output;
    }
 
    template<typename Out = double, typename T, typename IDX, typename Group>
    Results<Out> run(const tatami::Matrix<T, IDX>* mat, const Group* group, size_t reference) const {
        Results<Out> output(mat->ncol());
        run(mat, group, reference, output.factors.data());
        return output;
    }
};
 
}
 
#endif