ResidualPca.hpp

Go to the documentation of this file.

#ifndef SCRAN_RESIDUAL_PCA_HPP
#define SCRAN_RESIDUAL_PCA_HPP
 
#include "../utils/macros.hpp"
 
#include "tatami/tatami.hpp"
 
#include "irlba/irlba.hpp"
#include "Eigen/Dense"
 
#include <vector>
#include <cmath>
 
#include "utils.hpp"
#include "convert.hpp"
#include "wrappers.hpp"
#include "blocking.hpp"
 
namespace scran {
 
class ResidualPca {
public:
    struct Defaults {
        static constexpr int rank = 10;
 
        static constexpr bool scale = false;
 
        static constexpr bool transpose = true;
 
        static constexpr int num_threads = 1;
 
        static constexpr WeightPolicy block_weight_policy = WeightPolicy::VARIABLE;
 
        static constexpr VariableBlockWeightParameters variable_block_weight_parameters = VariableBlockWeightParameters();
 
        static constexpr bool return_rotation = false;
 
        static constexpr bool return_center = false;
 
        static constexpr bool return_scale = false;
    };
 
private:
    bool scale = Defaults::scale;
    bool transpose = Defaults::transpose;
    int rank = Defaults::rank;
 
    bool return_rotation = Defaults::return_rotation;
    bool return_center = Defaults::return_center;
    bool return_scale = Defaults::return_scale;
 
    WeightPolicy block_weight_policy = Defaults::block_weight_policy;
    VariableBlockWeightParameters variable_block_weight_parameters = Defaults::variable_block_weight_parameters;
 
    int nthreads = Defaults::num_threads;
 
public:
    ResidualPca& set_rank(int r = Defaults::rank) {
        rank = r;
        return *this;
    }
 
    ResidualPca& set_scale(bool s = Defaults::scale) {
        scale = s;
        return *this;
    }
 
    ResidualPca& set_transpose(bool t = Defaults::transpose) {
        transpose = t;
        return *this;
    }
 
    ResidualPca& set_return_rotation(bool r = Defaults::return_rotation) {
        return_rotation = r;
        return *this;
    }
 
    ResidualPca& set_return_center(bool r = Defaults::return_center) {
        return_center = r;
        return *this;
    }
 
    ResidualPca& set_return_scale(bool r = Defaults::return_scale) {
        return_scale = r;
        return *this;
    }
 
    ResidualPca& set_block_weight_policy(WeightPolicy w = Defaults::block_weight_policy) {
        block_weight_policy = w;
        return *this;
    }
 
    ResidualPca& set_variable_block_weight_parameters(VariableBlockWeightParameters v = Defaults::variable_block_weight_parameters) {
        variable_block_weight_parameters = v;
        return *this;
    }
 
    ResidualPca& set_num_threads(int n = Defaults::num_threads) {
        nthreads = n;
        return *this;
    }
 
private:
    template<bool weight_, typename Data_, typename Index_, typename Block_>
    void run_sparse(
        const tatami::Matrix<Data_, Index_>* mat, 
        const Block_* block, 
        const pca_utils::BlockingDetails<weight_>& block_details, 
        const irlba::Irlba& irb,
        Eigen::MatrixXd& pcs, 
        Eigen::MatrixXd& rotation, 
        Eigen::VectorXd& variance_explained, 
        Eigen::MatrixXd& center_m,
        Eigen::VectorXd& scale_v,
        double& total_var) 
    const {
        auto ngenes = mat->nrow(), ncells = mat->ncol(); 
        auto extracted = pca_utils::extract_sparse_for_pca(mat, nthreads); // row-major filling.
        pca_utils::SparseMatrix emat(ncells, ngenes, std::move(extracted.values), std::move(extracted.indices), std::move(extracted.ptrs), nthreads); // CSC with genes in columns.
 
        auto nblocks = block_details.num_blocks();
        center_m.resize(nblocks, ngenes);
        scale_v.resize(ngenes);
        pca_utils::compute_mean_and_variance_regress<weight_>(emat, block, block_details, center_m, scale_v, nthreads);
        total_var = pca_utils::process_scale_vector(scale, scale_v);
 
        pca_utils::RegressWrapper<decltype(emat), Block_> centered(&emat, block, &center_m);
        if constexpr(weight_) {
            if (scale) {
                irlba::Scaled<decltype(centered)> scaled(&centered, &scale_v);
                pca_utils::SampleScaledWrapper<decltype(scaled)> weighted(&scaled, &(block_details.expanded_weights));
                irb.run(weighted, pcs, rotation, variance_explained);
            } else {
                pca_utils::SampleScaledWrapper<decltype(centered)> weighted(&centered, &(block_details.expanded_weights));
                irb.run(weighted, pcs, rotation, variance_explained);
            }
 
            // This transposes 'pcs' to be a NDIM * NCELLS matrix.
            pca_utils::project_sparse_matrix(emat, pcs, rotation, scale, scale_v, nthreads);
 
            // Subtracting each block's mean from the PCs.
            Eigen::MatrixXd centering;
            if (scale) {
                centering = (center_m * (rotation.array().colwise() / scale_v.array()).matrix()).adjoint();
            } else {
                centering = (center_m * rotation).adjoint();
            }
            for (size_t i = 0, iend = pcs.cols(); i < iend; ++i) {
                pcs.col(i) -= centering.col(block[i]);
            }
 
            pca_utils::clean_up_projected<true>(pcs, variance_explained);
            if (!transpose) {
                pcs.adjointInPlace();
            }
 
        } else {
            if (scale) {
                irlba::Scaled<decltype(centered)> scaled(&centered, &scale_v);
                irb.run(scaled, pcs, rotation, variance_explained);
            } else {
                irb.run(centered, pcs, rotation, variance_explained);
            }
 
            pca_utils::clean_up(mat->ncol(), pcs, variance_explained);
            if (transpose) {
                pcs.adjointInPlace();
            }
        }
    }
 
