simple_pca.hpp

Go to the documentation of this file.

#ifndef SCRAN_PCA_SIMPLE_PCA_HPP
#define SCRAN_PCA_SIMPLE_PCA_HPP
 
#include <vector>
#include <type_traits>
#include <algorithm>
#include <memory>
#include <optional>
 
#include "tatami/tatami.hpp"
#include "tatami_stats/tatami_stats.hpp"
#include "quickstats/quickstats.hpp"
#include "irlba/irlba.hpp"
#include "irlba/parallel.hpp"
#include "irlba_tatami/irlba_tatami.hpp"
#include "Eigen/Dense"
#include "sanisizer/sanisizer.hpp"
 
#include "utils.hpp"
 
namespace scran_pca {
 
template<typename EigenVector_ = Eigen::VectorXd>
struct SimplePcaOptions {
    SimplePcaOptions() {
        // Avoid throwing an error if too many PCs are requested.
        irlba_options.cap_number = true;
    }
    int number = 25;
 
    bool scale = false;
 
    bool transpose = true;
 
    bool realize_matrix = true;
 
    int num_threads = 1;
 
    irlba::Options<EigenVector_> irlba_options;
};
 
template<typename Value_, typename Index_, class EigenVector_>
void compute_row_means_and_variances(const tatami::Matrix<Value_, Index_>& mat, const int num_threads, EigenVector_& center_v, EigenVector_& scale_v) {
    tatami_stats::RssOptions<typename EigenVector_::Scalar> opts;
    opts.num_threads = num_threads;
    opts.mean_placeholder = 0; // rss() emits NaNs if there are no cells, we replace them with zeros to avoid downstream problems with propagation.
 
    tatami_stats::RssBuffers<typename EigenVector_::Scalar> buffers;
    buffers.mean = center_v.data();
    buffers.rss = scale_v.data();
    tatami_stats::rss(true, mat, buffers, opts);
 
    const auto ncells = mat.ncol();
    if (ncells > 1) {
        scale_v /= ncells - 1;
    }
}
 
template<class EigenVector_, class EigenMatrix_>
std::unique_ptr<irlba::Matrix<EigenVector_, EigenMatrix_> > prepare_deferred_matrix_for_irlba(
    std::unique_ptr<irlba::Matrix<EigenVector_, EigenMatrix_> > ptr,
    const SimplePcaOptions<EigenVector_>& options,
    const EigenVector_& center_v,
    const EigenVector_& scale_v
) {
    std::unique_ptr<irlba::Matrix<EigenVector_, EigenMatrix_> > alt;
    alt.reset(new irlba::CenteredMatrix<EigenVector_, EigenMatrix_, I<decltype(ptr)>, I<decltype(&center_v)> >(std::move(ptr), &center_v));
    ptr.swap(alt);
 
    if (options.scale) {
        alt.reset(new irlba::ScaledMatrix<EigenVector_, EigenMatrix_, I<decltype(ptr)>, I<decltype(&scale_v)> >(std::move(ptr), &scale_v, true, true));
        ptr.swap(alt);
    }
 
    return ptr;
}
 
template<class EigenMatrix_, typename Value_, typename Index_, class EigenVector_>
std::unique_ptr<irlba::Matrix<EigenVector_, EigenMatrix_> > prepare_sparse_matrix_for_irlba(
    const tatami::Matrix<Value_, Index_>& mat, 
    const SimplePcaOptions<EigenVector_>& options,
    EigenVector_& center_v,
    EigenVector_& scale_v,
    typename EigenVector_::Scalar& total_var
) {
    const auto ngenes = mat.nrow();
    sanisizer::resize(center_v, ngenes);
    sanisizer::resize(scale_v, ngenes);
    std::unique_ptr<irlba::Matrix<EigenVector_, EigenMatrix_> > output;
 
    if (options.realize_matrix) {
        // 'extracted' contains row-major contents...
        auto extracted = tatami::retrieve_compressed_sparse_contents<Value_, Index_>(
            mat, 
            /* row = */ true, 
            [&]{
                tatami::RetrieveCompressedSparseContentsOptions opt;
                opt.two_pass = false;
                opt.num_threads = options.num_threads;
                return opt;
            }()
        );
 
        // But we effectively transpose it to CSC with genes in columns.
        const Index_ ncells = mat.ncol();
        const auto sparse_ptr = new irlba::ParallelSparseMatrix<
            EigenVector_,
            EigenMatrix_,
            I<decltype(extracted.value)>,
            I<decltype(extracted.index)>,
            I<decltype(extracted.pointers)>
        >(
            ncells,
            ngenes,
            std::move(extracted.value),
            std::move(extracted.index),
            std::move(extracted.pointers), 
            true,
            options.num_threads
        );
        output.reset(sparse_ptr);
 
        tatami::parallelize([&](const int, const Index_ start, const Index_ length) -> void {
            const auto& pointers = sparse_ptr->get_pointers();
            const auto& values = sparse_ptr->get_values();
            quickstats::RssWorkspace<typename EigenVector_::Scalar> work;
 
            for (Index_ g = start, end = start + length; g < end; ++g) {
                const auto offset = pointers[g];
                const auto next_offset = pointers[g + 1]; // increment won't overflow as 'g + 1 <= end'.
                const Index_ num_nonzero = next_offset - offset;
                const auto results = quickstats::rss(ncells, num_nonzero, values.data() + offset, work);
                center_v.coeffRef(g) = results.mean;
                scale_v.coeffRef(g) = results.rss;
            }
        }, ngenes, options.num_threads);
 
        if (ncells > 1) { 
            // if there are fewer than 2 cells, scale_v will naturally be set to zero.
            scale_v /= ncells - 1;
        } else if (!ncells) {
            // override quickstats::rss()'s setting of the mean to NaN if there are no cells.
            std::fill(center_v.begin(), center_v.end(), 0);
        }
 
        total_var = process_scale_vector(options.scale, scale_v);
 
    } else {
        compute_row_means_and_variances(mat, options.num_threads, center_v, scale_v);
        total_var = process_scale_vector(options.scale, scale_v);
        output.reset(
            new irlba_tatami::Transposed<EigenVector_, EigenMatrix_, Value_, Index_, decltype(&mat)>(&mat, options.num_threads)
        ); 
    }
 
    return prepare_deferred_matrix_for_irlba(std::move(output), options, center_v, scale_v);
}
 
template<class EigenMatrix_, typename Value_, typename Index_, class EigenVector_>
std::unique_ptr<irlba::Matrix<EigenVector_, EigenMatrix_> > prepare_dense_matrix_for_irlba(
    const tatami::Matrix<Value_, Index_>& mat, 
    const SimplePcaOptions<EigenVector_>& options,
    EigenVector_& center_v,
    EigenVector_& scale_v,
    typename EigenVector_::Scalar& total_var
) {
    const Index_ ngenes = mat.nrow();
    sanisizer::resize(center_v, ngenes);
    sanisizer::resize(scale_v, ngenes);
 
    if (options.realize_matrix) {
        // Create a matrix with genes in columns.
        const Index_ ncells = mat.ncol();
        auto emat = std::make_unique<EigenMatrix_>(
            sanisizer::cast<I<decltype(std::declval<EigenMatrix_>().rows())> >(ncells),
            sanisizer::cast<I<decltype(std::declval<EigenMatrix_>().cols())> >(ngenes)
        );
 
        // By default, Eigen's matrices are column major. In such cases, because we want to do
        // a transposition, we pretend it's row major during the conversion.
        static_assert(!EigenMatrix_::IsRowMajor);
        tatami::convert_to_dense(
            mat,
            /* row_major = */ true,
            emat->data(),
            [&]{
                tatami::ConvertToDenseOptions opt;
                opt.num_threads = options.num_threads;
                return opt;
            }()
        );
 
        center_v.array() = emat->array().colwise().sum();
        if (ncells) { // if there are no cells, center_v will naturally be set to zero.
            center_v /= ncells;
        }
        emat->array().rowwise() -= center_v.adjoint().array(); // applying it to avoid wasting time with deferred operations inside IRLBA.
 
        scale_v.array() = emat->array().colwise().squaredNorm();
        if (ncells > 1) { // if there are fewer than 2 cells, scale_v will naturally be set to zero.
            scale_v /= ncells - 1;
        }
 
        total_var = process_scale_vector(options.scale, scale_v);
        if (options.scale) {
            emat->array().rowwise() /= scale_v.adjoint().array();
        }
 
        return std::unique_ptr<irlba::Matrix<EigenVector_, EigenMatrix_> >( 
            new irlba::SimpleMatrix<EigenVector_, EigenMatrix_, decltype(emat)>(std::move(emat))
        );
 
    } else {
        compute_row_means_and_variances(mat, options.num_threads, center_v, scale_v);
        total_var = process_scale_vector(options.scale, scale_v);
        std::unique_ptr<irlba::Matrix<EigenVector_, EigenMatrix_> > output(
            new irlba_tatami::Transposed<EigenVector_, EigenMatrix_, Value_, Index_, decltype(&mat)>(&mat, options.num_threads)
        ); 
        return prepare_deferred_matrix_for_irlba(std::move(output), options, center_v, scale_v);
    }
}
template<typename EigenMatrix_, typename EigenVector_>
struct SimplePcaResults {
    EigenMatrix_ components;
 
