normalize_counts.hpp

Go to the documentation of this file.

#ifndef SCRAN_NORM_NORMALIZE_COUNTS_HPP
#define SCRAN_NORM_NORMALIZE_COUNTS_HPP
 
#include <type_traits>
#include <vector>
#include <memory>
#include <cassert>
 
#include "tatami/tatami.hpp"
 
namespace scran_norm {
 
template<typename OutputValue_, typename InputValue_, typename Index_, typename ReciprocalSizeFactors_>
class DelayedLogNormalizeHelper final : public tatami::DelayedUnaryIsometricOperationHelper<OutputValue_, InputValue_, Index_> {
public:
    DelayedLogNormalizeHelper(ReciprocalSizeFactors_ reciprocal_size_factors, OutputValue_ log_base, OutputValue_ pseudo_count) : 
        my_reciprocal_sf(std::move(reciprocal_size_factors)),
        my_reciprocal_denom(1.0 / std::log(log_base)),
        my_pseudo(pseudo_count)
    {
        sanisizer::cast<Index_>(my_reciprocal_sf.size()); // check that cast is safe in ncol().
        if (my_pseudo != 1) {
            my_sparse = false;
        }
        for (const auto x : my_reciprocal_sf) {
            if (!std::isfinite(x)) {
                my_sparse = false;
                my_has_weird_sf = true;
                break;
            }
        }
    }
 
private:
    ReciprocalSizeFactors_ my_reciprocal_sf;
    OutputValue_ my_reciprocal_denom, my_pseudo;
    bool my_has_weird_sf = false, my_sparse = true;
 
public:
    std::optional<Index_> nrow() const {
        return std::nullopt;
    }
 
    std::optional<Index_> ncol() const {
        return my_reciprocal_sf.size();
    }
 
public:
    bool zero_depends_on_row() const {
        return false;
    }
 
    bool zero_depends_on_column() const {
        return my_has_weird_sf;
    }
 
    bool non_zero_depends_on_row() const {
        return false;
    }
 
    bool non_zero_depends_on_column() const {
        return true;
    }
 
private:
    void log_normalize(const Index_ idx, const Index_ length, const InputValue_* input, OutputValue_* const output) const {
        const auto current_rsf = my_reciprocal_sf[idx];
        for (Index_ i = 0; i < length; ++i) {
            output[i] = std::log(input[i] * current_rsf + my_pseudo) * my_reciprocal_denom;
        }
    }
 
public:
    void dense(const bool row, const Index_ idx, const Index_ start, const Index_ length, const InputValue_* input, OutputValue_* const output) const {
        if constexpr(std::is_same<InputValue_, OutputValue_>::value) {
            input = output; // basically an assertion to the compiler to skip aliasing protection.
        }
 
        if (row) {
            for (Index_ i = 0; i < length; ++i) {
                output[i] = std::log(input[i] * my_reciprocal_sf[i + start] + my_pseudo) * my_reciprocal_denom;
            }
        } else {
            log_normalize(idx, length, input, output);
        }
    }
 
    void dense(const bool row, const Index_ idx, const std::vector<Index_>& indices, const InputValue_* input, OutputValue_* const output) const {
        if constexpr(std::is_same<InputValue_, OutputValue_>::value) {
            input = output; // basically an assertion to the compiler to skip aliasing protection.
        }
        const Index_ length = indices.size();
 
        if (row) {
            for (Index_ i = 0; i < length; ++i) {
                output[i] = std::log(input[i] * my_reciprocal_sf[indices[i]] + my_pseudo) * my_reciprocal_denom;
            }
        } else {
            log_normalize(idx, length, input, output);
        }
    }
 
public:
    bool is_sparse() const {
        return my_sparse;
    }
 
    void sparse(
        const bool row,
        const Index_ idx,
        const Index_ number,
        const InputValue_* input_value,
        const Index_* const index,
        OutputValue_* const output_value
    ) const {
        if constexpr(std::is_same<InputValue_, OutputValue_>::value) {
            input_value = output_value; // basically an assertion to the compiler to skip aliasing protection.
        }
 
        if (row) {
            for (Index_ i = 0; i < number; ++i) {
                output_value[i] = std::log(input_value[i] * my_reciprocal_sf[index[i]] + my_pseudo) * my_reciprocal_denom;
            }
        } else {
            log_normalize(idx, number, input_value, output_value);
        }
    }
 
    OutputValue_ fill(const bool row, const Index_ idx) const {
        if (row) {
            // This should never be called in the presence of size factors of zero, as these will lead to NaNs.
            assert(!my_has_weird_sf);
            return std::log(my_pseudo) * my_reciprocal_denom;
        } else {
            return std::log(static_cast<InputValue_>(0) * my_reciprocal_sf[idx] + my_pseudo) * my_reciprocal_denom;
        }
    }
};
 
struct NormalizeCountsOptions {
    double pseudo_count = 1;
 
    bool preserve_sparsity = false;
 
    bool log = true;
 
    double log_base = 2;
};
 
template<typename OutputValue_ = double, typename InputValue_, typename Index_, class SizeFactors_>
std::shared_ptr<tatami::Matrix<OutputValue_, Index_> > normalize_counts(
    std::shared_ptr<const tatami::Matrix<InputValue_, Index_> > counts, 
    SizeFactors_ size_factors, 
    const NormalizeCountsOptions& options) 
{
    auto current_pseudo = options.pseudo_count;
    if (options.preserve_sparsity && current_pseudo != 1 && options.log) {
        for (auto& x : size_factors) { 
            x *= current_pseudo;
        }
        current_pseudo = 1;
    }
 
    static_assert(std::is_floating_point<OutputValue_>::value);
    if (static_cast<size_t>(size_factors.size()) != static_cast<size_t>(counts->ncol())) {
        throw std::runtime_error("length of 'size_factors' should be equal to the number of columns of 'counts'");
    }
 
    for (auto& x : size_factors) { 
        x = static_cast<OutputValue_>(1.0) / x;
    }
 
    if (!options.log) {
        return std::make_shared<tatami::DelayedUnaryIsometricOperation<OutputValue_, InputValue_, Index_> >(
            std::move(counts), 
            std::make_shared<tatami::DelayedUnaryIsometricMultiplyVectorHelper<OutputValue_, InputValue_, Index_, SizeFactors_> >(std::move(size_factors), false)
        );
    } else {
        return std::make_shared<tatami::DelayedUnaryIsometricOperation<OutputValue_, InputValue_, Index_> >(
            std::move(counts), 
            std::make_shared<DelayedLogNormalizeHelper<OutputValue_, InputValue_, Index_, SizeFactors_> >(std::move(size_factors), options.log_base, current_pseudo)
        );
    }
};
 
// Overload for template deduction.
template<typename OutputValue_ = double, typename InputValue_, typename Index_, class SizeFactors_>
std::shared_ptr<tatami::Matrix<OutputValue_, Index_> > normalize_counts(
    std::shared_ptr<tatami::Matrix<InputValue_, Index_> > counts,
    SizeFactors_ size_factors,
    const NormalizeCountsOptions& options)
{
    return normalize_counts(std::shared_ptr<const tatami::Matrix<InputValue_, Index_> >(std::move(counts)), std::move(size_factors), options);
}
}
 
#endif