Add lassosum RSS method with C++ coordinate descent

NAMESPACE

@@ -51,7 +51,10 @@ export(get_susie_result)
 export(glmnet_weights)
 export(harmonize_gwas)
 export(harmonize_twas)
+export(invert_minmax_scaling)
 export(lasso_weights)
+export(lassosum_rss)
+export(lassosum_rss_weights)
 export(lbf_to_alpha)
 export(load_LD_matrix)
 export(load_genotype_region)

R/RcppExports.R

@@ -23,6 +23,10 @@ dentist_iterative_impute <- function(LD_mat, nSample, zScore, pValueThreshold, p
     .Call('_pecotmr_dentist_iterative_impute', PACKAGE = 'pecotmr', LD_mat, nSample, zScore, pValueThreshold, propSVD, gcControl, nIter, gPvalueThreshold, ncpus, correct_chen_et_al_bug, verbose)
 }
 
+lassosum_rss_rcpp <- function(z, LD, lambda, thr, maxiter) {
+    .Call('_pecotmr_lassosum_rss_rcpp', PACKAGE = 'pecotmr', z, LD, lambda, thr, maxiter)
+}
+
 prs_cs_rcpp <- function(a, b, phi, bhat, maf, n, ld_blk, n_iter, n_burnin, thin, verbose, seed) {
     .Call('_pecotmr_prs_cs_rcpp', PACKAGE = 'pecotmr', a, b, phi, bhat, maf, n, ld_blk, n_iter, n_burnin, thin, verbose, seed)
 }

R/misc.R

@@ -153,6 +153,10 @@ is_zero_variance <- function(x) length(unique(x)) == 1
 #'   drops trailing samples so that \code{nrow(X)} is a multiple of 4, matching
 #'   PLINK .bed file chunk processing. Ignored when \code{method = "sample"}.
 #'   Default is \code{FALSE}.
+#' @param shrinkage Numeric in (0, 1]. Shrink the LD matrix toward the identity:
+#'   \code{R_s = (1 - shrinkage) * R + shrinkage * I}. Useful for regularizing
+#'   LD for summary-statistics-based methods such as lassosum (Mak et al 2017).
+#'   Default is 0 (no shrinkage).
 #'
 #' @return A symmetric correlation matrix with row and column names taken from
 #'   \code{colnames(X)}.