    EigenVector_ variance_explained;
 
    typename EigenVector_::Scalar total_variance = 0;
 
    EigenMatrix_ rotation;
 
    EigenVector_ center;
 
    std::optional<EigenVector_> scale;
 
    irlba::Metrics metrics;
};
 
template<typename Value_, typename Index_, typename EigenMatrix_, class EigenVector_, class SubsetFunction_>
void simple_pca_internal(
    const tatami::Matrix<Value_, Index_>& mat,
    const SimplePcaOptions<EigenVector_>& options,
    SimplePcaResults<EigenMatrix_, EigenVector_>& output,
    SubsetFunction_ subset_fun
) {
    irlba::EigenThreadScope t(options.num_threads);
 
    auto scale = tatami::create_container_of_Index_size<EigenVector_>(mat.nrow());
    std::unique_ptr<irlba::Matrix<EigenVector_, EigenMatrix_> > ptr;
    if (mat.sparse()) {
        ptr = prepare_sparse_matrix_for_irlba<EigenMatrix_>(mat, options, output.center, scale, output.total_variance);
    } else {
        ptr = prepare_dense_matrix_for_irlba<EigenMatrix_>(mat, options, output.center, scale, output.total_variance);
    }
 
    output.metrics = irlba::compute(*ptr, options.number, output.components, output.rotation, output.variance_explained, options.irlba_options);
 
    subset_fun(output.components, output.variance_explained);
 
    clean_up(mat.ncol(), output.components, output.variance_explained);
    if (options.transpose) {
        output.components.adjointInPlace();
    }
 
    if (options.scale) {
        output.scale = std::move(scale);
    }
}
template<typename Value_, typename Index_, typename EigenMatrix_, class EigenVector_>
void simple_pca(const tatami::Matrix<Value_, Index_>& mat, const SimplePcaOptions<EigenVector_>& options, SimplePcaResults<EigenMatrix_, EigenVector_>& output) {
    simple_pca_internal(
        mat,
        options,
        output, 
        [](const EigenMatrix_&, const EigenVector_&) -> void {}
    );
}
 
template<typename EigenMatrix_ = Eigen::MatrixXd, class EigenVector_ = Eigen::VectorXd, typename Value_, typename Index_>
SimplePcaResults<EigenMatrix_, EigenVector_> simple_pca(const tatami::Matrix<Value_, Index_>& mat, const SimplePcaOptions<EigenVector_>& options) {
    SimplePcaResults<EigenMatrix_, EigenVector_> output;
    simple_pca(mat, options, output);
    return output;
}
 
}
 
#endif