// #include <Rcpp.h>
#include <RcppArmadillo.h>
#include "beachmat/numeric_matrix.h"
#include "beachmat/integer_matrix.h"
using namespace Rcpp;
// [[Rcpp::depends(RcppArmadillo)]]
// This correction factor is necessary to avoid estimates of
// theta that are basically +Inf. The problem is that for
// some combination of the y, mu, and X the term
// lgamma(1/theta) and the log(det(t(X) %*% W %*% X))
// with W = diag(1/(1/mu + theta)) canceled each other
// exactly out for large theta.
const double cr_correction_factor = 0.99;
// [[Rcpp::export]]
List make_table_if_small(const NumericVector& x, int stop_if_larger){
std::unordered_map<long, size_t> counts;
counts.reserve(stop_if_larger);
for (double v : x){
++counts[(long) v];
if(counts.size() > stop_if_larger){
return List::create(NumericVector::create(), NumericVector::create());
}
}
NumericVector keys(counts.size());
NumericVector values(counts.size());
transform(counts.begin(), counts.end(), keys.begin(), [](std::pair<int, size_t> pair){return (double) pair.first;});
transform(counts.begin(), counts.end(), values.begin(), [](std::pair<int, size_t> pair){return (double) pair.second;});
return List::create(keys, values);
}
//--------------------------------------------------------------------------------------------------//
// The following code was originally copied from https://github.com/mikelove/DESeq2/blob/master/src/DESeq2.cpp
// I adapted it to the needs of this project by:
// * renaming alpha -> theta for consitency
// * removing the part for the prior on theta
// * renaming x -> model_matrix
// * additional small changes
// * adding capability to calculate digamma/trigamma only
// on unique counts
/*
* DESeq2 C++ functions
*
* Author: Michael I. Love, Constantin Ahlmann-Eltze
* Last modified: May 21, 2020
* License: LGPL (>= 3)
*
* Note: The canonical, up-to-date DESeq2.cpp lives in
* the DESeq2 library, the development branch of which
* can be viewed here:
*
* https://github.com/mikelove/DESeq2/blob/master/src/DESeq2.cpp
*/
// this function returns the log posterior of dispersion parameter alpha, for negative binomial variables
// given the counts y, the expected means mu, the design matrix x (used for calculating the Cox-Reid adjustment),
// and the parameters for the normal prior on log alpha
// [[Rcpp::export]]
double conventional_loglikelihood_fast(NumericVector y, NumericVector mu, double log_theta, const arma::mat& model_matrix, bool do_cr_adj,
NumericVector unique_counts = NumericVector::create(),
NumericVector count_frequencies = NumericVector::create()) {
double theta = exp(log_theta);
double cr_term = 0.0;
if(do_cr_adj){
arma::vec w_diag = 1.0 / (1.0 / mu + theta);
arma::mat b = model_matrix.t() * (model_matrix.each_col() % w_diag);
// cr_term = -0.5 * log(det(b)) * cr_correction_factor;
arma::mat L, U, P;
arma::lu(L, U, P, b);
double ld = sum(log(arma::diagvec(L)));
arma::vec u_diag = arma::diagvec(U);
for(double e : u_diag){
ld += e < 1e-50 ? log(1e-50) : log(e);
}
cr_term = -0.5 * ld * cr_correction_factor;
}
double theta_neg1 = R_pow_di(theta, -1);
double lgamma_term = 0;
// If summarized counts are available use those to calculate sum(lgamma(y + theta_neg1))
if(unique_counts.size() > 0 && unique_counts.size() == count_frequencies.size()){
for(size_t iter = 0; iter < count_frequencies.size(); ++iter){
lgamma_term += count_frequencies[iter] * lgamma(unique_counts[iter] + theta_neg1);
}
}else{
lgamma_term = sum(lgamma(y + theta_neg1));
}
lgamma_term -= y.size() * lgamma(theta_neg1);
double ll_part = 0.0;
for(size_t i = 0; i < y.size(); ++i){
ll_part += (-y[i] - theta_neg1) * log(mu[i] + theta_neg1);
}
ll_part -= y.size() * theta_neg1 * log(theta);
return lgamma_term + ll_part + cr_term;
}
// this function returns the derivative of the log posterior with respect to the log of the
// dispersion parameter alpha, given the same inputs as the previous function
// [[Rcpp::export]]
double conventional_score_function_fast(NumericVector y, NumericVector mu, double log_theta, const arma::mat& model_matrix, bool do_cr_adj,
NumericVector unique_counts = NumericVector::create(),
NumericVector count_frequencies = NumericVector::create()) {
double theta = exp(log_theta);
double theta_neg1 = 1.0 / theta;
double cr_term = 0.0;
if(do_cr_adj){
arma::vec w_diag = 1.0 / (1.0 / mu + theta);
arma::vec dw_diag = -1 * w_diag % w_diag;
arma::mat b = model_matrix.t() * (model_matrix.each_col() % w_diag);
arma::mat db = model_matrix.t() * (model_matrix.each_col() % dw_diag);
// The diag(1e-6) protects against singular matrices
arma::mat b_inv = inv_sympd(b + arma::eye(b.n_rows, b.n_cols) * 1e-6);
cr_term = -0.5 * trace(b_inv * db) * cr_correction_factor;
}
double digamma_term = 0;
// If summarized counts are available use those to calculate sum(digamma(y + theta_neg1))
if(unique_counts.size() > 0 && unique_counts.size() == count_frequencies.size()){
double max_y = 0.0;
double sum_y = 0.0;
double sum_prod_y = 0.0;
for(size_t iter = 0; iter < count_frequencies.size(); ++iter){
digamma_term += count_frequencies[iter] * Rf_digamma(unique_counts[iter] + theta_neg1);
sum_y += count_frequencies[iter] * unique_counts[iter];
sum_prod_y += count_frequencies[iter] * (unique_counts[iter] - 1) * unique_counts[iter];
max_y = std::max(max_y, unique_counts[iter]);
}
double corr = theta_neg1 > 1e5 ? sum_prod_y / (2 * theta_neg1) : 0.0;
if(max_y * 1e6 < theta_neg1){
// This approximation is based on the fact that for large x
// (sum(digamma(y + x)) - length(y) * digamma(x)) * x \approx sum(y)
// Due to numerical imprecision the digamma_term reaches sum(y) sometimes
// quicker than the ll_term, thus I subtract the first term of the
// Laurent series expansion at x -> inf
digamma_term = sum_y - corr;
}else{
digamma_term -= y.size() * Rf_digamma(theta_neg1);
digamma_term *= theta_neg1;
digamma_term = std::min(digamma_term, sum_y - corr);
}
}else{
double max_y = 0.0;
double sum_y = 0.0;
double sum_prod_y = 0.0;
for(size_t iter = 0; iter < y.size(); ++iter){
digamma_term += Rf_digamma(y[iter] + theta_neg1);
sum_y += y[iter];
sum_prod_y += (y[iter] - 1) * y[iter];
max_y = std::max(max_y, y[iter]);
}
double corr = theta_neg1 > 1e5 ? sum_prod_y / (2 * theta_neg1) : 0.0;
if(max_y * 1e6 < theta_neg1){
digamma_term = sum_y - corr;
}else{
digamma_term -= y.size() * Rf_digamma(theta_neg1);
digamma_term *= theta_neg1;
digamma_term = std::min(digamma_term, sum_y - corr);
}
}
double ll_part = 0.0;
for(size_t i = 0; i < y.size(); ++i){
double mu_theta = (mu[i] * theta);
if(mu_theta < 1e-10){
ll_part += mu_theta * mu_theta * (1 / (1 + mu_theta) - 0.5);
}else if(mu_theta < 1e-4){
// The bounds are based on the Taylor expansion of log(1 + x) for x = 0.
double inv = 1 / (1 + mu_theta);
double upper_bound = mu_theta * mu_theta * inv;
double lower_bound = mu_theta * mu_theta * (inv - 0.5);
double suggest = (log(1 + mu_theta) - mu[i] / (mu[i] + theta_neg1)) ;
ll_part += std::max(std::min(suggest, upper_bound), lower_bound);
}else{
ll_part += log(1 + mu_theta) - mu[i] / (mu[i] + theta_neg1);
}
ll_part += y[i] / (mu[i] + theta_neg1);
}
ll_part *= theta_neg1;
return ll_part - digamma_term + cr_term * theta;
}
// this function returns the second derivative of the log posterior with respect to the log of the
// dispersion parameter alpha, given the same inputs as the previous function
// [[Rcpp::export]]
double conventional_deriv_score_function_fast(NumericVector y, NumericVector mu, double log_theta, const arma::mat& model_matrix, bool do_cr_adj,
NumericVector unique_counts = NumericVector::create(),
NumericVector count_frequencies = NumericVector::create()) {
double theta = exp(log_theta);
double cr_term = 0.0;
double cr_term2 = 0.0;
if(do_cr_adj){
arma::vec w_diag = 1/(1/mu + theta);
arma::vec dw_diag = -1 * w_diag % w_diag;
arma::vec d2w_diag = -2 * dw_diag % w_diag;
arma::mat b = model_matrix.t() * (model_matrix.each_col() % w_diag);
arma::mat db = model_matrix.t() * (model_matrix.each_col() % dw_diag);
arma::mat d2b = model_matrix.t() * (model_matrix.each_col() % d2w_diag);
// The diag(1e-6) protects against singular matrices
arma::mat b_inv = inv_sympd(b + arma::eye(b.n_rows, b.n_cols) * 1e-6);
arma::mat d_i_db = b_inv * db;
double ddetb = trace(d_i_db);
double d2detb = ((R_pow_di(ddetb, 2) - trace(d_i_db * d_i_db) + trace(b_inv * d2b)) );
cr_term = (0.5 * R_pow_di(ddetb, 2) - 0.5 * d2detb) * cr_correction_factor;
cr_term2 = -0.5 * ddetb * cr_correction_factor;
}
double theta_neg1 = R_pow_di(theta, -1);
double theta_neg2 = R_pow_di(theta, -2);
double digamma_term = 0.0;
double trigamma_term = 0.0;
// If summarized counts are available use those to calculate sum(digamma()) and sum(trigamma())
if(unique_counts.size() > 0 && unique_counts.size() == count_frequencies.size()){
for(size_t iter = 0; iter < count_frequencies.size(); ++iter){
digamma_term += count_frequencies[iter] * Rf_digamma(unique_counts[iter] + theta_neg1);
trigamma_term += count_frequencies[iter] * Rf_trigamma(unique_counts[iter] + theta_neg1);
}
trigamma_term *= theta_neg2;
digamma_term -= y.size() * Rf_digamma(theta_neg1);
trigamma_term -= theta_neg2 * y.size() * Rf_trigamma(theta_neg1);
}else{
digamma_term = sum(digamma(y + theta_neg1));
digamma_term -= y.size() * Rf_digamma(theta_neg1);
trigamma_term = theta_neg2 * sum(trigamma(y + theta_neg1));
trigamma_term -= theta_neg2 * y.size() * Rf_trigamma(theta_neg1);
}
double ll_part_1 = 0.0;
double ll_part_2 = 0.0;
for(size_t i = 0; i < y.size(); ++i){
ll_part_1 += log(1 + mu[i] * theta) + (y[i] - mu[i]) / (mu[i] + theta_neg1);
ll_part_2 += (mu[i] * mu[i] * theta + y[i]) / (1 + mu[i] * theta) / (1 + mu[i] * theta);
}
double ll_part = -2 * theta_neg1 * (ll_part_1 - digamma_term) + (ll_part_2 + trigamma_term);
double res = ll_part + cr_term * R_pow_di(theta, 2) + (ll_part_1 - digamma_term) * theta_neg1 + cr_term2 * theta;
return res;
}
// ------------------------------------------------------------------------------------------------
template<class NumericType>
List estimate_overdispersions_fast_internal(RObject Y, RObject mean_matrix, NumericMatrix model_matrix, bool do_cox_reid_adjustment,
double n_subsamples, int max_iter){
auto Y_bm = beachmat::create_matrix<NumericType>(Y);
auto mean_mat_bm = beachmat::create_numeric_matrix(mean_matrix);
int n_samples = Y_bm->get_ncol();
int n_genes = Y_bm->get_nrow();
NumericVector estimates(n_genes);
NumericVector iterations(n_genes);
CharacterVector messages(n_genes);
if(n_genes != mean_mat_bm->get_nrow() || n_samples != mean_mat_bm->get_ncol()){
throw std::runtime_error("Dimensions of Y and mean_matrix do not match");
}
// This is calling back to R, which simplifies my code a lot
Environment glmGamPoiEnv = Environment::namespace_env("glmGamPoi");
Function overdispersion_mle_impl = glmGamPoiEnv["overdispersion_mle_impl"];
for(int gene_idx = 0; gene_idx < n_genes; gene_idx++){
if (gene_idx % 100 == 0) checkUserInterrupt();
typename NumericType::vector counts(n_samples);
Y_bm->get_row(gene_idx, counts.begin());
NumericVector mu(n_samples);
mean_mat_bm->get_row(gene_idx, mu.begin());
// Check if the first value is NA, if yes all of them will be
if(n_samples > 0 && Rcpp::traits::is_na<REALSXP>(mu[0])){
estimates(gene_idx) = NA_REAL;
iterations(gene_idx) = max_iter;
messages(gene_idx) = "Mean estimate was NA. Cannot estimate overdispersion";
}else{
List dispRes = Rcpp::as<List>(overdispersion_mle_impl(counts, mu, model_matrix, do_cox_reid_adjustment, n_subsamples, max_iter));
estimates(gene_idx) = Rcpp::as<double>(dispRes["estimate"]);
iterations(gene_idx) = Rcpp::as<double>(dispRes["iterations"]);
messages(gene_idx) = Rcpp::as<String>(dispRes["message"]);
}
}
return List::create(
Named("estimate", estimates),
Named("iterations", iterations),
Named("message", messages));;
}
// [[Rcpp::export]]
List estimate_overdispersions_fast(RObject Y, RObject mean_matrix, NumericMatrix model_matrix, bool do_cox_reid_adjustment,
double n_subsamples, int max_iter){
auto mattype=beachmat::find_sexp_type(Y);
if (mattype==INTSXP) {
return estimate_overdispersions_fast_internal<beachmat::integer_matrix>(Y, mean_matrix, model_matrix, do_cox_reid_adjustment, n_subsamples, max_iter);
} else if (mattype==REALSXP) {
return estimate_overdispersions_fast_internal<beachmat::numeric_matrix>(Y, mean_matrix, model_matrix, do_cox_reid_adjustment, n_subsamples, max_iter);
} else {
throw std::runtime_error("unacceptable matrix type");
}
}
template<class NumericType>
NumericVector estimate_global_overdispersions_fast_internal(RObject Y, RObject mean_matrix, const arma::mat model_matrix, const bool do_cox_reid_adjustment,
const NumericVector log_thetas){
const auto Y_bm = beachmat::create_matrix<NumericType>(Y);
const auto mean_mat_bm = beachmat::create_numeric_matrix(mean_matrix);
int n_samples = Y_bm->get_ncol();
int n_genes = Y_bm->get_nrow();
int n_spline_points = log_thetas.size();
NumericVector log_likelihoods(n_spline_points);
for(int gene_idx = 0; gene_idx < n_genes; gene_idx++){
if (gene_idx % 100 == 0) checkUserInterrupt();
NumericVector counts(n_samples);
Y_bm->get_row(gene_idx, counts.begin());
NumericVector mu(n_samples);
mean_mat_bm->get_row(gene_idx, mu.begin());
ListOf<NumericVector> tab = List::create(NumericVector::create(), NumericVector::create());
tab = make_table_if_small(counts, /*stop_if_larger = */ n_samples / 2);
for(int point_idx = 0; point_idx < n_spline_points; point_idx++){
log_likelihoods[point_idx] += conventional_loglikelihood_fast(counts, mu, log_thetas[point_idx], model_matrix,
do_cox_reid_adjustment, tab[0], tab[1]);
}
}
return log_likelihoods;
}
// [[Rcpp::export]]
NumericVector estimate_global_overdispersions_fast(RObject Y, RObject mean_matrix, const arma::mat model_matrix, const bool do_cox_reid_adjustment,
const NumericVector log_thetas){
auto mattype=beachmat::find_sexp_type(Y);
if (mattype==INTSXP) {
return estimate_global_overdispersions_fast_internal<beachmat::integer_matrix>(Y, mean_matrix, model_matrix, do_cox_reid_adjustment, log_thetas);
} else if (mattype==REALSXP) {
return estimate_global_overdispersions_fast_internal<beachmat::numeric_matrix>(Y, mean_matrix, model_matrix, do_cox_reid_adjustment, log_thetas);
} else {
throw std::runtime_error("unacceptable matrix type");
}
}