    template<bool weight_, typename Data_, typename Index_, typename Block_>
    void run_dense(
        const tatami::Matrix<Data_, Index_>* mat, 
        const Block_* block,
        const pca_utils::BlockingDetails<weight_>& block_details, 
        const irlba::Irlba& irb,
        Eigen::MatrixXd& pcs, 
        Eigen::MatrixXd& rotation, 
        Eigen::VectorXd& variance_explained, 
        Eigen::MatrixXd& center_m,
        Eigen::VectorXd& scale_v,
        double& total_var) 
    const {
        auto emat = pca_utils::extract_dense_for_pca(mat, nthreads); // get a column-major matrix with genes in columns.
 
        auto ngenes = emat.cols();
        auto nblocks = block_details.num_blocks();
        center_m.resize(nblocks, ngenes);
        scale_v.resize(ngenes);
        pca_utils::compute_mean_and_variance_regress<weight_>(emat, block, block_details, center_m, scale_v, nthreads);
        total_var = pca_utils::process_scale_vector(scale, scale_v);
 
        // Applying the centering and scaling directly so that we can run the PCA with no or fewer layers.
        tatami::parallelize([&](size_t, size_t start, size_t length) -> void {
            size_t ncells = emat.rows();
            double* ptr = emat.data() + static_cast<size_t>(start) * ncells;
            for (size_t g = start, end = start + length; g < end; ++g, ptr += ncells) {
                for (size_t c = 0; c < ncells; ++c) {
                    ptr[c] -= center_m.coeff(block[c], g);
                }
 
                if (scale) {
                    auto sd = scale_v[g];
                    for (size_t c = 0; c < ncells; ++c) {
                        ptr[c] /= sd; // process_scale_vector should already protect against division by zero.
                    }
                }
            }
        }, ngenes, nthreads);
 
        if constexpr(weight_) {
            pca_utils::SampleScaledWrapper<decltype(emat)> weighted(&emat, &(block_details.expanded_weights));
            irb.run(weighted, pcs, rotation, variance_explained);
            pcs.noalias() = emat * rotation;
            pca_utils::clean_up_projected<false>(pcs, variance_explained);
        } else {
            irb.run(emat, pcs, rotation, variance_explained);
            pca_utils::clean_up(pcs.rows(), pcs, variance_explained);
        }
 
        if (transpose) {
            pcs.adjointInPlace();
        }
    }
 
    template<typename Data_, typename Index_, typename Block_>
    void run_internal(
        const tatami::Matrix<Data_, Index_>* mat, 
        const Block_* block,
        Eigen::MatrixXd& pcs, 
        Eigen::MatrixXd& rotation, 
        Eigen::VectorXd& variance_explained, 
        Eigen::MatrixXd& center_m,
        Eigen::VectorXd& scale_v,
        double& total_var) 
    const {
        irlba::EigenThreadScope t(nthreads);
        irlba::Irlba irb;
        irb.set_number(rank);
        irb.set_cap_number(true);
 
        if (block_weight_policy == WeightPolicy::NONE) {
            auto bdetails = pca_utils::compute_blocking_details(mat->ncol(), block);
            if (mat->sparse()) {
                run_sparse<false>(mat, block, bdetails, irb, pcs, rotation, variance_explained, center_m, scale_v, total_var);
            } else {
                run_dense<false>(mat, block, bdetails, irb, pcs, rotation, variance_explained, center_m, scale_v, total_var);
            }
 
        } else {
            auto bdetails = pca_utils::compute_blocking_details(mat->ncol(), block, block_weight_policy, variable_block_weight_parameters);
            if (mat->sparse()) {
                run_sparse<true>(mat, block, bdetails, irb, pcs, rotation, variance_explained, center_m, scale_v, total_var);
            } else {
                run_dense<true>(mat, block, bdetails, irb, pcs, rotation, variance_explained, center_m, scale_v, total_var);
            }
        }
    }
 
public:
    struct Results {
        Eigen::MatrixXd pcs;
 
        Eigen::VectorXd variance_explained;
 
        double total_variance = 0;
 
        Eigen::MatrixXd rotation;
 
        Eigen::MatrixXd center;
 
        Eigen::VectorXd scale;
    };
 
    template<typename Data_, typename Index_, typename Block_>
    Results run(const tatami::Matrix<Data_, Index_>* mat, const Block_* block) const {
        Results output;
        Eigen::MatrixXd rotation, center;
        Eigen::VectorXd scale;
 
        run_internal(mat, block, output.pcs, rotation, output.variance_explained, center, scale, output.total_variance);
 
        // Shifting them if we want to keep them.
        if (return_rotation) {
            output.rotation = std::move(rotation);
        }
        if (return_center) {
            output.center = center.adjoint();
        }
        if (return_scale) {
            output.scale = std::move(scale);
        }
 
        return output;
    }
 
    template<typename Data_, typename Index_, typename Block_, typename Subset_>
    Results run(const tatami::Matrix<Data_, Index_>* mat, const Block_* block, const Subset_* features) const {
        if (!features) {
            return run(mat, block);
        } else {
            auto subsetted = pca_utils::subset_matrix_by_features(mat, features);
            return run(subsetted.get(), block);
        }
    }
};
 
}
 
#endif