@@ -188,7 +192,7 @@ is_zero_variance <- function(x) length(unique(x)) == 1
 #'
 #' @export
 compute_LD <- function(X, method = c("sample", "population"),
-                       trim_samples = FALSE) {
+                       trim_samples = FALSE, shrinkage = 0) {
   if (is.null(X)) {
     stop("X must be provided.")
   }
@@ -235,6 +239,13 @@ compute_LD <- function(X, method = c("sample", "population"),
   # Ensure clean output
   diag(R) <- 1.0
   R[is.na(R) | is.nan(R)] <- 0
+
+  # Optional shrinkage toward identity: R_s = (1 - shrinkage) * R + shrinkage * I
+  # Used e.g. by lassosum (Mak et al 2017) to regularize LD for RSS methods.
+  if (shrinkage > 0 && shrinkage <= 1) {
+    R <- (1 - shrinkage) * R + shrinkage * diag(nrow(R))
+  }
+
   colnames(R) <- rownames(R) <- nms
   R
 }

R/regularized_regression.R

@@ -782,3 +782,79 @@ bayes_c_rss_weights <- function(sumstats, LD, ...) {
 bayes_r_rss_weights <- function(sumstats, LD, ...) {
   return(bayes_alphabet_rss_weights(sumstats, LD, method = "bayesR", ...))
 }
+
+#' Lassosum RSS: LASSO on summary statistics with LD reference
+#'
+#' Coordinate descent to solve the penalized regression on summary statistics:
+#' \deqn{f(\beta) = \beta' R \beta - 2\beta' r + 2\lambda ||\beta||_1}
+#' where \eqn{R} is the LD matrix (pre-shrunk if desired) and \eqn{r = \hat\beta / \sqrt{n}}.
+#'
+#' Based on Mak et al (2017) "Polygenic scores via penalized regression on summary statistics",
+#' Genetic Epidemiology 41(6):469-480.
+#'
+#' @param bhat A vector of marginal effect sizes.
+#' @param LD A list of LD blocks, where each element is a matrix representing an LD block.
+#'   If shrinkage is desired, apply it before passing (e.g., \code{(1-s)*R + s*I}).
+#' @param n Sample size of the GWAS.
+#' @param lambda A vector of L1 penalty values. Default: 20 values from 0.001 to 0.1 on log scale.
+#' @param thr Convergence threshold. Default: 1e-4.
+#' @param maxiter Maximum number of iterations. Default: 10000.
+#'
+#' @return A list containing:
+#'   \item{beta_est}{Posterior estimates of SNP effect sizes at best lambda.}
+#'   \item{beta}{Matrix of estimates (p x nlambda).}
+#'   \item{lambda}{The lambda values used.}
+#'   \item{conv}{Convergence indicators (1 = converged).}
+#'   \item{loss}{Quadratic loss at each lambda.}
+#'   \item{fbeta}{Full objective value at each lambda.}
+#'   \item{nparams}{Number of non-zero coefficients at each lambda.}
+#'
+#' @examples
+#' set.seed(42)
+#' p <- 10
+#' n <- 100
+#' bhat <- rnorm(p, sd = 0.1)
+#' R <- diag(p)
+#' for (i in 1:(p - 1)) {
+#'   R[i, i + 1] <- 0.3
+#'   R[i + 1, i] <- 0.3
+#' }
+#' LD <- list(blk1 = R)
+#' out <- lassosum_rss(bhat, LD, n)
+#' @export
+lassosum_rss <- function(bhat, LD, n,
+                         lambda = exp(seq(log(0.001), log(0.1), length.out = 20)),
+                         thr = 1e-4, maxiter = 10000) {
+  if (!is.list(LD)) {
+    stop("Please provide a valid list of LD blocks using 'LD'.")
+  }
+  if (missing(n) || n <= 0) {
+    stop("Please provide a valid sample size using 'n'.")
+  }
+  total_rows_in_LD <- sum(sapply(LD, nrow))
+  if (length(bhat) != total_rows_in_LD) {
+    stop("The length of 'bhat' must be the same as the sum of the number of rows of all elements in the 'LD' list.")
+  }
+
+  z <- bhat / sqrt(n)
+  order <- order(lambda, decreasing = TRUE)
+  result <- lassosum_rss_rcpp(z, LD, lambda[order], thr, maxiter)
+
+  # Reorder back to original lambda order
+  result$beta[, order] <- result$beta
+  result$conv[order] <- result$conv
+  result$loss[order] <- result$loss
+  result$fbeta[order] <- result$fbeta
+  result$lambda <- lambda
+  result$nparams <- as.integer(colSums(result$beta != 0))
+  result$beta_est <- as.numeric(result$beta[, which.min(result$fbeta)])
+  result
+}
+
+#' Extract weights from lassosum_rss function
+#' @return A numeric vector of the posterior SNP coefficients.
+#' @export
+lassosum_rss_weights <- function(stat, LD, ...) {
+  model <- lassosum_rss(bhat = stat$b, LD = list(blk1 = LD), n = median(stat$n), ...)
+  return(model$beta_est)
+}

src/RcppExports.cpp

@@ -44,6 +44,21 @@ BEGIN_RCPP
     return rcpp_result_gen;
 END_RCPP
 }
+// lassosum_rss_rcpp
+Rcpp::List lassosum_rss_rcpp(const arma::vec& z, const Rcpp::List& LD, const arma::vec& lambda, double thr, int maxiter);
+RcppExport SEXP _pecotmr_lassosum_rss_rcpp(SEXP zSEXP, SEXP LDSEXP, SEXP lambdaSEXP, SEXP thrSEXP, SEXP maxiterSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< const arma::vec& >::type z(zSEXP);
+    Rcpp::traits::input_parameter< const Rcpp::List& >::type LD(LDSEXP);
+    Rcpp::traits::input_parameter< const arma::vec& >::type lambda(lambdaSEXP);
+    Rcpp::traits::input_parameter< double >::type thr(thrSEXP);
+    Rcpp::traits::input_parameter< int >::type maxiter(maxiterSEXP);
+    rcpp_result_gen = Rcpp::wrap(lassosum_rss_rcpp(z, LD, lambda, thr, maxiter));
+    return rcpp_result_gen;
+END_RCPP
+}
 // prs_cs_rcpp
 Rcpp::List prs_cs_rcpp(double a, double b, Rcpp::Nullable<double> phi, Rcpp::NumericVector bhat, Rcpp::Nullable<Rcpp::NumericVector> maf, int n, Rcpp::List ld_blk, int n_iter, int n_burnin, int thin, bool verbose, Rcpp::Nullable<unsigned int> seed);
 RcppExport SEXP _pecotmr_prs_cs_rcpp(SEXP aSEXP, SEXP bSEXP, SEXP phiSEXP, SEXP bhatSEXP, SEXP mafSEXP, SEXP nSEXP, SEXP ld_blkSEXP, SEXP n_iterSEXP, SEXP n_burninSEXP, SEXP thinSEXP, SEXP verboseSEXP, SEXP seedSEXP) {
@@ -113,6 +128,7 @@ END_RCPP
 static const R_CallMethodDef CallEntries[] = {
     {"_pecotmr_compute_LD_gcta_cpp", (DL_FUNC) &_pecotmr_compute_LD_gcta_cpp, 2},
     {"_pecotmr_dentist_iterative_impute", (DL_FUNC) &_pecotmr_dentist_iterative_impute, 11},
+    {"_pecotmr_lassosum_rss_rcpp", (DL_FUNC) &_pecotmr_lassosum_rss_rcpp, 5},
     {"_pecotmr_prs_cs_rcpp", (DL_FUNC) &_pecotmr_prs_cs_rcpp, 12},
     {"_pecotmr_qtl_enrichment_rcpp", (DL_FUNC) &_pecotmr_qtl_enrichment_rcpp, 7},
     {"_pecotmr_sdpr_rcpp", (DL_FUNC) &_pecotmr_sdpr_rcpp, 16},