From c22daa6458583c6c662b0a98adec89a60cd3c728 Mon Sep 17 00:00:00 2001 From: Daniel Nachun Date: Wed, 13 May 2026 13:31:13 -0700 Subject: [PATCH 1/7] TWAS ensemble fixes --- .github/environment/pixi.toml | 1 + .github/recipe/recipe.yaml | 2 + NAMESPACE | 1 + R/regularized_regression.R | 46 +- R/twas_weights.R | 149 +++- inst/prototype/ensemble_twas_weights.ipynb | 992 +++++++++++++++++++++ man/bayes_alphabet_weights.Rd | 1 + man/bayes_b_weights.Rd | 4 +- man/bayes_c_weights.Rd | 2 +- man/dpr_weights.Rd | 8 +- man/estimate_sparsity.Rd | 49 + man/mrash_weights.Rd | 2 +- man/twas_weights.Rd | 2 +- man/twas_weights_pipeline.Rd | 16 +- pixi.toml | 2 + 15 files changed, 1242 insertions(+), 35 deletions(-) create mode 100644 inst/prototype/ensemble_twas_weights.ipynb create mode 100644 man/estimate_sparsity.Rd diff --git a/.github/environment/pixi.toml b/.github/environment/pixi.toml index fbd2e2fe..d0672df9 100644 --- a/.github/environment/pixi.toml +++ b/.github/environment/pixi.toml @@ -39,5 +39,6 @@ r45 = {features = ["r45"]} "r-markdown" = "*" "r-pkgdown" = "*" "r-rcmdcheck" = "*" +"r-simxqtl" = "*" "r-testthat" = "*" "r-tidyverse" = "*" diff --git a/.github/recipe/recipe.yaml b/.github/recipe/recipe.yaml index 211c37ab..f933a2ab 100644 --- a/.github/recipe/recipe.yaml +++ b/.github/recipe/recipe.yaml @@ -53,6 +53,7 @@ requirements: - r-magrittr - r-mashr - r-matrixstats + - r-mr.ash.alpha - r-mr.mashr - r-mvsusier - r-ncvreg @@ -97,6 +98,7 @@ requirements: - r-magrittr - r-mashr - r-matrixstats + - r-mr.ash.alpha - r-mr.mashr - r-mvsusier - r-ncvreg diff --git a/NAMESPACE b/NAMESPACE index ec610e43..ceb82386 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -32,6 +32,7 @@ export(dpr_weights) export(enet_weights) export(enforce_design_full_rank) export(ensemble_weights) +export(estimate_sparsity) export(extract_cs_info) export(extract_flatten_sumstats_from_nested) export(extract_top_pip_info) diff --git a/R/regularized_regression.R b/R/regularized_regression.R index 0d86c658..9d022993 100644 --- a/R/regularized_regression.R +++ b/R/regularized_regression.R @@ -260,7 +260,7 @@ sdpr_weights <- function(stat, LD, ...) { # @param ... Additional arguments forwarded to susie_wrapper. #' @importFrom susieR coef.susie #' @noRd -.susie_extract_weights <- function(fit, X, y, required_fields, fit_args = list(), ...) { +.susie_extract_weights <- function(fit, X, y, required_fields, fit_args = list(), retain_fit = FALSE, ...) { if (is.null(fit)) { fit <- do.call(susie_wrapper, c(list(X = X, y = y), fit_args, list(...))) } @@ -272,30 +272,35 @@ sdpr_weights <- function(stat, LD, ...) { } if (all(required_fields %in% names(fit))) { fit$intercept <- 0 - return(coef.susie(fit)[-1]) + weights <- coef.susie(fit)[-1] } else { - return(rep(0, length(fit$pip))) + weights <- rep(0, length(fit$pip)) } + if (retain_fit) attr(weights, "fit") <- fit + return(weights) } #' @export -susie_weights <- function(X = NULL, y = NULL, susie_fit = NULL, ...) { +susie_weights <- function(X = NULL, y = NULL, susie_fit = NULL, retain_fit = FALSE, ...) { .susie_extract_weights(susie_fit, X, y, - required_fields = c("alpha", "mu", "X_column_scale_factors"), ...) + required_fields = c("alpha", "mu", "X_column_scale_factors"), + retain_fit = retain_fit, ...) } #' @export -susie_ash_weights <- function(X = NULL, y = NULL, susie_ash_fit = NULL, ...) { +susie_ash_weights <- function(X = NULL, y = NULL, susie_ash_fit = NULL, retain_fit = FALSE, ...) { .susie_extract_weights(susie_ash_fit, X, y, required_fields = c("alpha", "mu", "theta", "X_column_scale_factors"), - fit_args = list(unmappable_effects = "ash", convergence_method = "pip"), ...) + fit_args = list(unmappable_effects = "ash", convergence_method = "pip"), + retain_fit = retain_fit, ...) } #' @export -susie_inf_weights <- function(X = NULL, y = NULL, susie_inf_fit = NULL, ...) { +susie_inf_weights <- function(X = NULL, y = NULL, susie_inf_fit = NULL, retain_fit = FALSE, ...) { .susie_extract_weights(susie_inf_fit, X, y, required_fields = c("alpha", "mu", "theta", "X_column_scale_factors"), - fit_args = list(unmappable_effects = "inf", convergence_method = "pip"), ...) + fit_args = list(unmappable_effects = "inf", convergence_method = "pip"), + retain_fit = retain_fit, ...) } #' @export @@ -396,7 +401,7 @@ lasso_weights <- function(X, y) glmnet_weights(X, y, 1) #' @importFrom susieR mr.ash #' @importFrom stats predict #' @export -mrash_weights <- function(X, y, init_prior_sd = TRUE, ...) { +mrash_weights <- function(X, y, init_prior_sd = TRUE, retain_fit = FALSE, ...) { eff.wgt <- rep(0, ncol(X)) keep <- .drop_zero_variance(X, "mrash_weights") X_keep <- X[, keep, drop = FALSE] @@ -408,6 +413,7 @@ mrash_weights <- function(X, y, init_prior_sd = TRUE, ...) { } fit.mr.ash <- do.call(mr.ash, c(list(X = X_keep, y = y, sa2 = if (init_prior_sd) init_prior_sd(X_keep, y)^2 else NULL), args_list)) eff.wgt[keep] <- predict(fit.mr.ash, type = "coefficients")[-1] + if (retain_fit) attr(eff.wgt, "fit") <- fit.mr.ash return(eff.wgt) } #' Extract Coefficients From Bayesian Linear Regression @@ -435,7 +441,7 @@ mrash_weights <- function(X, y, init_prior_sd = TRUE, ...) { #' bayes_l_weights(y = y, X = X, Z = Z) #' bayes_r_weights(y = y, X = X, Z = Z) #' @export -bayes_alphabet_weights <- function(X, y, method, Z = NULL, nit = 5000, nburn = 1000, nthin = 5, ...) { +bayes_alphabet_weights <- function(X, y, method, Z = NULL, h2 = NULL, nit = 5000, nburn = 1000, nthin = 5, ...) { # Make sure qgg is installed if (!requireNamespace("qgg", quietly = TRUE)) { stop("To use this function, please install qgg: https://cran.r-project.org/web/packages/qgg/index.html") @@ -459,6 +465,7 @@ bayes_alphabet_weights <- function(X, y, method, Z = NULL, nit = 5000, nburn = 1 W = X[, keep, drop = FALSE], X = Z, method = method, + h2 = h2, nit = nit, nburn = nburn, ... @@ -484,8 +491,8 @@ bayes_a_weights <- function(X, y, Z = NULL, ...) { } #' Use a rounded spike prior (low-variance Gaussian). #' @export -bayes_c_weights <- function(X, y, Z = NULL, ...) { - return(bayes_alphabet_weights(X, y, method = "bayesC", Z, ...)) +bayes_c_weights <- function(X, y, Z = NULL, pi = 0.1, ...) { + return(bayes_alphabet_weights(X, y, method = "bayesC", Z, pi = pi, ...)) } #' Use a hierarchical Bayesian mixture model with four Gaussian components. Variances are scaled #' by 0, 0.0001 , 0.001 , and 0.01 . @@ -1137,11 +1144,11 @@ bglr_weights <- function(X, y, model, nIter, burnIn, thin, eta_args = list(), .. #' @param nIter Number of MCMC iterations. Default is 10000. #' @param burnIn Number of burn-in iterations. Default is 2000. #' @param thin Thinning interval. Default is 5. -#' @param probIn Prior inclusion probability for each marker. Default is 0.05. +#' @param probIn Prior inclusion probability for each marker. Default is 0.2. #' @param ... Additional arguments passed through to `BGLR::BGLR`. #' @return A numeric vector of length `ncol(X)` of variant weights. #' @export -bayes_b_weights <- function(X, y, nIter = 10000, burnIn = 2000, thin = 5, probIn = 0.05, ...) { +bayes_b_weights <- function(X, y, nIter = 10000, burnIn = 2000, thin = 5, probIn = 0.2, ...) { bglr_weights( X, y, model = "BayesB", nIter = nIter, burnIn = burnIn, thin = thin, @@ -1194,7 +1201,7 @@ b_lasso_weights <- function(X, y, nIter = 10000, burnIn = 2000, thin = 5, ...) { #' @param ... Additional arguments passed through to `RcppDPR::fit_model`. #' @return A numeric vector of length `ncol(X)` of variant weights. #' @export -dpr_weights <- function(X, y, fitting_method = "VB", ...) { +dpr_weights <- function(X, y, fitting_method = "VB", retain_fit = FALSE, ...) { if (!requireNamespace("RcppDPR", quietly = TRUE)) { stop("To use this function, please install RcppDPR: https://cran.r-project.org/package=RcppDPR") } @@ -1206,17 +1213,18 @@ dpr_weights <- function(X, y, fitting_method = "VB", ...) { rotate_variables = FALSE, fitting_method = fitting_method, ... ) eff.wgt[keep] <- as.numeric(fit$beta + fit$alpha) + if (retain_fit) attr(eff.wgt, "fit") <- fit return(eff.wgt) } #' @rdname dpr_weights #' @export -dpr_vb_weights <- function(X, y, ...) dpr_weights(X, y, fitting_method = "VB", ...) +dpr_vb_weights <- function(X, y, n_k = 8, retain_fit = FALSE, ...) dpr_weights(X, y, fitting_method = "VB", n_k = n_k, retain_fit = retain_fit, ...) #' @rdname dpr_weights #' @export -dpr_gibbs_weights <- function(X, y, ...) dpr_weights(X, y, fitting_method = "Gibbs", ...) +dpr_gibbs_weights <- function(X, y, s_step = 5000, retain_fit = FALSE, ...) dpr_weights(X, y, fitting_method = "Gibbs", s_step = s_step, retain_fit = retain_fit, ...) #' @rdname dpr_weights #' @export -dpr_adaptive_gibbs_weights <- function(X, y, ...) dpr_weights(X, y, fitting_method = "Adaptive_Gibbs", ...) +dpr_adaptive_gibbs_weights <- function(X, y, retain_fit = FALSE, ...) dpr_weights(X, y, fitting_method = "Adaptive_Gibbs", retain_fit = retain_fit, ...) diff --git a/R/twas_weights.R b/R/twas_weights.R index 336d6aef..6718fc24 100644 --- a/R/twas_weights.R +++ b/R/twas_weights.R @@ -406,7 +406,7 @@ twas_weights_cv <- function(X, Y, fold = NULL, sample_partitions = NULL, weight_ #' @importFrom furrr future_map furrr_options #' @importFrom purrr map exec #' @importFrom rlang !!! -twas_weights <- function(X, Y, weight_methods, num_threads = 1) { +twas_weights <- function(X, Y, weight_methods, num_threads = 1, retain_fits = FALSE) { if (!is.matrix(X) || (!is.matrix(Y) && !is.vector(Y))) { stop("X must be a matrix and Y must be a matrix or a vector.") } @@ -432,13 +432,20 @@ twas_weights <- function(X, Y, weight_methods, num_threads = 1) { multivariate_weight_methods <- c("mrmash_weights", "mvsusie_weights") args <- weight_methods[[method_name]] + # Only pass retain_fit to functions that accept it + if (retain_fits && "retain_fit" %in% names(formals(method_name))) { + args$retain_fit <- TRUE + } + # Remove columns with zero variance valid_columns <- .nonzero_var_columns(X) X_filtered <- as.matrix(X[, valid_columns, drop = FALSE]) + method_fit <- NULL if (method_name %in% multivariate_weight_methods) { # Apply multivariate method weights_matrix <- do.call(method_name, c(list(X = X_filtered, Y = Y), args)) + if (retain_fits) method_fit <- attr(weights_matrix, "fit") if (nrow(weights_matrix) != length(valid_columns)) weights_matrix <- weights_matrix[names(valid_columns), , drop = FALSE] } else { # Apply univariate method to each column of Y @@ -447,12 +454,17 @@ twas_weights <- function(X, Y, weight_methods, num_threads = 1) { for (k in 1:ncol(Y)) { weights_vector <- do.call(method_name, c(list(X = X_filtered, y = Y[, k]), args)) + if (retain_fits && is.null(method_fit)) { + method_fit <- attr(weights_vector, "fit") + } if (is.matrix(weights_vector)) weights_vector <- weights_vector[, k] weights_matrix[, k] <- weights_vector } } - return(.embed_weights(weights_matrix, valid_columns, ncol(X), ncol(Y), colnames(X), colnames(Y))) + result <- .embed_weights(weights_matrix, valid_columns, ncol(X), ncol(Y), colnames(X), colnames(Y)) + if (!is.null(method_fit)) attr(result, "fit") <- method_fit + return(result) } if (num_cores >= 2) { @@ -466,7 +478,9 @@ twas_weights <- function(X, Y, weight_methods, num_threads = 1) { if (!is.null(colnames(X))) { weights_list <- lapply(weights_list, function(x) { + fit <- attr(x, "fit") rownames(x) <- colnames(X) + if (!is.null(fit)) attr(x, "fit") <- fit return(x) }) } @@ -500,6 +514,56 @@ twas_predict <- function(X, weights_list) { setNames(lapply(weights_list, function(w) X %*% w), gsub("_weights", "_predicted", names(weights_list))) } +# Short method names allowed for sparsity estimation (non-full-Bayesian methods). +# @noRd +.sparsity_estimation_methods <- "mrash" + +# Map short method names to weight function names for sparsity estimation. +# @noRd +.sparsity_method_fn_map <- c(mrash = "mrash_weights") + +#' Estimate Sparsity from mr.ash Mixture Proportions +#' +#' Computes an empirical estimate of the proportion of non-zero effects +#' (sparsity) from the mr.ash fit. mr.ash fits a mixture model with a +#' point mass at zero (spike) plus continuous components (slab), and +#' learns the mixture proportions via variational EM. The sparsity +#' estimate \code{1 - pi[1]} is the empirical Bayes estimate of the +#' non-null proportion, which can be used as a data-driven prior for +#' the inclusion probability parameters (\code{pi} for bayesC, +#' \code{probIn} for BayesB) of spike-and-slab Bayesian methods. +#' +#' @param weight_results Named list of weight vectors or matrices as +#' returned by \code{\link{twas_weights}}. The mr.ash element should +#' have a \code{"fit"} attribute containing the model fit object +#' (set \code{retain_fits = TRUE} in \code{twas_weights} to obtain this). +#' @param methods Must be \code{"mrash"}. +#' +#' @return A scalar sparsity estimate (proportion of non-zero effects). +#' @export +estimate_sparsity <- function(weight_results, methods = "mrash") { + invalid <- setdiff(methods, .sparsity_estimation_methods) + if (length(invalid) > 0) { + stop("Invalid sparsity estimation method(s): ", paste(invalid, collapse = ", "), + ". Allowed: ", paste(.sparsity_estimation_methods, collapse = ", ")) + } + + fn_name <- .sparsity_method_fn_map["mrash"] + w <- weight_results[[fn_name]] + if (is.null(w)) { + stop("mr.ash weights ('mrash_weights') not found in weight_results.") + } + + fit <- attr(w, "fit") + if (is.null(fit) || is.null(fit$pi)) { + stop("mr.ash fit object not found. Run twas_weights() with retain_fits = TRUE ", + "and ensure mrash_weights is included.") + } + + # fit$pi[1] is the weight on the spike (sa2[1] = 0); 1 - pi[1] = non-null proportion + return(1 - fit$pi[1]) +} + #' TWAS Weights Pipeline #' #' This function performs weights computation for Transcriptome-Wide Association Study (TWAS) @@ -527,6 +591,13 @@ twas_predict <- function(X, weights_list) { #' \code{\link{ensemble_weights}}. Defaults to \code{"quadprog"}. #' @param ensemble_alpha Elastic net mixing parameter, used only when #' \code{ensemble_solver = "glmnet"}. Defaults to 1 (lasso). +#' @param estimate_pi_methods Character vector of short method names to +#' use for empirical sparsity estimation. The estimated sparsity is +#' used as the \code{pi} prior for bayesC and the \code{probIn} prior +#' for BayesB, unless the user provides explicit values in +#' \code{weight_methods}. Set to \code{NULL} to disable empirical +#' estimation. Defaults to \code{c("mrash")}. See +#' \code{\link{estimate_sparsity}} for allowed method names. #' #' @return A list containing results from the TWAS pipeline, including TWAS weights, predictions, and optionally cross-validation results. #' @export @@ -546,7 +617,8 @@ twas_weights_pipeline <- function(X, ensemble = FALSE, ensemble_r2_threshold = 0.01, ensemble_solver = "quadprog", - ensemble_alpha = 1) { + ensemble_alpha = 1, + estimate_pi_methods = c("mrash")) { if (is.character(weight_methods)) { weight_methods <- .twas_method_lookup(weight_methods) } @@ -555,16 +627,81 @@ twas_weights_pipeline <- function(X, st <- proc.time() message("Performing TWAS weights computation for univariate analysis methods ...") - # TWAS weights and predictions + # Store susie intermediate info if susie_fit is provided if (!is.null(susie_fit) && !is.null(weight_methods$susie_weights)) { - weight_methods$susie_weights <- list(susie_fit = susie_fit) res$susie_weights_intermediate <- susie_fit[c("mu", "lbf_variable", "X_column_scale_factors", "pip")] if (!is.null(susie_fit$sets$cs)) { res$susie_weights_intermediate$cs_variants <- setNames(lapply(susie_fit$sets$cs, function(L) colnames(X)[L]), names(susie_fit$sets$cs)) res$susie_weights_intermediate$cs_purity <- susie_fit$sets$purity } } - res$twas_weights <- twas_weights(X, y, weight_methods = weight_methods) + + # Check if empirical pi estimation is needed for spike-and-slab methods + bayes_c_needs_pi <- "bayes_c_weights" %in% names(weight_methods) && + !"pi" %in% names(weight_methods$bayes_c_weights) + bayes_b_needs_pi <- "bayes_b_weights" %in% names(weight_methods) && + !"probIn" %in% names(weight_methods$bayes_b_weights) + needs_pi_estimation <- (bayes_c_needs_pi || bayes_b_needs_pi) && !is.null(estimate_pi_methods) + + if (needs_pi_estimation) { + # Identify which estimation methods are already in the pipeline + est_fn_names <- .sparsity_method_fn_map[estimate_pi_methods] + phase1_methods <- weight_methods[names(weight_methods) %in% est_fn_names] + + # Add any requested estimation methods not already in weight_methods + missing <- setdiff(est_fn_names, names(weight_methods)) + for (m in missing) phase1_methods[[m]] <- list() + + # Pass susie_fit to phase 1 if SuSiE is an estimation method + if (!is.null(susie_fit) && "susie_weights" %in% names(phase1_methods)) { + phase1_methods$susie_weights <- c( + list(susie_fit = susie_fit), + weight_methods[["susie_weights"]][setdiff(names(weight_methods[["susie_weights"]]), "susie_fit")] + ) + } + + message(" Estimating sparsity from: ", paste(estimate_pi_methods, collapse = ", "), " ...") + phase1_weights <- twas_weights(X, y, weight_methods = phase1_methods, retain_fits = TRUE) + + empirical_pi <- estimate_sparsity(phase1_weights, methods = estimate_pi_methods) + message(sprintf(" Empirical sparsity estimate: %.4f", empirical_pi)) + res$empirical_pi <- empirical_pi + + # Inject into spike-and-slab methods that need it + if (bayes_c_needs_pi) weight_methods$bayes_c_weights$pi <- as.numeric(empirical_pi) + if (bayes_b_needs_pi) weight_methods$bayes_b_weights$probIn <- as.numeric(empirical_pi) + + # Phase 2: run remaining methods (those not already computed in phase 1) + phase2_fn_names <- setdiff(names(weight_methods), names(phase1_methods)) + + # Pass susie_fit for phase 2 if SuSiE is NOT an estimation method but is in weight_methods + if (!is.null(susie_fit) && "susie_weights" %in% phase2_fn_names) { + weight_methods$susie_weights <- list(susie_fit = susie_fit) + } + + if (length(phase2_fn_names) > 0) { + phase2_methods <- weight_methods[phase2_fn_names] + phase2_weights <- twas_weights(X, y, weight_methods = phase2_methods) + res$twas_weights <- c(phase1_weights, phase2_weights) + } else { + res$twas_weights <- phase1_weights + } + + # Clean up fit attributes from final weight results + res$twas_weights <- lapply(res$twas_weights, function(w) { attr(w, "fit") <- NULL; w }) + + # Remove any estimation-only methods that were not in the original weight_methods + if (length(missing) > 0) { + res$twas_weights <- res$twas_weights[setdiff(names(res$twas_weights), missing)] + } + } else { + # Original flow: run all methods at once + if (!is.null(susie_fit) && !is.null(weight_methods$susie_weights)) { + weight_methods$susie_weights <- list(susie_fit = susie_fit) + } + res$twas_weights <- twas_weights(X, y, weight_methods = weight_methods) + } + res$twas_predictions <- twas_predict(X, res$twas_weights) if (cv_folds > 1) { diff --git a/inst/prototype/ensemble_twas_weights.ipynb b/inst/prototype/ensemble_twas_weights.ipynb new file mode 100644 index 00000000..9277af02 --- /dev/null +++ b/inst/prototype/ensemble_twas_weights.ipynb @@ -0,0 +1,992 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "kernel": "SoS" + }, + "source": [ + "# Ensemble TWAS Weights via Stacked Regression\n", + "\n", + "## Overview\n", + "\n", + "`pecotmr` offers 10+ TWAS weight methods (SuSiE, LASSO, Elastic Net, mr.ash, BayesR, DPR, etc.), each suited to different genetic architectures. Picking one method is arbitrary; running all and testing each incurs a multiple-testing penalty. **Stacked regression** ([SR-TWAS; Dai et al.,2024](https://doi.org/10.1038/s41467-024-50983-w)) provides a principled alternative: learn one convex combination of methods per gene, producing a single weight vector for downstream TWAS testing with no method-selection bias.\n", + "\n", + "This notebook:\n", + "1. Presents the mathematical framework and algorithm for ensemble weight learning\n", + "2. Benchmarks the approach via simulation using **genotypes** from LD sketch reference panels, so that realistic LD patterns are preserved\n", + "3. Produces diagnostic plots comparing MSE against ground truth and MSE against observed $Y$\n", + "\n", + "| Step | Description | Parallelism |\n", + "|------|-------------|-------------|\n", + "| `simulate_and_fit` | Load X from sketch LD, simulate Y, run methods + ensemble, save RDS | Per-region (`for_each`) |\n", + "| `analyze_results` | Load RDS files, compute metrics, generate diagnostic plots | Single task |" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "kernel": "SoS" + }, + "source": [ + "## Mathematical Framework\n", + "\n", + "### Gene Expression Model\n", + "\n", + "For a target gene $g$ with $p$ cis-SNPs, the expression level of individual $i$ is:\n", + "\n", + "$$E_i = \\mathbf{G}_i \\mathbf{w} + \\epsilon_i, \\quad \\epsilon_i \\sim N(0, \\sigma^2_\\epsilon)$$\n", + "\n", + "where $\\mathbf{G}_i \\in \\mathbb{R}^{1 \\times p}$ is the genotype vector and $\\mathbf{w} \\in \\mathbb{R}^p$ is the true cis-eQTL effect-size vector.\n", + "\n", + "### Base Models\n", + "\n", + "We train $K$ prediction models $\\{m_1, \\ldots, m_K\\}$, each producing a weight estimate $\\hat{\\mathbf{w}}_k \\in \\mathbb{R}^p$. Under $F$-fold cross-validation, the out-of-fold prediction for sample $i$ from method $k$ is:\n", + "\n", + "$$\\hat{E}_i^{(k)} = \\mathbf{G}_i \\hat{\\mathbf{w}}_k^{(-f_i)}$$\n", + "\n", + "where $f_i$ denotes the fold containing sample $i$ and $\\hat{\\mathbf{w}}_k^{(-f)}$ is trained on all samples except fold $f$. Stacking all $n$ out-of-fold predictions across $K$ methods yields the prediction matrix:\n", + "\n", + "$$\\mathbf{P} = \\begin{pmatrix} \\hat{E}_1^{(1)} & \\cdots & \\hat{E}_1^{(K)} \\\\ \\vdots & \\ddots & \\vdots \\\\ \\hat{E}_n^{(1)} & \\cdots & \\hat{E}_n^{(K)} \\end{pmatrix} \\in \\mathbb{R}^{n \\times K}$$\n", + "\n", + "### Stacked Regression (Constrained QP)\n", + "\n", + "We seek ensemble coefficients $\\boldsymbol{\\zeta} = (\\zeta_1, \\ldots, \\zeta_K)^\\top$ that minimize the prediction error of the convex combination:\n", + "\n", + "$$\\min_{\\boldsymbol{\\zeta}} \\;\\; \\left\\| \\mathbf{E} - \\sum_{k=1}^{K} \\zeta_k \\hat{\\mathbf{E}}^{(k)} \\right\\|^2 = \\left\\| \\mathbf{E} - \\mathbf{P} \\boldsymbol{\\zeta} \\right\\|^2$$\n", + "\n", + "$$\\text{subject to} \\quad \\zeta_k \\geq 0 \\;\\; \\forall k, \\qquad \\sum_{k=1}^{K} \\zeta_k = 1$$\n", + "\n", + "This is a convex quadratic program solved via `quadprog::solve.QP` with:\n", + "\n", + "$$\\mathbf{D} = \\mathbf{P}^\\top \\mathbf{P} + \\lambda \\mathbf{I}, \\quad \\mathbf{d} = \\mathbf{P}^\\top \\mathbf{E}$$\n", + "\n", + "where $\\lambda = 10^{-8} \\cdot \\text{mean}(\\text{diag}(\\mathbf{P}^\\top \\mathbf{P}))$ provides ridge regularization for numerical stability.\n", + "\n", + "### Ensemble Weight Vector\n", + "\n", + "The final ensemble TWAS weight vector combines the full-data weight estimates:\n", + "\n", + "$$\\tilde{\\mathbf{w}} = \\sum_{k=1}^{K} \\zeta_k^* \\hat{\\mathbf{w}}_k$$\n", + "\n", + "### Performance Metrics\n", + "\n", + "Two complementary evaluation criteria:\n", + "\n", + "- **Against ground truth** (oracle): $\\text{MSE}_{\\text{truth}} = \\frac{1}{n} \\| \\mathbf{G}\\hat{\\mathbf{w}} - \\mathbf{G}\\mathbf{w}_{\\text{true}} \\|^2$\n", + "- **Against observed $Y$** (practical): $R^2_Y = \\text{cor}(Y,\\; \\mathbf{G}\\hat{\\mathbf{w}})^2$ and $\\text{MSE}_Y = \\frac{1}{n}\\|Y - \\mathbf{G}\\hat{\\mathbf{w}}\\|^2$, where $Y = \\mathbf{G}\\mathbf{w}_{\\text{true}} + \\epsilon$" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "kernel": "SoS" + }, + "source": [ + "## Algorithm\n", + "\n", + "The algorithm below shows the existing pecotmr cross-validation framework (Steps 1--5) and the **new ensemble learning extension** (Steps 6--9, in bold). Step numbers reference the mathematical formulation above.\n", + "\n", + "```\n", + "Algorithm: TWAS Weight Estimation with Ensemble Learning\n", + "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n", + "\n", + "Input: Genotype matrix G (n x p), expression vector E (n x 1),\n", + " weight methods M = {m_1, ..., m_K}, number of CV folds F\n", + "Output: Per-method weights {w_k}, ensemble coefficients {zeta_k},\n", + " combined weight vector w_tilde, CV performance metrics\n", + "\n", + "── Existing Framework: twas_weights_cv() + twas_weights() ──────\n", + "\n", + "1. Partition n samples into F folds: {S_1, ..., S_F}\n", + "2. For each fold f = 1, ..., F:\n", + " For each method k = 1, ..., K:\n", + " a. Train: w_k^(-f) <- m_k(G[T_f,], E[T_f])\n", + " b. Predict: E_hat^(k)[S_f] <- G[S_f,] * w_k^(-f)\n", + "3. Assemble out-of-fold prediction matrix P (n x K)\n", + " P[i,k] = E_hat_i^(k) from the fold containing sample i\n", + "4. Compute per-method performance:\n", + " R^2_k = cor(E, P[,k])^2 for k = 1, ..., K\n", + "5. Train final weights on full data:\n", + " w_k <- m_k(G, E) for all k\n", + "\n", + "── NEW: Ensemble Learning via ensemble_weights() ──────────────\n", + "\n", + " 6. Solve constrained QP for ensemble coefficients:\n", + " D <- P^T P + lambda * I (ridge regularization)\n", + " d <- P^T E\n", + " zeta* <- argmin 1/2 zeta^T D zeta - d^T zeta\n", + " s.t. sum(zeta_k) = 1, zeta_k >= 0 for all k\n", + " 7. Compute ensemble prediction:\n", + " E_hat_ens = P zeta*\n", + " 8. Compute ensemble performance:\n", + " R^2_ens = cor(E, E_hat_ens)^2\n", + " 9. Combine final weight vectors:\n", + " w_tilde = sum_k zeta_k * w_k\n", + "\n", + "Return: {w_k}, {zeta_k}, w_tilde, {R^2_k}, R^2_ens\n", + "```\n", + "\n", + "**Implementation**: Steps 1--5 are handled internally by `pecotmr::twas_weights_pipeline()`. Steps 6--9 are `pecotmr::ensemble_weights()`, now exported in PR #486 (no local source file needed)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "kernel": "SoS" + }, + "source": [ + "## Weight Methods\n", + "\n", + "Two method sets are provided depending on computational budget:\n", + "\n", + "### Expensive (default)\n", + "\n", + "| # | Method | Type | pecotmr function | Key parameters |\n", + "|---|--------|------|------------------|----------------|\n", + "| 1 | SuSiE | Bayesian sparse (L0-like) | `susie_weights` | `refine=FALSE, init_L=5, max_L=10` |\n", + "| 2 | SuSiE-ash | Bayesian sparse + ash shrinkage | `susie_ash_weights` | default |\n", + "| 3 | SuSiE-inf | Bayesian sparse + infinitesimal | `susie_inf_weights` | default |\n", + "| 4 | mr.ash | Bayesian adaptive shrinkage | `mrash_weights` | `init_prior_sd=TRUE, max.iter=100` |\n", + "| 5 | Elastic Net | L1+L2 penalized (glmnet) | `enet_weights` | default (`alpha=0.5`) |\n", + "| 6 | LASSO | L1 penalized (glmnet) | `lasso_weights` | default |\n", + "| 7 | BayesR | Bayesian mixture of normals | `bayes_r_weights` | default (`qgg`) |\n", + "| 8 | BayesL | Bayesian LASSO | `bayes_l_weights` | default |\n", + "| 9 | BayesA | Bayesian ridge-like | `bayes_a_weights` | default |\n", + "| 10 | BayesB | Bayesian variable selection | `bayes_b_weights` | default |\n", + "| 11 | BayesC | Bayesian spike-and-slab | `bayes_c_weights` | default |\n", + "| 12 | BayesN | Bayesian normal mixture | `bayes_n_weights` | default |\n", + "| 13 | B-LASSO | Bayesian LASSO variant | `b_lasso_weights` | default |\n", + "| 14 | DPR-VB | Dirichlet process regression (VB) | `dpr_vb_weights` | default (`RcppDPR`) |\n", + "| 15 | DPR-Gibbs | Dirichlet process regression (Gibbs) | `dpr_gibbs_weights` | default |\n", + "| 16 | DPR-Adaptive Gibbs | Dirichlet process regression (adaptive) | `dpr_adaptive_gibbs_weights` | default |\n", + "| 17 | SCAD | SCAD penalized | `scad_weights` | default |\n", + "| 18 | MCP | MCP penalized | `mcp_weights` | default |\n", + "| 19 | L0Learn | Best-subset approximation | `l0learn_weights` | default |\n", + "\n", + "### Cheap (fast iteration)\n", + "\n", + "| # | Method | Type | pecotmr function |\n", + "|---|--------|------|------------------|\n", + "| 1 | SuSiE | Bayesian sparse | `susie_weights` |\n", + "| 2 | mr.ash | Bayesian adaptive shrinkage | `mrash_weights` |\n", + "| 3 | LASSO | L1 penalized | `lasso_weights` |\n", + "| 4 | Elastic Net | L1+L2 penalized | `enet_weights` |\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "kernel": "SoS" + }, + "source": [ + "## Simulation Design\n", + "\n", + "### Genotypes\n", + "\n", + "Each replicate uses **genotypes** loaded from LD sketch reference panels via `pecotmr::load_LD_matrix(meta, region, return_genotype = TRUE)`. Preserving LD patterns is critical for faithfully benchmarking TWAS weight methods — no synthetic genotype simulation can reproduce the complex correlation structure of actual genomic data.\n", + "\n", + "The pipeline requires two input files:\n", + "\n", + "1. **`ld_meta_file`** — Sketch LD metadata TSV pointing to plink2 files (one row per chromosome, `start=0, end=0` sentinel). Generated by the `rss_ld_sketch` pipeline. Example:\n", + " ```\n", + " #chrom start end path\n", + " 22 0 0 ADSP.R5.EUR.chr22\n", + " ```\n", + " The `path` column is the plink2 file prefix (resolved relative to the TSV directory). The corresponding `.pgen`, `.pvar`, `.psam`, `.afreq` files must exist.\n", + "\n", + "2. **`ld_block_file`** — BED or TSV file defining LD block boundaries (one row per block). Each block becomes a replicate region. Example (BED format, 0-based half-open):\n", + " ```\n", + " chr22 16051249 17614263\n", + " chr22 17614263 18834263\n", + " ...\n", + " ```\n", + " Or TSV with `#chrom start end` header. These define the genomic windows passed to `load_LD_matrix()` to extract variants per region.\n", + "\n", + "### Phenotypes\n", + "\n", + "- **Sparse**: `n_sparse` variants with large effects\n", + "- **Oligogenic**: `n_oligogenic` variants with moderate effects\n", + "- **Infinitesimal**: `n_inf` variants with small effects\n", + "- Total heritability: $h^2_g$ (default 0.2)\n", + "\n", + "### Replicates\n", + "\n", + "Each genomic region in the `ld_block_file` defines one replicate. The `n_seeds` parameter controls the number of independent phenotype realizations per region (different random seeds produce different causal architectures on the same genotype matrix).\n", + "\n", + "### Data Saving\n", + "\n", + "Set `save_data = TRUE` to save the raw genotype matrix, simulated phenotype, and ground-truth causal information in a separate `data.rds` file per replicate. This is useful for debugging but produces large files (one per region x seed). Disabled by default.\n", + "\n", + "### Evaluation Metrics\n", + "\n", + "Two complementary metrics per replicate:\n", + "\n", + "1. **Against ground truth** — MSE of predicted genetic component vs. true genetic component: $\\frac{1}{n}\\|\\mathbf{G}\\hat{\\mathbf{w}} - \\mathbf{G}\\mathbf{w}_{\\text{true}}\\|^2$. This measures how well the method recovers the true signal.\n", + "\n", + "2. **Against observed $Y$** — $R^2$ and MSE of predictions vs. observed phenotype $Y = \\mathbf{G}\\mathbf{w}_{\\text{true}} + \\epsilon$. This is the practical metric: how well would the weights predict expression in a held-out sample?" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "kernel": "SoS" + }, + "source": [ + "## Usage\n", + "\n", + "### Quick test (cheap methods, few regions)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "kernel": "SoS" + }, + "outputs": [], + "source": [ + "sos run pipeline/ensemble_twas_weights.ipynb simulate_and_fit \\\n", + " --ld-meta-file data/ld_meta_file.tsv \\\n", + " --ld-block-file /restricted/projectnb/casa/oaolayin/ROSMAP_NIA_geno/EUR_LD_blocks.bed \\\n", + " --chrom 21 \\\n", + " --n-seeds 1 \\\n", + " --method-set cheap \\\n", + " -J 20" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "kernel": "SoS" + }, + "source": [ + "### Full benchmark (expensive methods, 10 seeds per region)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "kernel": "SoS" + }, + "outputs": [], + "source": [ + "sos run pipeline/ensemble_twas_weights.ipynb simulate_and_fit \\\n", + " --ld-meta-file /path/to/sketch_ld_meta.tsv \\\n", + " --ld-block-file /path/to/EUR_LD_blocks.bed \\\n", + " --chrom 22 \\\n", + " --n-seeds 10 \\\n", + " --method-set expensive \\\n", + " -J 40" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "kernel": "SoS" + }, + "source": [ + "### With data saving (for debugging)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "kernel": "SoS" + }, + "outputs": [], + "source": [ + "sos run pipeline/ensemble_twas_weights.ipynb simulate_and_fit \\\n", + " --ld-meta-file /path/to/sketch_ld_meta.tsv \\\n", + " --ld-block-file /path/to/EUR_LD_blocks.bed \\\n", + " --chrom 22 \\\n", + " --n-seeds 1 \\\n", + " --save-data TRUE \\\n", + " -J 20" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "kernel": "SoS" + }, + "source": [ + "\n", + "### Analysis and plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jp-MarkdownHeadingCollapsed": true, + "kernel": "SoS" + }, + "outputs": [], + "source": [ + "sos run pipeline/ensemble_twas_weights.ipynb analyze_results \\\n", + " --cwd /restricted/projectnb/xqtl/jaempawi/xqtl_protocol/toy_example/output/ensemble_twas \\\n", + " --ld-meta-file /restricted/projectnb/xqtl/jaempawi/xqtl_protocol/toy_example/output/rss_ld_sketch/sketch_ld_meta.tsv \\\n", + " --ld-block-file /restricted/projectnb/casa/oaolayin/ROSMAP_NIA_geno/EUR_LD_blocks.bed" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "kernel": "SoS" + }, + "source": [ + "# Pipeline Implementation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "kernel": "SoS" + }, + "outputs": [], + "source": [ + "[global]\n", + "# Output directory\n", + "parameter: cwd = path(\"output/ensemble_twas\")\n", + "\n", + "\n", + "# ── Genotype source ──────────────────────────────────────────\n", + "# Sketch LD metadata TSV (plink2 format: chrom, start=0, end=0, path=prefix)\n", + "# Generated by the rss_ld_sketch pipeline. Used by load_LD_matrix() to load genotypes.\n", + "parameter: ld_meta_file = path(\"\")\n", + "# LD block file defining replicate regions (BED or TSV with chrom/start/end)\n", + "# Each block becomes one replicate. Can be EUR_LD_blocks.bed or similar.\n", + "parameter: ld_block_file = path(\"\")\n", + "# Chromosome filter (0 = all chromosomes)\n", + "parameter: chrom = 0\n", + "# Max variants to subsample per region (0 = no limit)\n", + "parameter: max_variants = 0\n", + "\n", + "# ── Method selection ─────────────────────────────────────────\n", + "# \"expensive\": full 19-method list (SuSiE variants, mr.ash, BayesR family, DPR family, LASSO, Elastic Net, SCAD, MCP, L0Learn)\n", + "# \"cheap\": \"fast_default\" preset — SuSiE, mr.ash, Elastic Net, LASSO\n", + "parameter: method_set = \"expensive\"\n", + "\n", + "# ── Simulation parameters ────────────────────────────────────\n", + "parameter: h2g = 0.2\n", + "parameter: n_sparse = 3\n", + "parameter: n_oligogenic = 5\n", + "parameter: n_inf = 15\n", + "parameter: cv_folds = 5\n", + "parameter: max_cv_variants = 500\n", + "\n", + "# ── Replicate control ────────────────────────────────────────\n", + "parameter: n_seeds = 1\n", + "parameter: seed_base = 42\n", + "\n", + "# ── Data saving (for debugging) ──────────────────────────────\n", + "# When TRUE, saves raw genotype matrix X, phenotype y, and sim object\n", + "# as a separate data.rds per replicate. Produces large files.\n", + "parameter: save_data = False\n", + "\n", + "# ── Cluster resources ────────────────────────────────────────\n", + "parameter: job_size = 1\n", + "parameter: walltime = \"4:00:00\"\n", + "parameter: mem = \"16G\"\n", + "parameter: numThreads = 4\n", + "\n", + "import os\n", + "cwd = path(f'{cwd:a}')\n", + "\n", + "def _read_ld_blocks(block_file, chrom_filter):\n", + " \"\"\"Parse LD block file (BED or TSV with chrom/start/end columns).\n", + " Returns list of region dicts with id, chrom, start, end, region string.\n", + " Supports BED (0-based half-open) and TSV with #chrom header.\n", + " \"\"\"\n", + " regions = []\n", + " with open(block_file) as fh:\n", + " for line in fh:\n", + " line = line.strip()\n", + " if not line or line.startswith('#chrom') or line.startswith('chrom'):\n", + " continue\n", + " parts = line.split()\n", + " if len(parts) < 3:\n", + " continue\n", + " c = parts[0]\n", + " if not c.startswith('chr'):\n", + " c = f'chr{c}'\n", + " try:\n", + " cnum = int(c.replace('chr', ''))\n", + " except ValueError:\n", + " continue\n", + " if not (1 <= cnum <= 22):\n", + " continue\n", + " if chrom_filter != 0 and cnum != chrom_filter:\n", + " continue\n", + " start, end = int(parts[1]), int(parts[2])\n", + " regions.append({\n", + " 'id': f'{c}_{start}_{end}',\n", + " 'chrom': c,\n", + " 'start': start,\n", + " 'end': end,\n", + " 'region': f'{c}:{start}-{end}'\n", + " })\n", + " return regions\n", + "\n", + "# Validate inputs\n", + "if not str(ld_meta_file) or not os.path.isfile(str(ld_meta_file)):\n", + " raise ValueError(\"ld_meta_file is required — must point to a sketch LD metadata TSV (plink2 format)\")\n", + "if not str(ld_block_file) or not os.path.isfile(str(ld_block_file)):\n", + " raise ValueError(\"ld_block_file is required — must point to an LD block BED/TSV defining replicate regions\")\n", + "\n", + "# Build replicate list: regions x seeds\n", + "_base_regions = _read_ld_blocks(str(ld_block_file), chrom)\n", + "if not _base_regions:\n", + " raise ValueError(f\"No regions found in {ld_block_file} for chrom={chrom}\")\n", + "\n", + "replicates = []\n", + "for seed_idx in range(n_seeds):\n", + " for i, r in enumerate(_base_regions):\n", + " rep = dict(r)\n", + " rep['seed'] = seed_base + seed_idx * len(_base_regions) + i\n", + " rep['rep_id'] = f'{r[\"id\"]}_seed{seed_idx}'\n", + " replicates.append(rep)\n", + "\n", + "print(f\" {len(_base_regions)} regions x {n_seeds} seeds = {len(replicates)} replicates\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "kernel": "SoS" + }, + "source": [ + "### `simulate_and_fit`\n", + "\n", + "For each replicate (region x seed):\n", + "1. Load genotype matrix $\\mathbf{G}$ from LD sketch via `pecotmr::load_LD_matrix()` using the sketch LD metadata and the region defined by the LD block file\n", + "2. Simulate expression $E$ via `simxQTL::generate_cis_qtl_data()` with mixed architecture on top of $\\mathbf{G}$\n", + "3. Run `pecotmr::twas_weights_pipeline()` with selected methods and 5-fold CV\n", + "4. Call `pecotmr::ensemble_weights()` (now exported, PR #486) to learn $\\boldsymbol{\\zeta}$ and combine weights\n", + "5. Compute metrics against ground truth and against observed $Y$\n", + "6. Save results as RDS; optionally save raw data (X, y, sim) as separate `data.rds` when `--save-data TRUE`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "kernel": "SoS" + }, + "outputs": [], + "source": [ + "[simulate_and_fit]\n", + "input: for_each = \"replicates\"\n", + "output_files = [f'{cwd}/{_replicates[\"rep_id\"]}/results.rds']\n", + "if save_data:\n", + " output_files.append(f'{cwd}/{_replicates[\"rep_id\"]}/data.rds')\n", + "output: output_files\n", + "# Pre-compute data path in Python so the R script can reference it unconditionally.\n", + "# (Using ${_output[1]} directly fails with IndexError when save_data=False\n", + "# because SoS's ${} substitution is Python-eager, regardless of R control flow.)\n", + "_data_path = f'{cwd}/{_replicates[\"rep_id\"]}/data.rds' if save_data else ''\n", + "task: trunk_workers = 1, trunk_size = job_size, walltime = walltime, mem = mem, cores = numThreads\n", + "R: expand = \"${ }\", stderr = f'{_output[0]:n}.stderr', stdout = f'{_output[0]:n}.stdout'\n", + "\n", + " library(pecotmr)\n", + " library(simxQTL)\n", + " set.seed(${_replicates['seed']})\n", + " cat(\"Replicate: ${_replicates['rep_id']}\\n\")\n", + " cat(\"Seed: ${_replicates['seed']}\\n\")\n", + " cat(\"Region: ${_replicates['region']}\\n\")\n", + "\n", + " # ── 1. Load genotype from LD sketch ─────────────────────\n", + " ld_data <- load_LD_matrix(\n", + " \"${ld_meta_file}\",\n", + " region = \"${_replicates['region']}\",\n", + " return_genotype = TRUE\n", + " )\n", + " X <- ld_data$LD_matrix\n", + " X <- scale(X)\n", + " X[is.na(X)] <- 0\n", + " n <- nrow(X); p <- ncol(X)\n", + " cat(sprintf(\"Genotype matrix: %d samples x %d variants\\n\", n, p))\n", + "\n", + " max_var <- ${max_variants}\n", + " if (max_var > 0 && p > max_var) {\n", + " keep_idx <- sort(sample(p, max_var))\n", + " X <- X[, keep_idx, drop = FALSE]\n", + " p <- ncol(X)\n", + " cat(sprintf(\"Subsampled to %d variants\\n\", p))\n", + " }\n", + "\n", + " # ── 2. Simulate training phenotype ───────────────────────\n", + " sim <- generate_cis_qtl_data(X, h2g = ${h2g},\n", + " n_sparse = ${n_sparse},\n", + " n_oligogenic = ${n_oligogenic},\n", + " n_inf = ${n_inf})\n", + " y <- sim$y # training Y: G*w_true + eps_train\n", + " cat(sprintf(\"Realized h2g: %.4f residual var: %.4f\\n\",\n", + " sim$h2g, sim$residual_variance))\n", + "\n", + " # ── 2a. Build ground-truth pieces from sim$beta ───────────\n", + " if (!is.null(sim$beta) && length(sim$beta) == p) {\n", + " true_beta <- as.numeric(sim$beta)\n", + " } else {\n", + " warning(\"sim$beta missing or wrong length (\", length(sim$beta),\n", + " \" vs p=\", p, \"); defaulting to zero vector\")\n", + " true_beta <- rep(0, p)\n", + " }\n", + " true_genetic <- as.vector(X %*% true_beta)\n", + "\n", + " # ── 2b. Fresh held-out Y (same X, same w_true, new noise) ─\n", + " # This is what lets us evaluate prediction without in-sample\n", + " # overfitting bias — the weights never saw this Y.\n", + " sigma2 <- sim$residual_variance\n", + " eps_new <- rnorm(n, mean = 0, sd = sqrt(sigma2))\n", + " y_test <- true_genetic + eps_new\n", + " cat(\"Generated fresh Y_test (independent noise draw)\\n\")\n", + "\n", + " # ── 2c. Optional raw-data dump for debugging ─────────────\n", + " data_path <- \"${_data_path}\"\n", + " if (nzchar(data_path)) {\n", + " data_out <- list(\n", + " X = X, y = y, y_test = y_test, sim = sim,\n", + " region = \"${_replicates['region']}\",\n", + " rep_id = \"${_replicates['rep_id']}\",\n", + " seed = ${_replicates['seed']},\n", + " variant_info = ld_data$ref_panel\n", + " )\n", + " dir.create(dirname(data_path), recursive = TRUE, showWarnings = FALSE)\n", + " saveRDS(data_out, data_path)\n", + " cat(\"Saved raw data:\", data_path, \"\\n\")\n", + " rm(data_out)\n", + " }\n", + "\n", + " # ── 3. Define weight methods ─────────────────────────────\n", + " method_set <- \"${method_set}\"\n", + " if (method_set == \"expensive\") {\n", + " weight_methods <- list(\n", + " susie_weights = list(refine = FALSE, init_L = 5, max_L = 10),\n", + " susie_ash_weights = list(),\n", + " susie_inf_weights = list(),\n", + " mrash_weights = list(init_prior_sd = TRUE, max.iter = 100),\n", + " enet_weights = list(),\n", + " lasso_weights = list(),\n", + " bayes_r_weights = list(),\n", + " bayes_l_weights = list(),\n", + " bayes_a_weights = list(),\n", + " bayes_b_weights = list(),\n", + " bayes_c_weights = list(),\n", + " bayes_n_weights = list(),\n", + " b_lasso_weights = list(),\n", + " dpr_vb_weights = list(),\n", + " dpr_gibbs_weights = list(),\n", + " dpr_adaptive_gibbs_weights = list(),\n", + " scad_weights = list(),\n", + " mcp_weights = list(),\n", + " l0learn_weights = list()\n", + "\n", + " )\n", + " } else {\n", + " weight_methods <- list(\n", + " susie_weights = list(refine = FALSE, init_L = 5, max_L = 10),\n", + " mrash_weights = list(init_prior_sd = TRUE, max.iter = 100),\n", + " lasso_weights = list(),\n", + " enet_weights = list()\n", + " )\n", + " }\n", + "\n", + " # ── 4. TWAS weights pipeline (CV + full-data weights) ────\n", + " # Note: twas_weights_pipeline() in pecotmr >=0.5 (PR #486) does NOT accept num_threads.\n", + " # ensemble_weights() is now exported; no local R/ensemble_weights.R source needed.\n", + " cat(\"Running twas_weights_pipeline with\", length(weight_methods), \"methods...\\n\")\n", + " res <- tryCatch(\n", + " twas_weights_pipeline(\n", + " X, y,\n", + " cv_folds = ${cv_folds},\n", + " weight_methods = weight_methods,\n", + " max_cv_variants = ${max_cv_variants}\n", + " ),\n", + " error = function(e) {\n", + " cat(\"Pipeline error:\", conditionMessage(e), \"\\n\")\n", + " NULL\n", + " }\n", + " )\n", + "\n", + " if (is.null(res)) {\n", + " cat(\"Pipeline failed — exiting non-zero so SoS flags this replicate\\n\")\n", + " saveRDS(list(results = NULL, ensemble = NULL,\n", + " truth = list(rep_id = \"${_replicates['rep_id']}\",\n", + " seed = ${_replicates['seed']},\n", + " error = TRUE)),\n", + " \"${_output[0]}\")\n", + " quit(save = \"no\", status = 1)\n", + " }\n", + "\n", + " # ── 5. Ensemble learning (steps 6-9 in the algorithm) ────\n", + " cat(\"Computing ensemble weights...\\n\")\n", + " ens <- tryCatch(\n", + " pecotmr::ensemble_weights(\n", + " cv_results = res$twas_cv_result,\n", + " Y = y,\n", + " twas_weight_list = res$twas_weights\n", + " ),\n", + " error = function(e) {\n", + " cat(\"Ensemble error:\", conditionMessage(e), \"\\n\")\n", + " NULL\n", + " }\n", + " )\n", + "\n", + " # ── 6. Compute metrics ───────────────────────────────────\n", + " # Four metrics per method and per ensemble:\n", + " # mse_w = mean((w_hat - w_true)^2) weight-space, LD-independent\n", + " # mse_pred_truth = mean((X*w_hat - X*w_true)^2) prediction-space, LD-aware (what TWAS uses)\n", + " # mse_y = mean((X*w_hat - y_test)^2) held-out Y (fresh eps), honest out-of-sample\n", + " # rsq_y = cor(X*w_hat, y_test)^2 held-out Y R^2\n", + " compute_metrics <- function(w_hat) {\n", + " if (is.matrix(w_hat)) w_hat <- w_hat[, 1]\n", + " pred <- as.vector(X %*% w_hat)\n", + " mse_w <- mean((w_hat - true_beta)^2)\n", + " mse_pred_truth <- mean((pred - true_genetic)^2)\n", + " mse_y <- mean((pred - y_test)^2)\n", + " rsq_y <- if (sd(pred) > 0) cor(pred, y_test)^2 else 0\n", + " c(mse_w = mse_w,\n", + " mse_pred_truth = mse_pred_truth,\n", + " mse_y = mse_y,\n", + " rsq_y = rsq_y)\n", + " }\n", + "\n", + " method_metrics <- sapply(names(res$twas_weights), function(wname) {\n", + " compute_metrics(res$twas_weights[[wname]])\n", + " })\n", + " colnames(method_metrics) <- gsub(\"_weights$\", \"\", colnames(method_metrics))\n", + "\n", + " ensemble_metrics <- c(mse_w = NA, mse_pred_truth = NA, mse_y = NA, rsq_y = NA)\n", + " if (!is.null(ens) && !is.null(ens$ensemble_twas_weights)) {\n", + " ensemble_metrics <- compute_metrics(ens$ensemble_twas_weights)\n", + " }\n", + "\n", + "\n", + " # ── 7. Save results ──────────────────────────────────────\n", + " truth <- list(\n", + " causal_sparse = sim$sparse_indices,\n", + " causal_oligo = sim$oligogenic_indices,\n", + " causal_inf = sim$infinitesimal_indices,\n", + " true_beta = true_beta,\n", + " true_y = y,\n", + " y_test = y_test,\n", + " residual_variance = sigma2,\n", + " h2g = sim$h2g,\n", + " rep_id = \"${_replicates['rep_id']}\",\n", + " region = \"${_replicates['region']}\",\n", + " seed = ${_replicates['seed']},\n", + " n = n, p = p,\n", + " method_metrics = method_metrics,\n", + " ensemble_metrics = ensemble_metrics,\n", + " )\n", + "\n", + " out <- list(results = res, ensemble = ens, truth = truth)\n", + " dir.create(dirname(\"${_output[0]}\"), recursive = TRUE, showWarnings = FALSE)\n", + " saveRDS(out, \"${_output[0]}\")\n", + " cat(\"Saved:\", \"${_output[0]}\", \"\\n\")\n", + " tryCatch({\n", + " cat(\"Ensemble MSE (weights): \", round(ensemble_metrics[\"mse_w\"], 6), \"\\n\")\n", + " cat(\"Ensemble MSE (pred vs truth):\",\n", + " round(ensemble_metrics[\"mse_pred_truth\"], 6), \"\\n\")\n", + " cat(\"Ensemble MSE (held-out Y): \",\n", + " round(ensemble_metrics[\"mse_y\"], 6), \"\\n\")\n", + " cat(\"Ensemble R^2 (held-out Y): \",\n", + " round(ensemble_metrics[\"rsq_y\"], 4), \"\\n\")\n", + " }, error = function(e) cat(\"Metric-print failed (non-fatal):\", conditionMessage(e), \"\\n\"))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "kernel": "SoS" + }, + "source": [ + "### `analyze_results`\n", + "\n", + "Load all per-replicate RDS files and produce:\n", + "1. **MSE of weights** — $\\frac{1}{p}\\|\\hat{\\mathbf{w}} - \\mathbf{w}_{\\text{true}}\\|^2$ (LD-independent)\n", + "2. **MSE of predicted genetic component** — $\\frac{1}{n}\\|\\mathbf{G}\\hat{\\mathbf{w}} - \\mathbf{G}\\mathbf{w}_{\\text{true}}\\|^2$ (LD-aware, TWAS-relevant)\n", + "3. **MSE vs held-out $Y$** — $\\frac{1}{n}\\|Y_{\\text{test}} - \\mathbf{G}\\hat{\\mathbf{w}}\\|^2$, where $Y_{\\text{test}} = \\mathbf{G}\\mathbf{w}_{\\text{true}} + \\varepsilon_{\\text{new}}$\n", + "4. **$R^2$ vs held-out $Y$** — $\\text{cor}(Y_{\\text{test}},\\; \\mathbf{G}\\hat{\\mathbf{w}})^2$\n", + "5. **Coefficient heatmap** — ensemble $\\zeta_k$ across replicates\n", + "6. **Scatter plot** — ensemble $R^2$ vs. best single-method $R^2$\n", + "7. **Summary CSV** — one row per replicate with all metrics" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "kernel": "SoS" + }, + "outputs": [], + "source": [ + "[analyze_results]\n", + "parameter: results_dir = path(f'{cwd}')\n", + "output: f'{results_dir}/diagnostics.pdf', f'{results_dir}/summary_table.csv'\n", + "task: trunk_workers = 1, walltime = \"1:00:00\", mem = \"8G\", cores = 1\n", + "R: expand = \"${ }\", stderr = f'{results_dir}/analyze.stderr', stdout = f'{results_dir}/analyze.stdout'\n", + "\n", + " library(ggplot2)\n", + " library(tidyr)\n", + " library(dplyr)\n", + "\n", + " results_dir <- \"${results_dir}\"\n", + "\n", + " # ── 1. Load all RDS files ────────────────────────────────\n", + " rds_files <- list.files(results_dir, pattern = \"results\\\\.rds$\",\n", + " recursive = TRUE, full.names = TRUE)\n", + " cat(sprintf(\"Found %d RDS files\\n\", length(rds_files)))\n", + "\n", + " all_data <- lapply(rds_files, function(f) {\n", + " tryCatch(readRDS(f), error = function(e) NULL)\n", + " })\n", + " all_data <- all_data[!sapply(all_data, is.null)]\n", + " all_data <- all_data[!sapply(all_data, function(x) isTRUE(x$truth$error))]\n", + " cat(sprintf(\"%d successful replicates\\n\", length(all_data)))\n", + "\n", + " if (length(all_data) == 0) stop(\"No valid results found\")\n", + "\n", + " # ── 2. Extract metrics ───────────────────────────────────\n", + " metrics_list <- lapply(seq_along(all_data), function(i) {\n", + " x <- all_data[[i]]\n", + " coefs <- if (!is.null(x$ensemble)) x$ensemble$method_coef else NULL\n", + " list(\n", + " rep_id = x$truth$rep_id, seed = x$truth$seed, region = x$truth$region,\n", + " method_coef = coefs,\n", + " method_metrics = x$truth$method_metrics,\n", + " ensemble_metrics = x$truth$ensemble_metrics\n", + " )\n", + " })\n", + "\n", + " # ── 3. Tidy data frames ──────────────────────────────────\n", + " build_metric_df <- function(metric_name) {\n", + " do.call(rbind, lapply(metrics_list, function(m) {\n", + " if (is.null(m$method_metrics)) return(NULL)\n", + " if (!(metric_name %in% rownames(m$method_metrics))) return(NULL)\n", + " methods <- colnames(m$method_metrics)\n", + " vals <- m$method_metrics[metric_name, ]\n", + " method_df <- data.frame(rep_id = m$rep_id, method = methods,\n", + " value = as.numeric(vals),\n", + " stringsAsFactors = FALSE)\n", + " ens_df <- data.frame(rep_id = m$rep_id, method = \"ensemble\",\n", + " value = as.numeric(m$ensemble_metrics[metric_name]),\n", + " stringsAsFactors = FALSE)\n", + " rbind(method_df, ens_df)\n", + " }))\n", + " }\n", + "\n", + "\n", + " mse_w_df <- build_metric_df(\"mse_w\")\n", + " mse_pred_truth_df <- build_metric_df(\"mse_pred_truth\")\n", + " mse_y_df <- build_metric_df(\"mse_y\")\n", + " rsq_y_df <- build_metric_df(\"rsq_y\")\n", + "\n", + " coef_mat <- do.call(rbind, lapply(metrics_list, function(m) {\n", + " if (is.null(m$method_coef)) return(NULL)\n", + " m$method_coef\n", + " }))\n", + " if (!is.null(coef_mat)) {\n", + " rownames(coef_mat) <- sapply(metrics_list[!sapply(metrics_list,\n", + " function(m) is.null(m$method_coef))], function(m) m$rep_id)\n", + " }\n", + "\n", + " scatter_df <- do.call(rbind, lapply(metrics_list, function(m) {\n", + " if (is.null(m$method_metrics)) return(NULL)\n", + " if (!(\"rsq_y\" %in% rownames(m$method_metrics))) return(NULL)\n", + " rsq_vals <- m$method_metrics[\"rsq_y\", ]\n", + " data.frame(\n", + " rep_id = m$rep_id,\n", + " best_single_rsq = max(rsq_vals, na.rm = TRUE),\n", + " best_method = names(which.max(rsq_vals)),\n", + " ensemble_rsq = as.numeric(m$ensemble_metrics[\"rsq_y\"]),\n", + " stringsAsFactors = FALSE\n", + " )\n", + " }))\n", + "\n", + " # ── 4. Diagnostic plots ──────────────────────────────────\n", + " make_boxplot <- function(df, y_label, title, descending = FALSE) {\n", + " if (is.null(df) || nrow(df) == 0) return(NULL)\n", + " ord <- df %>% group_by(method) %>%\n", + " summarise(med = median(value, na.rm = TRUE))\n", + " if (descending) ord <- ord %>% arrange(desc(med)) else ord <- ord %>% arrange(med)\n", + " df$method <- factor(df$method, levels = ord$method)\n", + " ggplot(df, aes(x = method, y = value,\n", + " fill = ifelse(method == \"ensemble\", \"Ensemble\", \"Single\"))) +\n", + " geom_boxplot(outlier.size = 0.8) +\n", + " scale_fill_manual(values = c(\"Ensemble\" = \"#E41A1C\",\n", + " \"Single\" = \"#377EB8\"), name = \"\") +\n", + " labs(title = title, x = \"Method\", y = y_label) +\n", + " theme_bw(base_size = 12) +\n", + " theme(axis.text.x = element_text(angle = 45, hjust = 1))\n", + " }\n", + "\n", + " pdf(\"${_output[0]}\", width = 12, height = 8)\n", + "\n", + "\n", + " p <- make_boxplot(mse_w_df, \"MSE of weights\",\n", + " \"MSE on Weights (LD-independent): mean((w_hat - w_true)^2)\",\n", + " descending = TRUE)\n", + " if (!is.null(p)) print(p)\n", + "\n", + " p <- make_boxplot(mse_pred_truth_df, \"MSE(G*w_hat vs G*w_true)\",\n", + " \"MSE on Predicted Genetic Component (LD-aware, TWAS-relevant)\",\n", + " descending = TRUE)\n", + " if (!is.null(p)) print(p)\n", + "\n", + " p <- make_boxplot(mse_y_df, \"MSE (vs held-out Y)\",\n", + " \"MSE Against Held-Out Y (fresh noise — honest out-of-sample)\",\n", + " descending = TRUE)\n", + " if (!is.null(p)) print(p)\n", + "\n", + " p <- make_boxplot(rsq_y_df, expression(R^2~(vs~held-out~Y)),\n", + " \"R-squared Against Held-Out Y\")\n", + " if (!is.null(p)) print(p)\n", + "\n", + " if (!is.null(coef_mat) && nrow(coef_mat) > 1) {\n", + " coef_long <- as.data.frame(coef_mat) %>%\n", + " mutate(replicate = rownames(coef_mat)) %>%\n", + " pivot_longer(-replicate, names_to = \"method\", values_to = \"zeta\")\n", + " p5 <- ggplot(coef_long, aes(x = method, y = replicate, fill = zeta)) +\n", + " geom_tile() +\n", + " scale_fill_gradient(low = \"white\", high = \"#E41A1C\",\n", + " name = expression(zeta[k]), limits = c(0, 1)) +\n", + " labs(title = \"Ensemble Coefficients Across Replicates\",\n", + " x = \"Method\", y = \"Replicate\") +\n", + " theme_bw(base_size = 10) +\n", + " theme(axis.text.x = element_text(angle = 45, hjust = 1),\n", + " axis.text.y = element_text(size = 6))\n", + " print(p5)\n", + " }\n", + "\n", + " if (!is.null(scatter_df) && nrow(scatter_df) > 0) {\n", + " p6 <- ggplot(scatter_df, aes(x = best_single_rsq, y = ensemble_rsq)) +\n", + " geom_point(alpha = 0.6, size = 2, color = \"#377EB8\") +\n", + " geom_abline(intercept = 0, slope = 1, linetype = \"dashed\",\n", + " color = \"grey40\") +\n", + " labs(title = \"Ensemble vs Best Single Method (R^2 vs held-out Y)\",\n", + " x = expression(Best~single~method~R^2),\n", + " y = expression(Ensemble~R^2)) +\n", + " theme_bw(base_size = 12) + coord_equal()\n", + " print(p6)\n", + " }\n", + "\n", + " dev.off()\n", + " cat(\"Saved plots:\", \"${_output[0]}\", \"\\n\")\n", + "\n", + " # ── 5. Summary table ─────────────────────────────────────\n", + " metric_rows <- c(\"mse_w\", \"mse_pred_truth\", \"mse_y\", \"rsq_y\")\n", + " summary_rows <- lapply(metrics_list, function(m) {\n", + " row <- data.frame(\n", + " rep_id = m$rep_id, region = m$region, seed = m$seed,\n", + " stringsAsFactors = FALSE\n", + " )\n", + " for (mr in metric_rows) {\n", + " row[[paste0(\"ensemble_\", mr)]] <- as.numeric(m$ensemble_metrics[mr])\n", + " }\n", + " if (!is.null(m$method_metrics)) {\n", + " for (mn in colnames(m$method_metrics)) {\n", + " for (mr in metric_rows) {\n", + " if (mr %in% rownames(m$method_metrics)) {\n", + " row[[paste0(mn, \"_\", mr)]] <- m$method_metrics[mr, mn]\n", + " }\n", + " }\n", + " }\n", + " }\n", + " if (!is.null(m$method_coef)) {\n", + " for (mn in names(m$method_coef)) {\n", + " row[[paste0(\"zeta_\", mn)]] <- m$method_coef[[mn]]\n", + " }\n", + " }\n", + " row\n", + " })\n", + " summary_df <- bind_rows(summary_rows)\n", + "\n", + " method_rsq_cols <- grep(\"_rsq_y$\", names(summary_df), value = TRUE)\n", + " method_rsq_cols <- setdiff(method_rsq_cols, \"ensemble_rsq_y\")\n", + " if (length(method_rsq_cols) > 0) {\n", + " summary_df$best_method_rsq_y <- apply(summary_df[method_rsq_cols], 1,\n", + " max, na.rm = TRUE)\n", + " summary_df$best_method <- apply(summary_df[method_rsq_cols], 1, function(r) {\n", + " gsub(\"_rsq_y$\", \"\", names(which.max(r)))\n", + " })\n", + " summary_df$rsq_y_improvement <-\n", + " summary_df$ensemble_rsq_y - summary_df$best_method_rsq_y\n", + " }\n", + "\n", + " write.csv(summary_df, \"${_output[1]}\", row.names = FALSE)\n", + " cat(\"Saved summary:\", \"${_output[1]}\", \"\\n\")\n", + "\n", + " cat(\"\\n=== Summary ===\\n\")\n", + " cat(sprintf(\"Replicates: %d\\n\", nrow(summary_df)))\n", + " cat(sprintf(\"Mean ensemble MSE (weights) : %.6f\\n\",\n", + " mean(summary_df$ensemble_mse_w, na.rm = TRUE)))\n", + " cat(sprintf(\"Mean ensemble MSE (pred truth) : %.6f\\n\",\n", + " mean(summary_df$ensemble_mse_pred_truth, na.rm = TRUE)))\n", + " cat(sprintf(\"Mean ensemble MSE (held-out Y) : %.6f\\n\",\n", + " mean(summary_df$ensemble_mse_y, na.rm = TRUE)))\n", + " cat(sprintf(\"Mean ensemble R^2 (held-out Y) : %.4f\\n\",\n", + " mean(summary_df$ensemble_rsq_y, na.rm = TRUE)))\n", + " if (\"rsq_y_improvement\" %in% names(summary_df)) {\n", + " cat(sprintf(\"Ensemble >= best in %.1f%% of replicates (R^2 vs held-out Y)\\n\",\n", + " 100 * mean(summary_df$rsq_y_improvement >= 0, na.rm = TRUE)))\n", + " }" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "kernel": "SoS" + }, + "source": [ + "## Expected Results\n", + "\n", + "With default simulation parameters ($h^2_g = 0.2$, mixed architecture) on genotype regions:\n", + "\n", + "- **MSE against truth vs. MSE against $Y$**: MSE against truth isolates weight estimation quality (unaffected by noise). MSE against $Y$ includes irreducible noise $\\sigma^2_\\epsilon$, so it is always larger.\n", + "- **Coefficient patterns**: Sparse architectures favor SuSiE/LASSO; polygenic architectures favor BayesR/DPR; the ensemble adapts by assigning larger $\\zeta_k$ to the method best suited to each region's architecture.\n", + "- **Scatter plot**: Most points should lie on or above the diagonal, indicating ensemble matches or improves on the best single method." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "SoS", + "language": "sos", + "name": "sos" + }, + "language_info": { + "codemirror_mode": "sos", + "file_extension": ".sos", + "mimetype": "text/x-sos", + "name": "sos", + "nbconvert_exporter": "sos_notebook.converter.SoS_Exporter", + "pygments_lexer": "sos" + }, + "sos": { + "kernels": [ + [ + "R", + "ir", + "R", + "", + "r" + ], + [ + "SoS", + "sos", + "", + "", + "sos" + ] + ] + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/man/bayes_alphabet_weights.Rd b/man/bayes_alphabet_weights.Rd index 8cf6cced..ef8da658 100644 --- a/man/bayes_alphabet_weights.Rd +++ b/man/bayes_alphabet_weights.Rd @@ -9,6 +9,7 @@ bayes_alphabet_weights( y, method, Z = NULL, + h2 = NULL, nit = 5000, nburn = 1000, nthin = 5, diff --git a/man/bayes_b_weights.Rd b/man/bayes_b_weights.Rd index a2aa9a95..df62e4e2 100644 --- a/man/bayes_b_weights.Rd +++ b/man/bayes_b_weights.Rd @@ -10,7 +10,7 @@ bayes_b_weights( nIter = 10000, burnIn = 2000, thin = 5, - probIn = 0.05, + probIn = 0.2, ... ) } @@ -25,7 +25,7 @@ bayes_b_weights( \item{thin}{Thinning interval. Default is 5.} -\item{probIn}{Prior inclusion probability for each marker. Default is 0.05.} +\item{probIn}{Prior inclusion probability for each marker. Default is 0.2.} \item{...}{Additional arguments passed through to `BGLR::BGLR`.} } diff --git a/man/bayes_c_weights.Rd b/man/bayes_c_weights.Rd index 946b0d4f..e97692fc 100644 --- a/man/bayes_c_weights.Rd +++ b/man/bayes_c_weights.Rd @@ -4,7 +4,7 @@ \alias{bayes_c_weights} \title{Use a rounded spike prior (low-variance Gaussian).} \usage{ -bayes_c_weights(X, y, Z = NULL, ...) +bayes_c_weights(X, y, Z = NULL, pi = 0.1, ...) } \description{ Use a rounded spike prior (low-variance Gaussian). diff --git a/man/dpr_weights.Rd b/man/dpr_weights.Rd index 0802aee5..f0592c72 100644 --- a/man/dpr_weights.Rd +++ b/man/dpr_weights.Rd @@ -7,13 +7,13 @@ \alias{dpr_adaptive_gibbs_weights} \title{Compute Weights Using Dirichlet Process Regression (RcppDPR)} \usage{ -dpr_weights(X, y, fitting_method = "VB", ...) +dpr_weights(X, y, fitting_method = "VB", retain_fit = FALSE, ...) -dpr_vb_weights(X, y, ...) +dpr_vb_weights(X, y, n_k = 8, retain_fit = FALSE, ...) -dpr_gibbs_weights(X, y, ...) +dpr_gibbs_weights(X, y, s_step = 5000, retain_fit = FALSE, ...) -dpr_adaptive_gibbs_weights(X, y, ...) +dpr_adaptive_gibbs_weights(X, y, retain_fit = FALSE, ...) } \arguments{ \item{X}{A numeric matrix of predictors.} diff --git a/man/estimate_sparsity.Rd b/man/estimate_sparsity.Rd new file mode 100644 index 00000000..718b4081 --- /dev/null +++ b/man/estimate_sparsity.Rd @@ -0,0 +1,49 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/twas_weights.R +\name{estimate_sparsity} +\alias{estimate_sparsity} +\title{Estimate Sparsity from Non-Bayesian TWAS Weight Methods} +\usage{ +estimate_sparsity( + weight_results, + methods = "susie", + aggregation = c("median", "mean"), + pip_threshold = 0.5 +) +} +\arguments{ +\item{weight_results}{Named list of weight vectors or matrices as +returned by \code{\link{twas_weights}}. Elements corresponding to +SuSiE or mr.ash methods should have a \code{"fit"} attribute +containing the model fit object (set \code{retain_fits = TRUE} in +\code{twas_weights} to obtain these).} + +\item{methods}{Character vector of short method names to use for +estimation. Must be a subset of the allowed non-full-Bayesian +methods: \code{"susie"}, \code{"susie_ash"}, \code{"susie_inf"}, +\code{"mrash"}, \code{"enet"}, \code{"lasso"}, \code{"scad"}, +\code{"mcp"}, \code{"l0learn"}, \code{"dpr_vb"}.} + +\item{aggregation}{How to combine estimates across methods: +\code{"median"} (default) or \code{"mean"}.} + +\item{pip_threshold}{PIP threshold for SuSiE methods. Default 0.5.} +} +\value{ +A scalar sparsity estimate (proportion of non-zero effects), + with a \code{"per_method"} attribute containing per-method estimates. +} +\description{ +Computes an empirical estimate of the proportion of non-zero effects +(sparsity) from the results of non-full-Bayesian weight methods. +This estimate can be used as a data-driven prior for the inclusion +probability parameters (`pi` for bayesC, `probIn` for BayesB) of +spike-and-slab Bayesian methods. +} +\details{ +For penalized regression methods (LASSO, elastic net, SCAD, MCP, +L0Learn, DPR VB), sparsity is estimated as the proportion of +non-zero coefficients. For SuSiE variants, posterior inclusion +probabilities (PIPs) are thresholded. For mr.ash, the estimated +mixture weight on the spike (zero) component is used. +} diff --git a/man/mrash_weights.Rd b/man/mrash_weights.Rd index 9d567905..7e02bd63 100644 --- a/man/mrash_weights.Rd +++ b/man/mrash_weights.Rd @@ -4,7 +4,7 @@ \alias{mrash_weights} \title{Compute Weights Using mr.ash Shrinkage} \usage{ -mrash_weights(X, y, init_prior_sd = TRUE, ...) +mrash_weights(X, y, init_prior_sd = TRUE, retain_fit = FALSE, ...) } \description{ This function fits the `mr.ash` model (adaptive shrinkage regression) to estimate weights diff --git a/man/twas_weights.Rd b/man/twas_weights.Rd index 49cbf643..9bdd8926 100644 --- a/man/twas_weights.Rd +++ b/man/twas_weights.Rd @@ -4,7 +4,7 @@ \alias{twas_weights} \title{Run multiple TWAS weight methods} \usage{ -twas_weights(X, Y, weight_methods, num_threads = 1) +twas_weights(X, Y, weight_methods, num_threads = 1, retain_fits = FALSE) } \arguments{ \item{X}{A matrix of samples by features, where each row represents a sample and each column a feature.} diff --git a/man/twas_weights_pipeline.Rd b/man/twas_weights_pipeline.Rd index 82ad949c..f07bc655 100644 --- a/man/twas_weights_pipeline.Rd +++ b/man/twas_weights_pipeline.Rd @@ -17,7 +17,9 @@ twas_weights_pipeline( ensemble = FALSE, ensemble_r2_threshold = 0.01, ensemble_solver = "quadprog", - ensemble_alpha = 1 + ensemble_alpha = 1, + estimate_pi_methods = c("susie"), + pi_aggregation = "median" ) } \arguments{ @@ -52,6 +54,18 @@ for ensemble learning. One of \code{"quadprog"}, \code{"nnls"}, \item{ensemble_alpha}{Elastic net mixing parameter, used only when \code{ensemble_solver = "glmnet"}. Defaults to 1 (lasso).} + +\item{estimate_pi_methods}{Character vector of short method names to +use for empirical sparsity estimation. The estimated sparsity is +used as the \code{pi} prior for bayesC and the \code{probIn} prior +for BayesB, unless the user provides explicit values in +\code{weight_methods}. Set to \code{NULL} to disable empirical +estimation. Defaults to \code{c("susie")}. See +\code{\link{estimate_sparsity}} for allowed method names.} + +\item{pi_aggregation}{How to combine sparsity estimates when +multiple estimation methods are used: \code{"median"} (default) +or \code{"mean"}.} } \value{ A list containing results from the TWAS pipeline, including TWAS weights, predictions, and optionally cross-validation results. diff --git a/pixi.toml b/pixi.toml index d5a0fb20..8589a366 100644 --- a/pixi.toml +++ b/pixi.toml @@ -39,6 +39,7 @@ r45 = {features = ["r45"]} "r-markdown" = "*" "r-pkgdown" = "*" "r-rcmdcheck" = "*" +"r-simxqtl" = "*" "r-testthat" = "*" "r-tidyverse" = "*" "bioconductor-iranges" = "*" @@ -71,6 +72,7 @@ r45 = {features = ["r45"]} "r-magrittr" = "*" "r-mashr" = "*" "r-matrixstats" = "*" +"r-mr.ash.alpha" = "*" "r-mr.mashr" = "*" "r-mvsusier" = "*" "r-ncvreg" = "*" From d517456ebdfd84afc0be49021bd402608ed4c9c8 Mon Sep 17 00:00:00 2001 From: Daniel Nachun Date: Wed, 13 May 2026 22:54:49 -0700 Subject: [PATCH 2/7] update defaults --- R/twas_weights.R | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/R/twas_weights.R b/R/twas_weights.R index 6718fc24..02daf555 100644 --- a/R/twas_weights.R +++ b/R/twas_weights.R @@ -33,7 +33,7 @@ if (methods == "fast_default") { methods <- c("susie", "mrash", "enet", "lasso") } else if (methods == "default") { - methods <- c("susie", "mrash", "enet", "lasso", "bayes_r", "dpr_gibbs") + methods <- c("susie", "mrash", "enet", "lasso", "bayes_r", "bayes_c") } } From 3429e00fbc4a662b85e30cd865f730b51647c4df Mon Sep 17 00:00:00 2001 From: Daniel Nachun Date: Wed, 13 May 2026 23:39:33 -0700 Subject: [PATCH 3/7] update documentation --- man/estimate_sparsity.Rd | 48 +++++++++++------------------------- man/twas_weights_pipeline.Rd | 9 ++----- 2 files changed, 16 insertions(+), 41 deletions(-) diff --git a/man/estimate_sparsity.Rd b/man/estimate_sparsity.Rd index 718b4081..ada2eb1b 100644 --- a/man/estimate_sparsity.Rd +++ b/man/estimate_sparsity.Rd @@ -2,48 +2,28 @@ % Please edit documentation in R/twas_weights.R \name{estimate_sparsity} \alias{estimate_sparsity} -\title{Estimate Sparsity from Non-Bayesian TWAS Weight Methods} +\title{Estimate Sparsity from mr.ash Mixture Proportions} \usage{ -estimate_sparsity( - weight_results, - methods = "susie", - aggregation = c("median", "mean"), - pip_threshold = 0.5 -) +estimate_sparsity(weight_results, methods = "mrash") } \arguments{ \item{weight_results}{Named list of weight vectors or matrices as -returned by \code{\link{twas_weights}}. Elements corresponding to -SuSiE or mr.ash methods should have a \code{"fit"} attribute -containing the model fit object (set \code{retain_fits = TRUE} in -\code{twas_weights} to obtain these).} +returned by \code{\link{twas_weights}}. The mr.ash element should +have a \code{"fit"} attribute containing the model fit object +(set \code{retain_fits = TRUE} in \code{twas_weights} to obtain this).} -\item{methods}{Character vector of short method names to use for -estimation. Must be a subset of the allowed non-full-Bayesian -methods: \code{"susie"}, \code{"susie_ash"}, \code{"susie_inf"}, -\code{"mrash"}, \code{"enet"}, \code{"lasso"}, \code{"scad"}, -\code{"mcp"}, \code{"l0learn"}, \code{"dpr_vb"}.} - -\item{aggregation}{How to combine estimates across methods: -\code{"median"} (default) or \code{"mean"}.} - -\item{pip_threshold}{PIP threshold for SuSiE methods. Default 0.5.} +\item{methods}{Must be \code{"mrash"}.} } \value{ -A scalar sparsity estimate (proportion of non-zero effects), - with a \code{"per_method"} attribute containing per-method estimates. +A scalar sparsity estimate (proportion of non-zero effects). } \description{ Computes an empirical estimate of the proportion of non-zero effects -(sparsity) from the results of non-full-Bayesian weight methods. -This estimate can be used as a data-driven prior for the inclusion -probability parameters (`pi` for bayesC, `probIn` for BayesB) of -spike-and-slab Bayesian methods. -} -\details{ -For penalized regression methods (LASSO, elastic net, SCAD, MCP, -L0Learn, DPR VB), sparsity is estimated as the proportion of -non-zero coefficients. For SuSiE variants, posterior inclusion -probabilities (PIPs) are thresholded. For mr.ash, the estimated -mixture weight on the spike (zero) component is used. +(sparsity) from the mr.ash fit. mr.ash fits a mixture model with a +point mass at zero (spike) plus continuous components (slab), and +learns the mixture proportions via variational EM. The sparsity +estimate \code{1 - pi[1]} is the empirical Bayes estimate of the +non-null proportion, which can be used as a data-driven prior for +the inclusion probability parameters (\code{pi} for bayesC, +\code{probIn} for BayesB) of spike-and-slab Bayesian methods. } diff --git a/man/twas_weights_pipeline.Rd b/man/twas_weights_pipeline.Rd index f07bc655..7f4e2ddf 100644 --- a/man/twas_weights_pipeline.Rd +++ b/man/twas_weights_pipeline.Rd @@ -18,8 +18,7 @@ twas_weights_pipeline( ensemble_r2_threshold = 0.01, ensemble_solver = "quadprog", ensemble_alpha = 1, - estimate_pi_methods = c("susie"), - pi_aggregation = "median" + estimate_pi_methods = c("mrash") ) } \arguments{ @@ -60,12 +59,8 @@ use for empirical sparsity estimation. The estimated sparsity is used as the \code{pi} prior for bayesC and the \code{probIn} prior for BayesB, unless the user provides explicit values in \code{weight_methods}. Set to \code{NULL} to disable empirical -estimation. Defaults to \code{c("susie")}. See +estimation. Defaults to \code{c("mrash")}. See \code{\link{estimate_sparsity}} for allowed method names.} - -\item{pi_aggregation}{How to combine sparsity estimates when -multiple estimation methods are used: \code{"median"} (default) -or \code{"mean"}.} } \value{ A list containing results from the TWAS pipeline, including TWAS weights, predictions, and optionally cross-validation results. From c08ea91320970d94ba296c722cfd5619e2177e22 Mon Sep 17 00:00:00 2001 From: Daniel Nachun Date: Wed, 13 May 2026 23:52:09 -0700 Subject: [PATCH 4/7] cleanup --- R/twas_weights.R | 84 +++++++++--------------------- man/estimate_sparsity.Rd | 4 +- man/twas_weights_pipeline.Rd | 10 +--- tests/testthat/test_twas_weights.R | 12 +++-- 4 files changed, 33 insertions(+), 77 deletions(-) diff --git a/R/twas_weights.R b/R/twas_weights.R index 02daf555..50d4d385 100644 --- a/R/twas_weights.R +++ b/R/twas_weights.R @@ -29,11 +29,12 @@ ) # Handle presets + fast_default <- c("susie", "susie_inf", "mrash", "enet", "lasso", "mcp", "scad", "l0learn") if (length(methods) == 1) { if (methods == "fast_default") { - methods <- c("susie", "mrash", "enet", "lasso") + methods <- fast_default } else if (methods == "default") { - methods <- c("susie", "mrash", "enet", "lasso", "bayes_r", "bayes_c") + methods <- c(fast_default, "bayes_r", "bayes_c") } } @@ -514,14 +515,6 @@ twas_predict <- function(X, weights_list) { setNames(lapply(weights_list, function(w) X %*% w), gsub("_weights", "_predicted", names(weights_list))) } -# Short method names allowed for sparsity estimation (non-full-Bayesian methods). -# @noRd -.sparsity_estimation_methods <- "mrash" - -# Map short method names to weight function names for sparsity estimation. -# @noRd -.sparsity_method_fn_map <- c(mrash = "mrash_weights") - #' Estimate Sparsity from mr.ash Mixture Proportions #' #' Computes an empirical estimate of the proportion of non-zero effects @@ -537,19 +530,11 @@ twas_predict <- function(X, weights_list) { #' returned by \code{\link{twas_weights}}. The mr.ash element should #' have a \code{"fit"} attribute containing the model fit object #' (set \code{retain_fits = TRUE} in \code{twas_weights} to obtain this). -#' @param methods Must be \code{"mrash"}. #' #' @return A scalar sparsity estimate (proportion of non-zero effects). #' @export -estimate_sparsity <- function(weight_results, methods = "mrash") { - invalid <- setdiff(methods, .sparsity_estimation_methods) - if (length(invalid) > 0) { - stop("Invalid sparsity estimation method(s): ", paste(invalid, collapse = ", "), - ". Allowed: ", paste(.sparsity_estimation_methods, collapse = ", ")) - } - - fn_name <- .sparsity_method_fn_map["mrash"] - w <- weight_results[[fn_name]] +estimate_sparsity <- function(weight_results) { + w <- weight_results[["mrash_weights"]] if (is.null(w)) { stop("mr.ash weights ('mrash_weights') not found in weight_results.") } @@ -591,13 +576,6 @@ estimate_sparsity <- function(weight_results, methods = "mrash") { #' \code{\link{ensemble_weights}}. Defaults to \code{"quadprog"}. #' @param ensemble_alpha Elastic net mixing parameter, used only when #' \code{ensemble_solver = "glmnet"}. Defaults to 1 (lasso). -#' @param estimate_pi_methods Character vector of short method names to -#' use for empirical sparsity estimation. The estimated sparsity is -#' used as the \code{pi} prior for bayesC and the \code{probIn} prior -#' for BayesB, unless the user provides explicit values in -#' \code{weight_methods}. Set to \code{NULL} to disable empirical -#' estimation. Defaults to \code{c("mrash")}. See -#' \code{\link{estimate_sparsity}} for allowed method names. #' #' @return A list containing results from the TWAS pipeline, including TWAS weights, predictions, and optionally cross-validation results. #' @export @@ -618,7 +596,7 @@ twas_weights_pipeline <- function(X, ensemble_r2_threshold = 0.01, ensemble_solver = "quadprog", ensemble_alpha = 1, - estimate_pi_methods = c("mrash")) { + estimate_pi = TRUE) { if (is.character(weight_methods)) { weight_methods <- .twas_method_lookup(weight_methods) } @@ -641,29 +619,16 @@ twas_weights_pipeline <- function(X, !"pi" %in% names(weight_methods$bayes_c_weights) bayes_b_needs_pi <- "bayes_b_weights" %in% names(weight_methods) && !"probIn" %in% names(weight_methods$bayes_b_weights) - needs_pi_estimation <- (bayes_c_needs_pi || bayes_b_needs_pi) && !is.null(estimate_pi_methods) + needs_pi_estimation <- (bayes_c_needs_pi || bayes_b_needs_pi) && estimate_pi if (needs_pi_estimation) { - # Identify which estimation methods are already in the pipeline - est_fn_names <- .sparsity_method_fn_map[estimate_pi_methods] - phase1_methods <- weight_methods[names(weight_methods) %in% est_fn_names] - - # Add any requested estimation methods not already in weight_methods - missing <- setdiff(est_fn_names, names(weight_methods)) - for (m in missing) phase1_methods[[m]] <- list() - - # Pass susie_fit to phase 1 if SuSiE is an estimation method - if (!is.null(susie_fit) && "susie_weights" %in% names(phase1_methods)) { - phase1_methods$susie_weights <- c( - list(susie_fit = susie_fit), - weight_methods[["susie_weights"]][setdiff(names(weight_methods[["susie_weights"]]), "susie_fit")] - ) - } + # Run mr.ash first to estimate sparsity + mrash_methods <- list(mrash_weights = weight_methods[["mrash_weights"]] %||% list()) - message(" Estimating sparsity from: ", paste(estimate_pi_methods, collapse = ", "), " ...") - phase1_weights <- twas_weights(X, y, weight_methods = phase1_methods, retain_fits = TRUE) + message(" Estimating sparsity from mr.ash ...") + mrash_weights <- twas_weights(X, y, weight_methods = mrash_methods, retain_fits = TRUE) - empirical_pi <- estimate_sparsity(phase1_weights, methods = estimate_pi_methods) + empirical_pi <- estimate_sparsity(mrash_weights) message(sprintf(" Empirical sparsity estimate: %.4f", empirical_pi)) res$empirical_pi <- empirical_pi @@ -671,31 +636,30 @@ twas_weights_pipeline <- function(X, if (bayes_c_needs_pi) weight_methods$bayes_c_weights$pi <- as.numeric(empirical_pi) if (bayes_b_needs_pi) weight_methods$bayes_b_weights$probIn <- as.numeric(empirical_pi) - # Phase 2: run remaining methods (those not already computed in phase 1) - phase2_fn_names <- setdiff(names(weight_methods), names(phase1_methods)) + # Run remaining methods (those not already computed) + remaining_fn_names <- setdiff(names(weight_methods), "mrash_weights") - # Pass susie_fit for phase 2 if SuSiE is NOT an estimation method but is in weight_methods - if (!is.null(susie_fit) && "susie_weights" %in% phase2_fn_names) { + if (!is.null(susie_fit) && "susie_weights" %in% remaining_fn_names) { weight_methods$susie_weights <- list(susie_fit = susie_fit) } - if (length(phase2_fn_names) > 0) { - phase2_methods <- weight_methods[phase2_fn_names] - phase2_weights <- twas_weights(X, y, weight_methods = phase2_methods) - res$twas_weights <- c(phase1_weights, phase2_weights) + if (length(remaining_fn_names) > 0) { + remaining_methods <- weight_methods[remaining_fn_names] + remaining_weights <- twas_weights(X, y, weight_methods = remaining_methods) + res$twas_weights <- c(mrash_weights, remaining_weights) } else { - res$twas_weights <- phase1_weights + res$twas_weights <- mrash_weights } # Clean up fit attributes from final weight results res$twas_weights <- lapply(res$twas_weights, function(w) { attr(w, "fit") <- NULL; w }) - # Remove any estimation-only methods that were not in the original weight_methods - if (length(missing) > 0) { - res$twas_weights <- res$twas_weights[setdiff(names(res$twas_weights), missing)] + # Remove mr.ash if it was not in the original weight_methods + if (!"mrash_weights" %in% names(weight_methods)) { + res$twas_weights[["mrash_weights"]] <- NULL } } else { - # Original flow: run all methods at once + # Run all methods at once if (!is.null(susie_fit) && !is.null(weight_methods$susie_weights)) { weight_methods$susie_weights <- list(susie_fit = susie_fit) } diff --git a/man/estimate_sparsity.Rd b/man/estimate_sparsity.Rd index ada2eb1b..147943cf 100644 --- a/man/estimate_sparsity.Rd +++ b/man/estimate_sparsity.Rd @@ -4,15 +4,13 @@ \alias{estimate_sparsity} \title{Estimate Sparsity from mr.ash Mixture Proportions} \usage{ -estimate_sparsity(weight_results, methods = "mrash") +estimate_sparsity(weight_results) } \arguments{ \item{weight_results}{Named list of weight vectors or matrices as returned by \code{\link{twas_weights}}. The mr.ash element should have a \code{"fit"} attribute containing the model fit object (set \code{retain_fits = TRUE} in \code{twas_weights} to obtain this).} - -\item{methods}{Must be \code{"mrash"}.} } \value{ A scalar sparsity estimate (proportion of non-zero effects). diff --git a/man/twas_weights_pipeline.Rd b/man/twas_weights_pipeline.Rd index 7f4e2ddf..0379b5d7 100644 --- a/man/twas_weights_pipeline.Rd +++ b/man/twas_weights_pipeline.Rd @@ -18,7 +18,7 @@ twas_weights_pipeline( ensemble_r2_threshold = 0.01, ensemble_solver = "quadprog", ensemble_alpha = 1, - estimate_pi_methods = c("mrash") + estimate_pi = TRUE ) } \arguments{ @@ -53,14 +53,6 @@ for ensemble learning. One of \code{"quadprog"}, \code{"nnls"}, \item{ensemble_alpha}{Elastic net mixing parameter, used only when \code{ensemble_solver = "glmnet"}. Defaults to 1 (lasso).} - -\item{estimate_pi_methods}{Character vector of short method names to -use for empirical sparsity estimation. The estimated sparsity is -used as the \code{pi} prior for bayesC and the \code{probIn} prior -for BayesB, unless the user provides explicit values in -\code{weight_methods}. Set to \code{NULL} to disable empirical -estimation. Defaults to \code{c("mrash")}. See -\code{\link{estimate_sparsity}} for allowed method names.} } \value{ A list containing results from the TWAS pipeline, including TWAS weights, predictions, and optionally cross-validation results. diff --git a/tests/testthat/test_twas_weights.R b/tests/testthat/test_twas_weights.R index af38a80c..b0076251 100644 --- a/tests/testthat/test_twas_weights.R +++ b/tests/testthat/test_twas_weights.R @@ -32,19 +32,21 @@ make_data <- function(n = 50, p = 10, seed = 42, add_zero_var_col = FALSE) { # # =========================================================================== -test_that(".twas_method_lookup: 'default' preset returns 6 methods", { +test_that(".twas_method_lookup: 'default' preset returns 10 methods", { result <- pecotmr:::.twas_method_lookup("default") expected_names <- c( - "susie_weights", "mrash_weights", "enet_weights", - "lasso_weights", "bayes_r_weights", "dpr_gibbs_weights" + "susie_weights", "susie_inf_weights", "mrash_weights", "enet_weights", + "lasso_weights", "mcp_weights", "scad_weights", "l0learn_weights", + "bayes_r_weights", "bayes_c_weights" ) expect_equal(sort(names(result)), sort(expected_names)) }) -test_that(".twas_method_lookup: 'fast_default' preset returns 4 methods", { +test_that(".twas_method_lookup: 'fast_default' preset returns 8 methods", { result <- pecotmr:::.twas_method_lookup("fast_default") expected_names <- c( - "susie_weights", "mrash_weights", "enet_weights", "lasso_weights" + "susie_weights", "susie_inf_weights", "mrash_weights", "enet_weights", + "lasso_weights", "mcp_weights", "scad_weights", "l0learn_weights" ) expect_equal(sort(names(result)), sort(expected_names)) }) From e818fdaf51963b7d8a13eb334f332251637565ac Mon Sep 17 00:00:00 2001 From: Daniel Nachun Date: Thu, 14 May 2026 00:25:22 -0700 Subject: [PATCH 5/7] fix tests --- tests/testthat/test_twas_weights.R | 86 ++++++++++++++++++++++-------- 1 file changed, 63 insertions(+), 23 deletions(-) diff --git a/tests/testthat/test_twas_weights.R b/tests/testthat/test_twas_weights.R index b0076251..64103e07 100644 --- a/tests/testthat/test_twas_weights.R +++ b/tests/testthat/test_twas_weights.R @@ -732,12 +732,17 @@ test_that("twas_weights_pipeline: returns list with expected structure (mocked)" enet_weights = function(X, y, ...) rep(0.1, ncol(X)), lasso_weights = function(X, y, ...) rep(0.2, ncol(X)), bayes_r_weights = function(X, y, ...) rep(0, ncol(X)), - dpr_gibbs_weights = function(X, y, ...) rep(0, ncol(X)), + bayes_c_weights = function(X, y, ...) rep(0, ncol(X)), mrash_weights = function(X, y, ...) rep(0, ncol(X)), - susie_weights = function(X, y, ...) rep(0, ncol(X)) + mcp_weights = function(X, y, ...) rep(0, ncol(X)), + scad_weights = function(X, y, ...) rep(0, ncol(X)), + l0learn_weights = function(X, y, ...) rep(0, ncol(X)), + susie_weights = function(X, y, ...) rep(0, ncol(X)), + susie_inf_weights = function(X, y, ...) rep(0, ncol(X)) ) - result <- twas_weights_pipeline(d$X, y_vec, susie_fit = NULL, cv_folds = 0) + result <- twas_weights_pipeline(d$X, y_vec, susie_fit = NULL, cv_folds = 0, + estimate_pi = FALSE) expect_true(is.list(result)) expect_true("twas_weights" %in% names(result)) @@ -748,8 +753,8 @@ test_that("twas_weights_pipeline: returns list with expected structure (mocked)" expect_true(all(enet_w[, 1] == 0.1)) lasso_w <- result$twas_weights[["lasso_weights"]] expect_true(all(lasso_w[, 1] == 0.2)) - # The number of weight methods should equal the 6 default methods - expect_equal(length(result$twas_weights), 6) + # The number of weight methods should equal the 10 default methods + expect_equal(length(result$twas_weights), 10) }) test_that("twas_weights_pipeline: twas_weights contains all default methods", { @@ -760,16 +765,23 @@ test_that("twas_weights_pipeline: twas_weights contains all default methods", { enet_weights = function(X, y, ...) rep(0.1, ncol(X)), lasso_weights = function(X, y, ...) rep(0.2, ncol(X)), bayes_r_weights = function(X, y, ...) rep(0, ncol(X)), - dpr_gibbs_weights = function(X, y, ...) rep(0, ncol(X)), + bayes_c_weights = function(X, y, ...) rep(0, ncol(X)), mrash_weights = function(X, y, ...) rep(0, ncol(X)), - susie_weights = function(X, y, ...) rep(0, ncol(X)) + mcp_weights = function(X, y, ...) rep(0, ncol(X)), + scad_weights = function(X, y, ...) rep(0, ncol(X)), + l0learn_weights = function(X, y, ...) rep(0, ncol(X)), + susie_weights = function(X, y, ...) rep(0, ncol(X)), + susie_inf_weights = function(X, y, ...) rep(0, ncol(X)) ) - result <- twas_weights_pipeline(d$X, y_vec, susie_fit = NULL, cv_folds = 0) + result <- twas_weights_pipeline(d$X, y_vec, susie_fit = NULL, cv_folds = 0, + estimate_pi = FALSE) expected_methods <- c( "enet_weights", "lasso_weights", "bayes_r_weights", - "dpr_gibbs_weights", "mrash_weights", "susie_weights" + "bayes_c_weights", "mrash_weights", "mcp_weights", + "scad_weights", "l0learn_weights", "susie_weights", + "susie_inf_weights" ) expect_true(all(expected_methods %in% names(result$twas_weights))) }) @@ -782,16 +794,23 @@ test_that("twas_weights_pipeline: predictions have _predicted suffix", { enet_weights = function(X, y, ...) rep(0, ncol(X)), lasso_weights = function(X, y, ...) rep(0, ncol(X)), bayes_r_weights = function(X, y, ...) rep(0, ncol(X)), - dpr_gibbs_weights = function(X, y, ...) rep(0, ncol(X)), + bayes_c_weights = function(X, y, ...) rep(0, ncol(X)), mrash_weights = function(X, y, ...) rep(0, ncol(X)), - susie_weights = function(X, y, ...) rep(0, ncol(X)) + mcp_weights = function(X, y, ...) rep(0, ncol(X)), + scad_weights = function(X, y, ...) rep(0, ncol(X)), + l0learn_weights = function(X, y, ...) rep(0, ncol(X)), + susie_weights = function(X, y, ...) rep(0, ncol(X)), + susie_inf_weights = function(X, y, ...) rep(0, ncol(X)) ) - result <- twas_weights_pipeline(d$X, y_vec, susie_fit = NULL, cv_folds = 0) + result <- twas_weights_pipeline(d$X, y_vec, susie_fit = NULL, cv_folds = 0, + estimate_pi = FALSE) expected_pred_names <- c( "enet_predicted", "lasso_predicted", "bayes_r_predicted", - "dpr_gibbs_predicted", "mrash_predicted", "susie_predicted" + "bayes_c_predicted", "mrash_predicted", "mcp_predicted", + "scad_predicted", "l0learn_predicted", "susie_predicted", + "susie_inf_predicted" ) expect_true(all(expected_pred_names %in% names(result$twas_predictions))) }) @@ -804,12 +823,17 @@ test_that("twas_weights_pipeline: cv_folds=0 skips cross-validation", { enet_weights = function(X, y, ...) rep(0, ncol(X)), lasso_weights = function(X, y, ...) rep(0, ncol(X)), bayes_r_weights = function(X, y, ...) rep(0, ncol(X)), - dpr_gibbs_weights = function(X, y, ...) rep(0, ncol(X)), + bayes_c_weights = function(X, y, ...) rep(0, ncol(X)), mrash_weights = function(X, y, ...) rep(0, ncol(X)), - susie_weights = function(X, y, ...) rep(0, ncol(X)) + mcp_weights = function(X, y, ...) rep(0, ncol(X)), + scad_weights = function(X, y, ...) rep(0, ncol(X)), + l0learn_weights = function(X, y, ...) rep(0, ncol(X)), + susie_weights = function(X, y, ...) rep(0, ncol(X)), + susie_inf_weights = function(X, y, ...) rep(0, ncol(X)) ) - result <- twas_weights_pipeline(d$X, y_vec, susie_fit = NULL, cv_folds = 0) + result <- twas_weights_pipeline(d$X, y_vec, susie_fit = NULL, cv_folds = 0, + estimate_pi = FALSE) expect_false("twas_cv_result" %in% names(result)) # All mock weights were zero, so all predictions should be zero @@ -849,7 +873,11 @@ test_that("twas_weights_pipeline: accepts 'fast_default' preset string", { enet_weights = function(X, y, ...) rep(0, ncol(X)), lasso_weights = function(X, y, ...) rep(0, ncol(X)), mrash_weights = function(X, y, ...) rep(0, ncol(X)), - susie_weights = function(X, y, ...) rep(0, ncol(X)) + mcp_weights = function(X, y, ...) rep(0, ncol(X)), + scad_weights = function(X, y, ...) rep(0, ncol(X)), + l0learn_weights = function(X, y, ...) rep(0, ncol(X)), + susie_weights = function(X, y, ...) rep(0, ncol(X)), + susie_inf_weights = function(X, y, ...) rep(0, ncol(X)) ) result <- twas_weights_pipeline( @@ -857,7 +885,9 @@ test_that("twas_weights_pipeline: accepts 'fast_default' preset string", { weight_methods = "fast_default" ) - expected_methods <- c("susie_weights", "mrash_weights", "enet_weights", "lasso_weights") + expected_methods <- c("susie_weights", "susie_inf_weights", "mrash_weights", + "enet_weights", "lasso_weights", "mcp_weights", + "scad_weights", "l0learn_weights") expect_equal(sort(names(result$twas_weights)), sort(expected_methods)) }) @@ -896,12 +926,17 @@ test_that("twas_weights_pipeline: with susie_fit stores intermediate results", { enet_weights = function(X, y, ...) rep(0, ncol(X)), lasso_weights = function(X, y, ...) rep(0, ncol(X)), bayes_r_weights = function(X, y, ...) rep(0, ncol(X)), - dpr_gibbs_weights = function(X, y, ...) rep(0, ncol(X)), + bayes_c_weights = function(X, y, ...) rep(0, ncol(X)), mrash_weights = function(X, y, ...) rep(0, ncol(X)), - susie_weights = function(X, y, ...) rep(0, ncol(X)) + mcp_weights = function(X, y, ...) rep(0, ncol(X)), + scad_weights = function(X, y, ...) rep(0, ncol(X)), + l0learn_weights = function(X, y, ...) rep(0, ncol(X)), + susie_weights = function(X, y, ...) rep(0, ncol(X)), + susie_inf_weights = function(X, y, ...) rep(0, ncol(X)) ) - result <- twas_weights_pipeline(d$X, y_vec, susie_fit = fake_susie, cv_folds = 0) + result <- twas_weights_pipeline(d$X, y_vec, susie_fit = fake_susie, cv_folds = 0, + estimate_pi = FALSE) expect_true("susie_weights_intermediate" %in% names(result)) expect_true("mu" %in% names(result$susie_weights_intermediate)) @@ -928,8 +963,12 @@ test_that("twas_weights_pipeline: with susie_fit, susie_weights gets susie_fit a enet_weights = function(X, y, ...) rep(0, ncol(X)), lasso_weights = function(X, y, ...) rep(0, ncol(X)), bayes_r_weights = function(X, y, ...) rep(0, ncol(X)), - dpr_gibbs_weights = function(X, y, ...) rep(0, ncol(X)), + bayes_c_weights = function(X, y, ...) rep(0, ncol(X)), mrash_weights = function(X, y, ...) rep(0, ncol(X)), + mcp_weights = function(X, y, ...) rep(0, ncol(X)), + scad_weights = function(X, y, ...) rep(0, ncol(X)), + l0learn_weights = function(X, y, ...) rep(0, ncol(X)), + susie_inf_weights = function(X, y, ...) rep(0, ncol(X)), susie_weights = function(X, y, ...) { args <- list(...) # The pipeline should pass susie_fit as an argument @@ -940,7 +979,8 @@ test_that("twas_weights_pipeline: with susie_fit, susie_weights gets susie_fit a } ) - result <- twas_weights_pipeline(d$X, y_vec, susie_fit = fake_susie, cv_folds = 0) + result <- twas_weights_pipeline(d$X, y_vec, susie_fit = fake_susie, cv_folds = 0, + estimate_pi = FALSE) expect_true(susie_received_fit) }) From 182ad6d17cd8a4384cb2597d0c4dfcc86f208a0c Mon Sep 17 00:00:00 2001 From: Daniel Nachun Date: Thu, 14 May 2026 00:48:18 -0700 Subject: [PATCH 6/7] fix tests --- R/susie_wrapper.R | 18 ++++++++++-------- R/univariate_pipeline.R | 8 +++++--- man/load_multitask_regional_data.Rd | 4 ++-- man/rss_analysis_pipeline.Rd | 6 ++++-- man/susie_rss_pipeline.Rd | 6 ++++-- man/susie_rss_wrapper.Rd | 8 ++++---- 6 files changed, 29 insertions(+), 21 deletions(-) diff --git a/R/susie_wrapper.R b/R/susie_wrapper.R index 8af0fa7b..d661d07f 100644 --- a/R/susie_wrapper.R +++ b/R/susie_wrapper.R @@ -163,9 +163,9 @@ susie_wrapper <- function(X, y, init_L = 5, max_L = 30, l_step = 5, ...) { #' @param L Initial number of causal configurations. #' @param max_L Maximum number of causal configurations. #' @param l_step Step size for increasing L when the limit is reached. -#' @param sketch_samples Sketch LD parameter passed to susie_rss. -#' NULL (default): no variance inflation. TRUE: infer sketch size from X -#' (requires X, not R). Integer: explicit sketch size (for R only). +#' @param R_finite Controls variance inflation to account for estimating +#' the R matrix from a finite reference panel. NULL (default): no +#' variance inflation. Passed directly to susie_rss. #' @param ... Extra parameters passed to susie_rss (e.g., var_y, coverage). #' @return SuSiE RSS fit object after dynamic L adjustment #' @importFrom susieR susie_rss @@ -173,14 +173,14 @@ susie_wrapper <- function(X, y, init_L = 5, max_L = 30, l_step = 5, ...) { susie_rss_wrapper <- function(z, R = NULL, X = NULL, n = NULL, L = 10, max_L = 30, l_step = 5, coverage = 0.95, - sketch_samples = NULL, ...) { + R_finite = NULL, ...) { # Validate: exactly one of R or X if (is.null(R) && is.null(X)) stop("Either R or X must be provided.") if (!is.null(R) && !is.null(X)) stop("Only one of R or X should be provided, not both.") # Build argument list for susie_rss base_args <- list(z = z, n = n, coverage = coverage, - sketch_samples = sketch_samples, ...) + R_finite = R_finite, ...) if (!is.null(X)) base_args$X <- X else base_args$R <- R run_with_L <- function(L_val) do.call(susie_rss, c(base_args, list(L = L_val))) @@ -214,7 +214,9 @@ susie_rss_wrapper <- function(z, R = NULL, X = NULL, n = NULL, #' @param secondary_coverage Secondary coverage levels (default: c(0.7, 0.5)). #' @param signal_cutoff PIP cutoff for susie_post_processor (default: 0.1). #' @param min_abs_corr Minimum absolute correlation for CS purity (default: 0.8). -#' @param sketch_samples Passed to susie_rss. NULL, TRUE, or integer. +#' @param R_finite Controls variance inflation to account for estimating +#' the R matrix from a finite reference panel. NULL (default): no +#' variance inflation. Passed directly to susie_rss. #' @param ... Additional parameters passed to susie_rss (e.g., var_y). #' @return A list with post-processed SuSiE RSS results. #' @importFrom magrittr %>% @@ -227,7 +229,7 @@ susie_rss_pipeline <- function(sumstats, LD_mat = NULL, X_mat = NULL, n = NULL, secondary_coverage = c(0.7, 0.5), signal_cutoff = 0.1, min_abs_corr = 0.8, - sketch_samples = NULL, ...) { + R_finite = NULL, ...) { analysis_method <- match.arg(analysis_method) if (!is.null(sumstats$z)) { @@ -240,7 +242,7 @@ susie_rss_pipeline <- function(sumstats, LD_mat = NULL, X_mat = NULL, n = NULL, # Common args for susie_rss_wrapper common <- list(z = z, n = n, coverage = coverage, - sketch_samples = sketch_samples, ...) + R_finite = R_finite, ...) if (!is.null(X_mat)) common$X <- X_mat else common$R <- LD_mat if (analysis_method == "single_effect") { diff --git a/R/univariate_pipeline.R b/R/univariate_pipeline.R index 958014dc..55db958b 100644 --- a/R/univariate_pipeline.R +++ b/R/univariate_pipeline.R @@ -191,7 +191,9 @@ load_study_LD <- function(ld_path, region) { #' @param impute Whether to impute missing variants via RAISS (default TRUE). #' @param impute_opts List of imputation options (rcond, R2_threshold, minimum_ld, lamb). #' @param pip_cutoff_to_skip PIP threshold for early stopping (default 0, no skip). -#' @param sketch_samples Passed to susie_rss. NULL (default), TRUE, or integer. +#' @param R_finite Controls variance inflation to account for estimating +#' the R matrix from a finite reference panel. NULL (default): no +#' variance inflation. Passed directly to susie_rss. #' @param keep_indel Whether to keep indel variants (default TRUE). #' @param comment_string Comment character for sumstat file (default "#"). #' @param diagnostics Whether to include diagnostic info (default FALSE). @@ -211,7 +213,7 @@ rss_analysis_pipeline <- function( min_abs_corr = 0.8 ), impute = TRUE, impute_opts = list(rcond = 0.01, R2_threshold = 0.6, minimum_ld = 5, lamb = 0.01), - pip_cutoff_to_skip = 0, sketch_samples = NULL, + pip_cutoff_to_skip = 0, R_finite = NULL, keep_indel = TRUE, comment_string = "#", diagnostics = FALSE) { # Detect genotype input: single X matrix or list of X matrices (mixture panel). # susie_rss accepts X=list(X1, X2, ...) for multi-panel mixture. @@ -306,7 +308,7 @@ rss_analysis_pipeline <- function( secondary_coverage = sec_coverage, signal_cutoff = finemapping_opts$signal_cutoff, min_abs_corr = finemapping_opts$min_abs_corr, - sketch_samples = sketch_samples + R_finite = R_finite ) if (!is.null(qc_method)) { res$outlier_number <- qc_results$outlier_number diff --git a/man/load_multitask_regional_data.Rd b/man/load_multitask_regional_data.Rd index 585c2070..46134bef 100644 --- a/man/load_multitask_regional_data.Rd +++ b/man/load_multitask_regional_data.Rd @@ -129,10 +129,10 @@ This function loads a mixture data sets for a specific region, including individ or summary statistics (sumstats, LD). Run \code{load_regional_univariate_data} and \code{load_rss_data} multiple times for different datasets } \section{Loading individual level data from multiple corhorts}{ -NA + } \section{Loading summary statistics from multiple corhorts or data set}{ -NA + } diff --git a/man/rss_analysis_pipeline.Rd b/man/rss_analysis_pipeline.Rd index 0d669b99..b259fde5 100644 --- a/man/rss_analysis_pipeline.Rd +++ b/man/rss_analysis_pipeline.Rd @@ -22,7 +22,7 @@ rss_analysis_pipeline( impute = TRUE, impute_opts = list(rcond = 0.01, R2_threshold = 0.6, minimum_ld = 5, lamb = 0.01), pip_cutoff_to_skip = 0, - sketch_samples = NULL, + R_finite = NULL, keep_indel = TRUE, comment_string = "#", diagnostics = FALSE @@ -65,7 +65,9 @@ signal_cutoff, min_abs_corr).} \item{pip_cutoff_to_skip}{PIP threshold for early stopping (default 0, no skip).} -\item{sketch_samples}{Passed to susie_rss. NULL (default), TRUE, or integer.} +\item{R_finite}{Controls variance inflation to account for estimating +the R matrix from a finite reference panel. NULL (default): no +variance inflation. Passed directly to susie_rss.} \item{keep_indel}{Whether to keep indel variants (default TRUE).} diff --git a/man/susie_rss_pipeline.Rd b/man/susie_rss_pipeline.Rd index e39c7e52..15cb5d84 100644 --- a/man/susie_rss_pipeline.Rd +++ b/man/susie_rss_pipeline.Rd @@ -17,7 +17,7 @@ susie_rss_pipeline( secondary_coverage = c(0.7, 0.5), signal_cutoff = 0.1, min_abs_corr = 0.8, - sketch_samples = NULL, + R_finite = NULL, ... ) } @@ -46,7 +46,9 @@ susie_rss_pipeline( \item{min_abs_corr}{Minimum absolute correlation for CS purity (default: 0.8).} -\item{sketch_samples}{Passed to susie_rss. NULL, TRUE, or integer.} +\item{R_finite}{Controls variance inflation to account for estimating +the R matrix from a finite reference panel. NULL (default): no +variance inflation. Passed directly to susie_rss.} \item{...}{Additional parameters passed to susie_rss (e.g., var_y).} } diff --git a/man/susie_rss_wrapper.Rd b/man/susie_rss_wrapper.Rd index ab860fc6..49978992 100644 --- a/man/susie_rss_wrapper.Rd +++ b/man/susie_rss_wrapper.Rd @@ -13,7 +13,7 @@ susie_rss_wrapper( max_L = 30, l_step = 5, coverage = 0.95, - sketch_samples = NULL, + R_finite = NULL, ... ) } @@ -33,9 +33,9 @@ When provided, susie_rss uses the low-rank X interface.} \item{l_step}{Step size for increasing L when the limit is reached.} -\item{sketch_samples}{Sketch LD parameter passed to susie_rss. -NULL (default): no variance inflation. TRUE: infer sketch size from X -(requires X, not R). Integer: explicit sketch size (for R only).} +\item{R_finite}{Controls variance inflation to account for estimating +the R matrix from a finite reference panel. NULL (default): no +variance inflation. Passed directly to susie_rss.} \item{...}{Extra parameters passed to susie_rss (e.g., var_y, coverage).} } From 6d46455bcb2f1016b651dc29f05491ef1f8f167c Mon Sep 17 00:00:00 2001 From: danielnachun Date: Thu, 14 May 2026 07:49:37 +0000 Subject: [PATCH 7/7] Update documentation --- man/load_multitask_regional_data.Rd | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/man/load_multitask_regional_data.Rd b/man/load_multitask_regional_data.Rd index 46134bef..585c2070 100644 --- a/man/load_multitask_regional_data.Rd +++ b/man/load_multitask_regional_data.Rd @@ -129,10 +129,10 @@ This function loads a mixture data sets for a specific region, including individ or summary statistics (sumstats, LD). Run \code{load_regional_univariate_data} and \code{load_rss_data} multiple times for different datasets } \section{Loading individual level data from multiple corhorts}{ - +NA } \section{Loading summary statistics from multiple corhorts or data set}{ - +NA }