diff --git a/.Rbuildignore b/.Rbuildignore index 75b92fa8..777dfec6 100644 --- a/.Rbuildignore +++ b/.Rbuildignore @@ -1,3 +1,4 @@ +^run_tests\.R$ ^\.github$ ^LICENSE\.md$ ^README\.md$ diff --git a/.gitignore b/.gitignore index 80fa8d64..f8e2c429 100644 --- a/.gitignore +++ b/.gitignore @@ -46,7 +46,8 @@ Rplot.pdf sccomp_draws_files .DS_Store docs - /.quarto/ **/*.quarto_ipynb .Renviron +*.exe +sccomp.Rcheck* diff --git a/DESCRIPTION b/DESCRIPTION index c9ced17c..59410f74 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,14 +1,13 @@ Package: sccomp Type: Package Title: Differential Composition and Variability Analysis for Single-Cell Data -Version: 2.1.31 -Date: 2026-04-05 +Version: 2.1.33 +Date: 2026-05-11 Authors@R: c(person("Stefano", "Mangiola", email = "stefano.mangiola@unimelb.edu.au", role = c("aut", "cre")), person("Alexandra J.", "Roth-Schulze", role = "aut"), person("Marie", "Trussart", role = "aut"), person("Enrique", "Zozaya-Valdés", role = "aut"), person("Mengyao", "Ma", role = "aut"), person("Zijie", "Gao", role = "aut"), person("Alan F.", "Rubin", role = "aut"), person("Terence P.", "Speed", role = "aut"), person("Heejung", "Shim", role = "aut"), person("Anthony T.", "Papenfuss", role = "aut")) Description: Comprehensive R package for differential composition and variability analysis in single-cell RNA sequencing, CyTOF, and microbiome data. Provides robust Bayesian modeling with outlier detection, random effects, and advanced statistical methods for cell type proportion analysis. Features include probabilistic outlier identification, mixed-effect modeling, differential variability testing, and comprehensive visualization tools. Perfect for cancer research, immunology, developmental biology, and single-cell genomics applications. License: GPL-3 Encoding: UTF-8 Roxygen: list(markdown = TRUE) -RoxygenNote: 7.3.3 Depends: R (>= 4.3.0), instantiate (>= 0.2.3) @@ -30,13 +29,16 @@ Imports: tibble, ggplot2, ggrepel, + ggside, patchwork, forcats, readr, stringr, glue, - SingleCellExperiment + SingleCellExperiment, + posterior Suggests: + mgcv, knitr, rmarkdown, BiocStyle, @@ -48,7 +50,6 @@ Suggests: tidyseurat, tidySingleCellExperiment, bayesplot, - posterior, remotes Additional_repositories: https://mc-stan.org/r-packages/ @@ -59,6 +60,7 @@ VignetteBuilder: knitr URL: https://github.com/MangiolaLaboratory/sccomp, https://mangiolalaboratory.github.io/sccomp/ BugReports: https://github.com/MangiolaLaboratory/sccomp/issues Config/testthat/edition: 3 -Config/testthat/parallel: true -Config/testthat/snapshot/parallel: true +Config/testthat/parallel: false +Config/testthat/snapshot/parallel: false Config/testthat/snapshot/parallel/workers: 2 +Config/roxygen2/version: 8.0.0 diff --git a/NAMESPACE b/NAMESPACE index 206c10fe..7354520f 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -21,6 +21,8 @@ export(plot_2D_intervals) export(sccomp_boxplot) export(sccomp_calculate_residuals) export(sccomp_estimate) +export(sccomp_plot_intervals_1D) +export(sccomp_plot_intervals_2D) export(sccomp_predict) export(sccomp_proportional_fold_change) export(sccomp_remove_outliers) @@ -42,6 +44,7 @@ importFrom(dplyr,all_of) importFrom(dplyr,any_of) importFrom(dplyr,arrange) importFrom(dplyr,as_tibble) +importFrom(dplyr,bind_rows) importFrom(dplyr,c_across) importFrom(dplyr,count) importFrom(dplyr,cummean) @@ -62,24 +65,44 @@ importFrom(dplyr,rename) importFrom(dplyr,row_number) importFrom(dplyr,rowwise) importFrom(dplyr,select) +importFrom(dplyr,slice) importFrom(dplyr,summarise) importFrom(dplyr,where) importFrom(dplyr,with_groups) +importFrom(forcats,fct_inorder) importFrom(forcats,fct_relevel) importFrom(forcats,fct_reorder) importFrom(ggplot2,aes) -importFrom(ggplot2,annotate) +importFrom(ggplot2,element_blank) +importFrom(ggplot2,element_rect) +importFrom(ggplot2,element_text) importFrom(ggplot2,facet_wrap) +importFrom(ggplot2,geom_blank) importFrom(ggplot2,geom_errorbar) importFrom(ggplot2,geom_hline) +importFrom(ggplot2,geom_line) importFrom(ggplot2,geom_point) importFrom(ggplot2,geom_vline) importFrom(ggplot2,ggplot) +importFrom(ggplot2,guide_legend) +importFrom(ggplot2,guides) +importFrom(ggplot2,labs) +importFrom(ggplot2,scale_alpha_manual) +importFrom(ggplot2,scale_color_manual) +importFrom(ggplot2,scale_fill_manual) +importFrom(ggplot2,theme) +importFrom(ggplot2,theme_bw) +importFrom(ggplot2,xlab) +importFrom(ggplot2,ylab) importFrom(ggrepel,geom_text_repel) +importFrom(ggside,geom_ysidedensity) +importFrom(ggside,scale_ysidex_continuous) +importFrom(ggside,theme_ggside_void) importFrom(glue,glue) importFrom(instantiate,stan_cmdstan_exists) importFrom(instantiate,stan_package_compile) importFrom(instantiate,stan_package_model) +importFrom(lifecycle,deprecate_soft) importFrom(lifecycle,deprecate_warn) importFrom(lifecycle,is_present) importFrom(magrittr,"%$%") @@ -91,6 +114,11 @@ importFrom(magrittr,not) importFrom(magrittr,subtract) importFrom(parallel,detectCores) importFrom(patchwork,wrap_plots) +importFrom(posterior,as_draws_df) +importFrom(posterior,ess_bulk) +importFrom(posterior,ess_tail) +importFrom(posterior,rhat) +importFrom(posterior,summarise_draws) importFrom(purrr,as_mapper) importFrom(purrr,map) importFrom(purrr,map2) diff --git a/R/model_fitting.R b/R/model_fitting.R index fcc1c9f4..e463af97 100644 --- a/R/model_fitting.R +++ b/R/model_fitting.R @@ -2,7 +2,8 @@ fit_model = function( data_for_model, model_name, censoring_iteration = 1, cores = detectCores(), quantile = 0.95, warmup_samples = 300, approximate_posterior_inference = NULL, inference_method, verbose = TRUE, - seed , pars = c("beta", "alpha", "prec_coeff","prec_sd"), output_samples = NULL, chains=NULL, max_sampling_iterations = 20000, + seed , pars = c("beta", "alpha", "prec_sd", "prec_intercept_1", "prec_slope_1", + "prec_intercept_2", "prec_slope_2"), output_samples = NULL, chains=NULL, max_sampling_iterations = 20000, output_directory = "sccomp_draws_files", sig_figs = 9, cache_stan_model = sccomp_stan_models_cache_dir, @@ -41,59 +42,69 @@ fit_model = function( min(cores) # chains = 3 - - init_list=list( - prec_coeff = c(5,0), - prec_sd = 1, - alpha = matrix(c(rep(5, data_for_model$M), rep(0, (data_for_model$A-1) *data_for_model$M)), nrow = data_for_model$A, byrow = TRUE), - beta_raw = matrix(0, data_for_model$C , data_for_model$M) , - mix_p = 0.1 + + init_list = list( + alpha = matrix(c(rep(5, data_for_model$M), + rep(0, (data_for_model$A - 1) * data_for_model$M)), + nrow = data_for_model$A, byrow = TRUE), + beta_raw = matrix(0, data_for_model$C, data_for_model$M), + log_prec_sd = rep(0, data_for_model$A), + mix_p = 0.5 ) - - if(data_for_model$n_random_eff>0){ - init_list$random_effect_raw = matrix(0, data_for_model$ncol_X_random_eff[1] , data_for_model$M) - init_list$random_effect_sigma_raw = matrix(0, data_for_model$M , data_for_model$how_many_factors_in_random_design[1]) - init_list$sigma_correlation_factor = array(0, dim = c( - data_for_model$M, - data_for_model$how_many_factors_in_random_design[1], - data_for_model$how_many_factors_in_random_design[1] - )) - - # init_list$random_effect_sigma_mu = 0.5 |> as.array() - # init_list$random_effect_sigma_sigma = 0.2 |> as.array() + + has_variability_intercept <- isTRUE(as.logical(data_for_model$intercept_in_design)) + bimodal <- data_for_model$bimodal_mean_variability_association == 1L + intercept_init <- if (bimodal) c(4, 5) else c(4) + non_intercept_init <- if (bimodal) c(0, 1) else c(0) + + init_list$prec_intercept <- + if (has_variability_intercept) { + c(list(intercept_init), rep(list(non_intercept_init), data_for_model$A - 1L)) + } else { + rep(list(intercept_init), data_for_model$A) + } + + init_list$prec_slope_1 = rep(0, data_for_model$A) + if (bimodal) { + init_list$prec_slope_2 = rep(0, data_for_model$A) + } + + # Random effect inits - 4 uniform slots (one per non-empty random-effect block). + # Each slot gets zero-initialised raws + an identity-like correlation Cholesky. + if (data_for_model$n_random_eff > 0) { init_list$zero_random_effect = rep(0, size = 1) |> as.array() - } - - if(data_for_model$n_random_eff>1){ - init_list$random_effect_raw_2 = matrix(0, data_for_model$ncol_X_random_eff[2] , data_for_model$M) - init_list$random_effect_sigma_raw_2 = matrix(0, data_for_model$M , data_for_model$how_many_factors_in_random_design[2]) - init_list$sigma_correlation_factor_2 = array(0, dim = c( - data_for_model$M, - data_for_model$how_many_factors_in_random_design[2], - data_for_model$how_many_factors_in_random_design[2] - )) - - } - + for (k in seq_len(4L)) { + if (data_for_model$ncol_X_random_eff[k] == 0) next + K = data_for_model$how_many_factors_in_random_design[k] + + init_list[[paste0("random_effect_raw_", k)]] = + matrix(0, data_for_model$ncol_X_random_eff[k], data_for_model$M) + init_list[[paste0("random_effect_sigma_raw_", k)]] = + matrix(0, data_for_model$M, K) + init_list[[paste0("sigma_correlation_factor_", k)]] = + array(0, dim = c(data_for_model$M, K, K)) + } + } + init = map(1:chains, ~ init_list) %>% setNames(as.character(1:chains)) - + #output_directory = "sccomp_draws_files" dir.create(output_directory, showWarnings = FALSE) - + # Fit mod = load_model(model_name, threads = cores, cache_dir = cache_stan_model) - + # Avoid 0 proportions if(data_for_model$is_proportion && min(data_for_model$y_proportion)==0){ warning("sccomp says: your proportion values include 0. Assuming that 0s derive from a precision threshold (e.g. deconvolution), 0s are converted to the smaller non 0 proportion value.") data_for_model$y_proportion[data_for_model$y_proportion==0] = min(data_for_model$y_proportion[data_for_model$y_proportion>0]) } - + if(inference_method == "hmc"){ - + # tryCatch({ mod$sample( data = data_for_model , @@ -113,16 +124,6 @@ fit_model = function( ... ) - # }, - # error = function(e) { - # - # # I don't know why thi is needed nd why the model sometimes is not compliled correctly - # if(e |> as.character() |> str_detect("Model not compiled")) - # model = load_model(model_name, force=TRUE, threads = cores) - # else - # stop(e) - # - # }) } else{ @@ -171,31 +172,31 @@ get_model_from_data = function(file_compiled_model, model_code){ #' #' This function attempts to load a precompiled Stan model using the `instantiate` package. #' If the model is not found in the cache or force recompilation is requested, it will locate -#' the Stan model file within the `sccomp` package, compile it using `cmdstanr`, and save the +#' the Stan model file within the `sccomp` package, compile it using `cmdstanr`, and save the #' compiled model to the cache directory for future use. #' #' @param name A character string representing the name of the Stan model (without the `.stan` extension). -#' @param cache_dir A character string representing the path to the cache directory where compiled models are saved. +#' @param cache_dir A character string representing the path to the cache directory where compiled models are saved. #' Defaults to `sccomp_stan_models_cache_dir`. -#' @param force A logical value. If `TRUE`, the model will be recompiled even if it exists in the cache. +#' @param force A logical value. If `TRUE`, the model will be recompiled even if it exists in the cache. #' Defaults to `FALSE`. -#' @param threads An integer specifying the number of threads to use for compilation. +#' @param threads An integer specifying the number of threads to use for compilation. #' Defaults to `1`. -# +# #' @return A compiled Stan model object from `cmdstanr`. -#' +#' #' @importFrom instantiate stan_package_model #' @importFrom instantiate stan_package_compile -#' +#' #' @noRd -#' +#' #' @examples #' \donttest{ #' model <- load_model("glm_multi_beta_binomial_", "~/cache", force = FALSE, threads = 1) #' } load_model <- function(name, cache_dir = sccomp_stan_models_cache_dir, force=FALSE, threads = 1) { - - + + # tryCatch({ # # Attempt to load a precompiled Stan model using the instantiate package # instantiate::stan_package_model( @@ -204,22 +205,22 @@ load_model <- function(name, cache_dir = sccomp_stan_models_cache_dir, force=FAL # ) # }, error = function(e) { # Try to load the model from cache - + # Handle cache directory - always add version to ensure version isolation sccomp_version <- as.character(packageVersion("sccomp")) cache_dir <- file.path(cache_dir, sccomp_version) - + # RDS compiled model cache_dir |> dir.create(showWarnings = FALSE, recursive = TRUE) cache_file <- file.path(cache_dir, paste0(name, ".rds")) - + # .STAN raw model stan_model_path <- system.file("stan", paste0(name, ".stan"), package = "sccomp") - + if (file.exists(cache_file) && !force) { mod <- readRDS(cache_file) stan_file <- tryCatch(mod$stan_file(), error = function(e) "") - + if (!is.null(stan_file) && nzchar(stan_file) && file.exists(stan_file)) { message("Loading model from cache...") return(mod) @@ -228,40 +229,40 @@ load_model <- function(name, cache_dir = sccomp_stan_models_cache_dir, force=FAL clear_stan_model_cache(cache_dir = cache_dir) } } - + # If loading the precompiled model fails, find the Stan model file within the package message("Precompiled model not found. Compiling the model...") - + # Compile the Stan model using cmdstanr with threading support enabled instantiate::stan_package_compile( - stan_model_path, + stan_model_path, cpp_options = list(stan_threads = TRUE), - force_recompile = TRUE, - threads = threads, + force_recompile = TRUE, + threads = threads, dir = system.file("stan", package = "sccomp") ) mod = instantiate::stan_package_model( - name = name, - package = "sccomp", + name = name, + package = "sccomp", compile = TRUE, cpp_options = list(stan_threads = TRUE) ) |> suppressWarnings() - + # Save the compiled model object to cache saveRDS(mod, file = cache_file) message("Model compiled and saved to cache successfully.") - + return(mod) # }) - + } #' Check and Install cmdstanr and CmdStan #' -#' This function checks if the `cmdstanr` package (version 0.9.0 or higher) and CmdStan are installed. +#' This function checks if the `cmdstanr` package (version 0.9.0 or higher) and CmdStan are installed. #' If they are not installed, it installs them automatically in non-interactive sessions #' or asks for permission to install them in interactive sessions. -#' +#' #' The function requires cmdstanr version 0.9.0 or higher for support of the new `sum_to_zero_vector` type. #' #' @importFrom instantiate stan_cmdstan_exists @@ -269,13 +270,13 @@ load_model <- function(name, cache_dir = sccomp_stan_models_cache_dir, force=FAL #' @importFrom rlang abort #' @importFrom rlang check_installed #' @return NULL -#' +#' #' @noRd check_and_install_cmdstanr <- function() { - + # Check if cmdstanr is installed # from https://github.com/wlandau/instantiate/blob/33989d74c26f349e292e5efc11c267b3a1b71d3f/R/utils_assert.R#L114 - + # tryCatch( rlang::check_installed( pkg = "cmdstanr", @@ -286,7 +287,7 @@ check_and_install_cmdstanr <- function() { "install.packages(pkgs = \"cmdstanr\",", "repos = c(\"https://mc-stan.org/r-packages/\", getOption(\"repos\"))" ), - + # I have to see if Bioconductor is compatible with this action = function(...) install.packages(..., repos = c('https://stan-dev.r-universe.dev', 'https://cloud.r-project.org')) ) @@ -300,9 +301,9 @@ check_and_install_cmdstanr <- function() { # Check if CmdStan is installed if (!stan_cmdstan_exists()) { - + clear_stan_model_cache() - + stop( "cmdstan is required to proceed.\n\n", "You can install CmdStan by running the following command:\n", @@ -342,10 +343,10 @@ vb_iterative = function(model, additional_parameters_to_save = c(), data, output_dir = output_dir, - seed, + seed, init = "random", inference_method, - cores = 1, + cores = 1, verbose = TRUE, psis_resample = FALSE, sig_figs = 9, @@ -355,9 +356,9 @@ vb_iterative = function(model, i = 0 while (is.null(res) & i < 5) { res = tryCatch({ - + if(inference_method=="pathfinder") - my_res = model |> + my_res = model |> sample_safe( pathfinder_fx, data = data, @@ -365,19 +366,19 @@ vb_iterative = function(model, output_dir = output_dir, seed = seed+i, # init = init, - num_paths=50, + num_paths=50, num_threads = cores, single_path_draws = output_samples / 50 , - max_lbfgs_iters=100, - history_size = 100, + max_lbfgs_iters=100, + history_size = 100, show_messages = verbose, psis_resample = psis_resample, sig_figs = sig_figs, ... ) - + else if(inference_method=="variational") - my_res = model |> + my_res = model |> sample_safe( variational_fx, data = data, @@ -392,7 +393,7 @@ vb_iterative = function(model, sig_figs = sig_figs, ... ) - + boolFalse <- TRUE return(my_res) }, @@ -401,8 +402,8 @@ vb_iterative = function(model, clear_stan_model_cache() model <<- load_model(model_name, force=TRUE, threads = cores, cache_dir = cache_stan_model) } - else writeLines(sprintf("Further attempt with Variational Bayes: %s", e)) - + else writeLines(sprintf("Further attempt with Variational Bayes: %s", e)) + return(NULL) }, finally = { diff --git a/R/plot.R b/R/plot.R index f0e15331..8abc72ef 100644 --- a/R/plot.R +++ b/R/plot.R @@ -2,22 +2,20 @@ #' #' @description This function plots a summary of the results of the model. #' -#' @importFrom ggrepel geom_text_repel -#' @importFrom tidyr pivot_longer -#' @import ggplot2 -#' @importFrom tidyr unite -#' @importFrom tidyr pivot_longer -#' @importFrom dplyr with_groups +#' @importFrom dplyr filter select #' @importFrom magrittr equals #' #' @param x A tibble including a cell_group name column | sample name column | read counts column | factor columns | Pvalue column | a significance column -#' @param significance_threshold Numeric value specifying the significance threshold for highlighting differences. Default is 0.025. +#' @param significance_threshold Numeric value specifying the significance threshold for highlighting differences. Default is 0.05. #' @param test_composition_above_logit_fold_change A positive integer. It is the effect threshold used for the hypothesis test. A value of 0.2 correspond to a change in cell proportion of 10% for a cell type with baseline proportion of 50%. That is, a cell type goes from 45% to 50%. When the baseline proportion is closer to 0 or 1 this effect thrshold has consistent value in the logit uncontrained scale. #' @param significance_statistic Character vector indicating which statistic to highlight. Default is "pH0". #' @param show_fdr_message Logical. Whether to show the Bayesian FDR interpretation message on the plot. Default is TRUE. -#' @param ... For internal use +#' @param add_marginal_density Logical. Whether to add marginal density plots on adjusted panels in 2D intervals. Default is TRUE. +#' @param omit_ci Logical. Whether to omit credible interval error bars from 2D interval plots. Default is FALSE. +#' @param sort_by Character vector indicating how to sort taxa. Options are "none" (default), "effect" (by effect size), "significance" (by FDR/pH0), or "alphabetical". +#' @param ... For internal use #' -#' @return A `ggplot` +#' @return A list containing ggplot objects #' #' @references #' S. Mangiola, A.J. Roth-Schulze, M. Trussart, E. Zozaya-Valdés, M. Ma, Z. Gao, A.F. Rubin, T.P. Speed, H. Shim, & A.T. Papenfuss, sccomp: Robust differential composition and variability analysis for single-cell data, Proc. Natl. Acad. Sci. U.S.A. 120 (33) e2203828120, https://doi.org/10.1073/pnas.2203828120 (2023). @@ -38,31 +36,42 @@ #' counts_obj, #' ~ type, ~1, "sample", "cell_group", "count", #' cores = 1 -#' ) +#' ) |> +#' sccomp_test() #' -#' # estimate |> plot() +#' plots = estimate |> plot() #' } #' } #' -plot.sccomp_tbl <- function(x, significance_threshold = 0.05, test_composition_above_logit_fold_change = .data |> attr("test_composition_above_logit_fold_change"), significance_statistic = c("pH0", "FDR"), show_fdr_message = TRUE, ...) { - +plot.sccomp_tbl <- function( + x, + significance_threshold = 0.05, + test_composition_above_logit_fold_change = x |> attr("test_composition_above_logit_fold_change"), + significance_statistic = c("pH0", "FDR"), + show_fdr_message = TRUE, + add_marginal_density = TRUE, + omit_ci = FALSE, + sort_by = c("none", "effect", "significance", "alphabetical"), + ... +) { + significance_statistic <- match.arg(significance_statistic) - + sort_by <- match.arg(sort_by) + + # Quosures from estimate (same names as legacy `.cell_group` args; stored on tbl attributes) + .sample <- attr(x, ".sample") + .cell_group <- attr(x, ".cell_group") + .count <- attr(x, ".count") + # Define the variables as NULL to avoid CRAN NOTES - parameter <- NULL - count_data <- x |> attr("count_data") v_effect <- NULL - - .cell_group = attr(x, ".cell_group") - .count = attr(x, ".count") - .sample = attr(x, ".sample") - + plots = list() - + # Check if test have been done if(x |> select(ends_with("FDR")) |> ncol() |> equals(0)) stop("sccomp says: to produce plots, you need to run the function sccomp_test() on your estimates.") - + data_proportion = x |> @@ -71,7 +80,8 @@ plot.sccomp_tbl <- function(x, significance_threshold = 0.05, test_composition_ pivot_wider(names_from = parameter, values_from = c(contains("c_"), contains("v_"))) |> left_join( - count_data, by = join_by(!!.cell_group) + attr(x, "count_data"), + by = join_by(!!.cell_group) ) |> with_groups(!!.sample, ~ mutate(.x, proportion = (!!.count)/sum(!!.count)) ) @@ -105,15 +115,13 @@ plot.sccomp_tbl <- function(x, significance_threshold = 0.05, test_composition_ # If variable is continuous if(data_proportion |> select(all_of(.x)) |> pull(1) |> is("numeric")) - my_plot = + my_plot = plot_scatterplot( .data = x, data_proportion = data_proportion, factor_of_interest = .x, - .cell_group = !!.cell_group, - .sample = !!.sample, - my_theme = sccomp_theme(), - significance_threshold = significance_threshold + significance_threshold = significance_threshold, + my_theme = sccomp_theme() ) # If discrete @@ -133,576 +141,30 @@ plot.sccomp_tbl <- function(x, significance_threshold = 0.05, test_composition_ } + # 1D intervals - plots$credible_intervals_1D = plot_1D_intervals( - .data = x, + plots$credible_intervals_1D = sccomp_plot_intervals_1D( + .data = x, significance_threshold = significance_threshold, + test_composition_above_logit_fold_change = test_composition_above_logit_fold_change, significance_statistic = significance_statistic, - show_fdr_message = show_fdr_message + show_fdr_message = show_fdr_message, + sort_by = sort_by ) - - # 2D intervals - if("v_effect" %in% colnames(x) && (x |> filter(!is.na(v_effect)) |> nrow()) > 0) plots$credible_intervals_2D = plot_2D_intervals( - .data = x, - significance_threshold = significance_threshold, - significance_statistic = significance_statistic, - show_fdr_message = show_fdr_message - ) - - plots - -} - -#' Plot 1D Intervals for Cell-group Effects -#' -#' This function creates a series of 1D interval plots for cell-group effects, highlighting significant differences based on a given significance threshold. -#' -#' @param .data Data frame containing the main data. -#' @param significance_threshold Numeric value specifying the significance threshold for highlighting differences. -#' @param test_composition_above_logit_fold_change A positive integer. It is the effect threshold used for the hypothesis test. A value of 0.2 correspond to a change in cell proportion of 10% for a cell type with baseline proportion of 50%. That is, a cell type goes from 45% to 50%. When the baseline proportion is closer to 0 or 1 this effect thrshold has consistent value in the logit uncontrained scale. -#' @param show_fdr_message Logical. Whether to show the Bayesian FDR interpretation message on the plot. Default is TRUE. -#' @param significance_statistic Character vector indicating which statistic to highlight. Default is "pH0". -#' @importFrom patchwork wrap_plots -#' @importFrom forcats fct_reorder -#' @importFrom tidyr drop_na -#' -#' @export -#' -#' @return A combined plot of 1D interval plots. -#' @examples -#' -#' print("cmdstanr is needed to run this example.") -#' -#' \donttest{ -#' if (instantiate::stan_cmdstan_exists()) { -#' data("counts_obj") -#' -#' estimate <- sccomp_estimate( -#' counts_obj, -#' ~ type, -#' ~1, -#' "sample", -#' "cell_group", -#' "count", -#' cores = 1 -#' ) |> -#' sccomp_test() -#' -#' # Example usage: -#' my_plot = plot_1D_intervals(estimate) -#' -#' } -#' } -#' -#' -plot_1D_intervals = function( - .data, - significance_threshold = 0.05, - test_composition_above_logit_fold_change = .data |> attr("test_composition_above_logit_fold_change"), - show_fdr_message = TRUE, - significance_statistic = c("pH0", "FDR") -) { - significance_statistic <- match.arg(significance_statistic) - - # Define the variables as NULL to avoid CRAN NOTES - parameter <- NULL - estimate <- NULL - value <- NULL - pH0 <- NULL - FDR <- NULL - - .cell_group = attr(.data, ".cell_group") - - # Check if test have been done - if(.data |> select(ends_with("FDR")) |> ncol() |> equals(0)) - stop("sccomp says: to produce plots, you need to run the function sccomp_test() on your estimates.") - - plot_list = - .data |> - filter(parameter != "(Intercept)") |> - - # Reshape data - select(-contains("n_eff"), -contains("R_k_hat"), -contains("_ess"), -contains("_rhat")) |> - pivot_longer(c(contains("c_"), contains("v_")), names_sep = "_", names_to = c("which", "estimate")) |> - pivot_wider(names_from = estimate, values_from = value) |> - - # Nest data by parameter and which - nest(data = -c(parameter, which)) |> - mutate(plot = pmap( - list(data, which, parameter), - ~ { - plot_data <- ..1 - # Check if there are any statistics to plot - if(plot_data |> filter(!effect |> is.na()) |> nrow() |> equals(0)) - return(NA) - - # Choose color variable and legend - if (significance_statistic == "FDR") { - color_var <- plot_data$FDR < significance_threshold - color_aes <- aes(xmin = lower, xmax = upper, color = FDR < significance_threshold) - color_scale <- scale_color_manual(values = c("grey40", "red")) - legend_title <- "FDR < significance_threshold" - } else { - color_var <- plot_data$pH0 < significance_threshold - color_aes <- aes(xmin = lower, xmax = upper, color = pH0 < significance_threshold) - color_scale <- scale_color_manual(values = c("grey40", "red")) - legend_title <- "pH0 < significance_threshold" - } - - ggplot(plot_data, aes(x = effect, y = fct_reorder(!!.cell_group, effect))) + - geom_vline(xintercept = test_composition_above_logit_fold_change, colour = "grey") + - geom_vline(xintercept = -test_composition_above_logit_fold_change, colour = "grey") + - geom_errorbar(color_aes) + - geom_point() + - color_scale + - xlab("Credible interval of the slope") + - ylab("Cell group") + - ggtitle(sprintf("%s %s", ..2, ..3)) + - sccomp_theme() + - theme(legend.position = "bottom") + - guides(color = guide_legend(title = legend_title)) - } - )) %>% - - # Filter out NA plots - filter(!plot |> is.na()) |> - pull(plot) - - # Combine all individual plots into one plot - combined_plot <- plot_list |> - wrap_plots(ncol = plot_list |> length() |> sqrt() |> ceiling()) - - # Only show the FDR message if significance_statistic == "FDR" and show_fdr_message is TRUE - if (significance_statistic == "FDR" && show_fdr_message) { - caption_text <- paste( - "Bayesian FDR: Stephens' method (doi: 10.1093/biostatistics/kxw041)", - "\nFDR-significant populations may cross fold change thresholds because Bayesian FDR considers posterior probabilities rather than p-values.", - "\nThe method sorts null hypothesis probabilities in ascending order and calculates cumulative averages for robust false discovery control.", - sep = "" - ) - combined_plot <- combined_plot + patchwork::plot_annotation( - caption = caption_text - ) - } - combined_plot -} - - -#' Plot 2D Intervals for Mean-Variance Association -#' -#' This function creates a 2D interval plot for mean-variance association, highlighting significant differences based on a given significance threshold. -#' -#' @param .data Data frame containing the main data. -#' @param significance_threshold Numeric value specifying the significance threshold for highlighting differences. Default is 0.025. -#' @param test_composition_above_logit_fold_change A positive integer. It is the effect threshold used for the hypothesis test. A value of 0.2 correspond to a change in cell proportion of 10% for a cell type with baseline proportion of 50%. That is, a cell type goes from 45% to 50%. When the baseline proportion is closer to 0 or 1 this effect thrshold has consistent value in the logit uncontrained scale. -#' @param show_fdr_message Logical. Whether to show the Bayesian FDR interpretation message on the plot. Default is TRUE. -#' @param significance_statistic Character vector indicating which statistic to highlight. Default is "pH0". -#' -#' -#' @importFrom dplyr filter arrange mutate if_else row_number -#' @importFrom ggplot2 ggplot geom_vline geom_hline geom_errorbar geom_point annotate aes facet_wrap -#' @importFrom ggrepel geom_text_repel -#' @importFrom scales trans_new -#' @importFrom stringr str_replace -#' @importFrom stats quantile -#' @importFrom magrittr equals -#' -#' @export -#' -#' @return A ggplot object representing the 2D interval plot. -#' -#' @examples -#' -#' print("cmdstanr is needed to run this example.") -#' -#' \donttest{ -#' if (instantiate::stan_cmdstan_exists()) { -#' data("counts_obj") -#' -#' estimate <- sccomp_estimate( -#' counts_obj, -#' ~ type, -#' ~type, -#' "sample", -#' "cell_group", -#' "count", -#' cores = 1 -#' ) |> -#' sccomp_test() -#' -#' # Example usage: -#' my_plot = plot_2D_intervals(estimate) -#' -#' } -#' } -#' -plot_2D_intervals = function( - .data, - significance_threshold = 0.05, - test_composition_above_logit_fold_change = - .data |> attr("test_composition_above_logit_fold_change"), - show_fdr_message = TRUE, - significance_statistic = c("pH0", "FDR") -){ - significance_statistic <- match.arg(significance_statistic) - - # Define the variables as NULL to avoid CRAN NOTES - v_effect <- NULL - parameter <- NULL - c_effect <- NULL - c_lower <- NULL - c_upper <- NULL - c_FDR <- NULL - v_lower <- NULL - v_upper <- NULL - v_FDR <- NULL - cell_type_label <- NULL - pH0 <- NULL - FDR <- NULL - - .cell_group = attr(.data, ".cell_group") - - # Check if test have been done - if(.data |> select(ends_with("FDR")) |> ncol() |> equals(0)) - stop("sccomp says: to produce plots, you need to run the function sccomp_test() on your estimates.") - - # Extract prec_coeff parameters from the fitted model for regression line - fit = attr(.data, "fit") - prec_coeff_summary = fit$summary("prec_coeff") - prec_coeff_intercept = prec_coeff_summary$mean[1] - prec_coeff_slope = prec_coeff_summary$mean[2] - - # Add adjusted intercept (unbiased variability) as a new facet - .data_adjusted <- .data %>% - filter(parameter == "(Intercept)") %>% - mutate( - v_effect = v_effect + prec_coeff_slope * c_effect, - v_lower = v_lower + prec_coeff_slope * c_effect, - v_upper = v_upper + prec_coeff_slope * c_effect, - parameter = "(Intercept, adjusted)" - ) - - # Bind adjusted data to original (uncorrected comes first) - .data_plot <- bind_rows(.data, .data_adjusted) - - # Always set parameter factor levels to only those present in the data - if ("parameter" %in% colnames(.data_plot)) { - .data_plot$parameter <- factor(.data_plot$parameter, levels = unique(.data_plot$parameter)) - } - - # Use .data_plot instead of .data in the rest of the function - plot <- .data_plot %>% - # Filter where variance is inferred - filter(!is.na(v_effect)) %>% - - # Add labels for significant cell groups - with_groups( - parameter, - ~ .x %>% - arrange(c_FDR) %>% - mutate(cell_type_label = if_else(row_number() <= 3 & c_FDR < significance_threshold & !parameter %in% c("(Intercept)", "(Intercept, adjusted)"), !!.cell_group, "")) - ) %>% - with_groups( - parameter, - ~ .x %>% - arrange(v_FDR) %>% - mutate(cell_type_label = if_else((row_number() <= 3 & v_FDR < significance_threshold & !parameter %in% c("(Intercept)", "(Intercept, adjusted)") ), !!.cell_group, cell_type_label)) - ) %>% - { - .x = (.) - - # Choose color variable and legend - if (significance_statistic == "FDR") { - color_c_aes <- aes(xmin = c_lower, xmax = c_upper, color = c_FDR < significance_threshold & !parameter %in% c("(Intercept)", "(Intercept, adjusted)"), alpha = c_FDR < significance_threshold & !parameter %in% c("(Intercept)", "(Intercept, adjusted)")) - color_v_aes <- aes(ymin = v_lower, ymax = v_upper, color = v_FDR < significance_threshold & !parameter %in% c("(Intercept)", "(Intercept, adjusted)"), alpha = v_FDR < significance_threshold & !parameter %in% c("(Intercept)", "(Intercept, adjusted)")) - color_scale <- scale_color_manual(values = c("#D3D3D3", "#E41A1C")) - alpha_scale <- scale_alpha_manual(values = c(0.4, 1)) - legend_title <- "FDR < significance_threshold" - } else { - color_c_aes <- aes(xmin = c_lower, xmax = c_upper, color = c_pH0 < significance_threshold & !parameter %in% c("(Intercept)", "(Intercept, adjusted)"), alpha = c_pH0 < significance_threshold & !parameter %in% c("(Intercept)", "(Intercept, adjusted)")) - color_v_aes <- aes(ymin = v_lower, ymax = v_upper, color = v_pH0 < significance_threshold & !parameter %in% c("(Intercept)", "(Intercept, adjusted)"), alpha = v_pH0 < significance_threshold & !parameter %in% c("(Intercept)", "(Intercept, adjusted)")) - color_scale <- scale_color_manual(values = c("#D3D3D3", "#377EB8")) - alpha_scale <- scale_alpha_manual(values = c(0.4, 1)) - legend_title <- "pH0 < significance_threshold" - } - - # Calculate range for regression line - c_range = range(.x$c_effect, na.rm = TRUE) - c_seq = seq(c_range[1], c_range[2], length.out = 100) - v_pred = -(prec_coeff_intercept + prec_coeff_slope * c_seq) - - # Create regression line data only for (Intercept) - regression_data = data.frame( - c_effect = c_seq, - v_effect = v_pred, - parameter = "(Intercept)" - ) - - # Plot - p <- ggplot(.x, aes(c_effect, v_effect)) + - # Add vertical and horizontal lines - geom_vline(xintercept = c(-test_composition_above_logit_fold_change, test_composition_above_logit_fold_change), colour = "grey", linetype = "dashed", linewidth = 0.3) + - geom_hline(yintercept = c(-test_composition_above_logit_fold_change, test_composition_above_logit_fold_change), colour = "grey", linetype = "dashed", linewidth = 0.3) - - # Add regression line only for (Intercept) facet - p <- p + geom_line( - data = regression_data, - mapping = aes(c_effect, v_effect), - color = "#0072B2", linewidth = 0.5, alpha = 0.8, - inherit.aes = FALSE - ) - # Add horizontal line for (Intercept, adjusted) facet - if ("(Intercept, adjusted)" %in% unique(.x$parameter)) { - mean_adjusted <- mean(.x$v_effect[.x$parameter == "(Intercept, adjusted)"]) - c_range_adjusted <- range(.x$c_effect[.x$parameter == "(Intercept, adjusted)"], na.rm = TRUE) - - # Create horizontal line data for adjusted intercept only - horizontal_line_data <- data.frame( - c_effect = c_range_adjusted, - v_effect = rep(mean_adjusted, 2), - parameter = "(Intercept, adjusted)" - ) - - p <- p + - geom_line( - data = horizontal_line_data, - mapping = aes(c_effect, v_effect), - color = "#0072B2", linewidth = 0.5, alpha = 0.8, - inherit.aes = FALSE - ) - } - - p <- p + - # Add error bars - geom_errorbar(color_c_aes, linewidth = 0.2) + - geom_errorbar(color_v_aes, linewidth = 0.2) + - # Add points - geom_point(size = 0.2) + - # Add annotations - # annotate("text", x = 0, y = 3.5, label = "Variability", size = 2) + # Disabled temporarily; revisit if annotations are needed for variability. - # annotate("text", x = 5, y = 0, label = "Abundance", size = 2, angle = 270) + # Disabled temporarily; revisit if annotations are needed for abundance. - # Add text labels for significant cell groups - geom_text_repel(aes(c_effect, -v_effect, label = cell_type_label), size = 2.5, data = .x %>% filter(cell_type_label != "")) + - # Set color and alpha scales - color_scale + - alpha_scale + - # Facet by parameter - facet_wrap( - ~ fct_relevel(parameter, c("(Intercept)", "(Intercept, adjusted)")) |> - - # I have to understand why this works but throws warning - suppressWarnings(), - scales = "free" - ) + - xlab("c_effect (Abundance effect)") + - ylab("v_effect (Variability effect)") + - # Apply custom theme - sccomp_theme() + - theme(legend.position = "bottom") + - guides(color = guide_legend(title = legend_title), alpha = "none") - p - } - - # Only show the FDR message if significance_statistic == "FDR" and show_fdr_message is TRUE - if (significance_statistic == "FDR" && show_fdr_message) { - caption_text <- paste( - "Bayesian FDR: Stephens' method (doi: 10.1093/biostatistics/kxw041)", - "\nFDR-significant populations may cross fold change thresholds because Bayesian FDR considers posterior probabilities rather than p-values.", - "\nThe method sorts null hypothesis probabilities in ascending order and calculates cumulative averages for robust false discovery control.", - sep = "" + # 2D intervals (only if variance effects exist) + if("v_effect" %in% colnames(x) && (x |> filter(!is.na(v_effect)) |> nrow()) > 0) { + plots$credible_intervals_2D = sccomp_plot_intervals_2D( + .data = x, + significance_threshold = significance_threshold, + test_composition_above_logit_fold_change = test_composition_above_logit_fold_change, + significance_statistic = significance_statistic, + show_fdr_message = show_fdr_message, + add_marginal_density = add_marginal_density, + omit_ci = omit_ci ) - plot <- plot + ggplot2::labs(caption = caption_text) - plot <- plot + theme(plot.caption = ggplot2::element_text(hjust = 0)) } - plot -} - - + plots -#' Plot Scatterplot of Cell-group Proportion -#' -#' This function creates a scatterplot of cell-group proportions, optionally highlighting significant differences based on a given significance threshold. -#' -#' @param .data Data frame containing the main data. -#' @param data_proportion Data frame containing proportions of cell groups. -#' @param factor_of_interest A factor indicating the biological condition of interest. -#' @param .cell_group The cell group to be analysed. -#' @param .sample The sample identifier. -#' @param significance_threshold Numeric value specifying the significance threshold for highlighting differences. Default is 0.025. -#' @param my_theme A ggplot2 theme object to be applied to the plot. -#' @importFrom scales trans_new -#' @importFrom stringr str_replace -#' @importFrom stats quantile -#' @importFrom magrittr equals -#' -#' -#' @return A ggplot object representing the scatterplot. -#' @examples -#' # Example usage: -#' # plot_scatterplot(.data, data_proportion, "condition", "cell_group", "sample", 0.025, theme_minimal()) -plot_scatterplot = function( - .data, data_proportion, factor_of_interest, .cell_group, - .sample, significance_threshold = 0.05, my_theme -){ - - # Define the variables as NULL to avoid CRAN NOTES - stats_name <- NULL - parameter <- NULL - stats_value <- NULL - count_data <- NULL - generated_proportions <- NULL - proportion <- NULL - name <- NULL - outlier <- NULL - - # Function to remove leading zero from labels - dropLeadingZero <- function(l){ stringr::str_replace(l, '0(?=.)', '') } - - # Define square root transformation and its inverse - S_sqrt <- function(x){sign(x)*sqrt(abs(x))} - IS_sqrt <- function(x){x^2*sign(x)} - S_sqrt_trans <- function() scales::trans_new("S_sqrt",S_sqrt,IS_sqrt) - - .cell_group = enquo(.cell_group) - .sample = enquo(.sample) - - # Prepare significance colors - significance_colors = - .data %>% - pivot_longer( - c(contains("c_"), contains("v_")), - names_pattern = "([cv])_([a-zA-Z0-9]+)", - names_to = c("which", "stats_name"), - values_to = "stats_value" - ) %>% - filter(stats_name == "FDR") %>% - filter(parameter != "(Intercept)") %>% - filter(stats_value < significance_threshold) %>% - filter(`factor` == factor_of_interest) - - if(nrow(significance_colors) > 0){ - - if(.data |> attr("contrasts") |> is.null()) - significance_colors = - significance_colors %>% - unite("name", c(which, parameter), remove = FALSE) %>% - distinct() %>% - - # Get clean parameter - mutate(!!as.symbol(factor_of_interest) := str_replace(parameter, sprintf("^%s", `factor`), "")) %>% - with_groups(c(!!.cell_group, !!as.symbol(factor_of_interest)), ~ .x %>% summarise(name = paste(name, collapse = ", "))) - else - significance_colors = - significance_colors |> - mutate( - factor_values = attr(.data, "count_data") |> - select(all_of(factor_of_interest)) |> - distinct() |> - pull(all_of(factor_of_interest)) - ) |> - unnest(factor_values) |> - - # Filter relevant parameters - mutate( !!as.symbol(factor_of_interest) := as.character(factor_values) ) |> - filter(str_detect(parameter, !!as.symbol(factor_of_interest) )) |> - - # Rename - select(!!.cell_group, !!as.symbol(factor_of_interest), name = parameter) |> - - # Merge contrasts - with_groups(c(!!.cell_group, !!as.symbol(factor_of_interest)), ~ .x %>% summarise(name = paste(name, collapse = ", "))) - } - - my_scatterplot = ggplot() - - if("fit" %in% names(attributes(.data))){ - - simulated_proportion = - .data |> - sccomp_replicate(number_of_draws = 1000) |> - left_join(data_proportion %>% distinct(!!as.symbol(factor_of_interest), !!.sample, !!.cell_group)) - - my_scatterplot = - my_scatterplot + - - # Add smoothed line for simulated proportions - geom_smooth( - aes(!!as.symbol(factor_of_interest), (generated_proportions)), - lwd=0.2, - data = - simulated_proportion %>% - inner_join(data_proportion %>% distinct(!!as.symbol(factor_of_interest), !!.cell_group, !!.sample)) , - color="blue", fill="blue", - span = 1 - ) - } - - if( - nrow(significance_colors)==0 || - - significance_colors |> - pull(!!as.symbol(factor_of_interest)) |> - intersect( - data_proportion |> - pull(!!as.symbol(factor_of_interest)) - ) |> - length() |> - equals(0) - ) { - - my_scatterplot= - my_scatterplot + - - # Add smoothed line without significance colors - geom_smooth( - aes(!!as.symbol(factor_of_interest), proportion, fill = NULL), - data = - data_proportion , - lwd=0.5, - color = "black", - span = 1 - ) - } else { - my_scatterplot= - my_scatterplot + - - # Add smoothed line with significance colors - geom_smooth( - aes(!!as.symbol(factor_of_interest), proportion, fill = name), - data = data_proportion , - fatten = 0.5, - lwd=0.5, - color = "black", - span = 1 - ) - } - - my_scatterplot + - - # Add jittered points for individual data - geom_point( - aes(!!as.symbol(factor_of_interest), proportion, shape=outlier, color=outlier), - data = data_proportion, - position=position_jitterdodge(jitter.height = 0, jitter.width = 0.2), - size = 0.5 - ) + - - # Facet wrap by cell group - facet_wrap( - vars(!!.cell_group), - scales = "free_y", - nrow = 4 - ) + - scale_color_manual(values = c("black", "#e11f28")) + - scale_y_continuous(trans=S_sqrt_trans(), labels = dropLeadingZero) + - scale_fill_discrete(na.value = "white") + - xlab("Biological condition") + - ylab("Cell-group proportion") + - guides(color="none", alpha="none", size="none") + - labs(fill="Significant difference") + - ggtitle("Note: Be careful judging significance (or outliers) visually for lowly abundant cell groups. \nVisualising proportion hides the uncertainty characteristic of count data, that a count-based statistical model can estimate.") + - my_theme + - theme(axis.text.x = element_text(angle=20, hjust = 1), title = element_text(size = 3)) } diff --git a/R/plot_scatterplot.R b/R/plot_scatterplot.R new file mode 100644 index 00000000..d09b73e1 --- /dev/null +++ b/R/plot_scatterplot.R @@ -0,0 +1,183 @@ + +#' Plot Scatterplot of Cell-group Proportion +#' +#' This function creates a scatterplot of cell-group proportions, optionally +#' highlighting significant differences based on a given significance threshold. +#' +#' @param .data Data frame containing the main data. +#' @param data_proportion Data frame containing proportions of cell groups. +#' @param factor_of_interest A factor indicating the biological condition of interest. +#' @param significance_threshold Numeric value specifying the significance threshold +#' for highlighting differences. Default is 0.05. +#' @param my_theme A ggplot2 theme object to be applied to the plot. +#' @importFrom scales trans_new +#' @importFrom stringr str_replace str_detect +#' +#' @return A ggplot object representing the scatterplot. +#' +#' @noRd +plot_scatterplot = function( + .data, data_proportion, factor_of_interest, + significance_threshold = 0.05, my_theme +){ + + # Define the variables as NULL to avoid CRAN NOTES + stats_name <- NULL + parameter <- NULL + stats_value <- NULL + count_data <- NULL + generated_proportions <- NULL + proportion <- NULL + name <- NULL + outlier <- NULL + + # Function to remove leading zero from labels + dropLeadingZero <- function(l){ stringr::str_replace(l, '0(?=.)', '') } + + # Define square root transformation and its inverse + S_sqrt <- function(x){sign(x)*sqrt(abs(x))} + IS_sqrt <- function(x){x^2*sign(x)} + S_sqrt_trans <- function() scales::trans_new("S_sqrt",S_sqrt,IS_sqrt) + + .cell_group = attr(.data, ".cell_group") + .count = attr(.data, ".count") + .sample = attr(.data, ".sample") + + # Prepare significance colors + significance_colors = + .data %>% + pivot_longer( + c(contains("c_"), contains("v_")), + names_pattern = "([cv])_([a-zA-Z0-9]+)", + names_to = c("which", "stats_name"), + values_to = "stats_value" + ) %>% + filter(stats_name == "FDR") %>% + filter(parameter != "(Intercept)") %>% + filter(stats_value < significance_threshold) %>% + filter(`factor` == factor_of_interest) + + if(nrow(significance_colors) > 0){ + + if(.data |> attr("contrasts") |> is.null()) + significance_colors = + significance_colors %>% + unite("name", c(which, parameter), remove = FALSE) %>% + distinct() %>% + + # Get clean parameter + mutate(!!as.symbol(factor_of_interest) := str_replace(parameter, sprintf("^%s", `factor`), "")) %>% + with_groups(c(!!.cell_group, !!as.symbol(factor_of_interest)), ~ .x %>% summarise(name = paste(name, collapse = ", "))) + else + significance_colors = + significance_colors |> + mutate( + factor_values = attr(.data, "count_data") |> + select(all_of(factor_of_interest)) |> + distinct() |> + pull(all_of(factor_of_interest)) + ) |> + unnest(factor_values) |> + + # Filter relevant parameters + mutate( !!as.symbol(factor_of_interest) := as.character(factor_values) ) |> + filter(str_detect(parameter, !!as.symbol(factor_of_interest) )) |> + + # Rename + select(!!.cell_group, !!as.symbol(factor_of_interest), name = parameter) |> + + # Merge contrasts + with_groups(c(!!.cell_group, !!as.symbol(factor_of_interest)), ~ .x %>% summarise(name = paste(name, collapse = ", "))) + } + + my_scatterplot = ggplot() + + if("fit" %in% names(attributes(.data))){ + + simulated_proportion = + .data |> + sccomp_replicate(number_of_draws = 1000) |> + left_join(data_proportion %>% distinct(!!as.symbol(factor_of_interest), !!.sample, !!.cell_group)) + + my_scatterplot = + my_scatterplot + + + # Add smoothed line for simulated proportions + geom_smooth( + aes(!!as.symbol(factor_of_interest), (generated_proportions)), + lwd=0.2, + data = + simulated_proportion %>% + inner_join(data_proportion %>% distinct(!!as.symbol(factor_of_interest), !!.cell_group, !!.sample)) , + color="blue", fill="blue", + span = 1 + ) + } + + if( + nrow(significance_colors)==0 || + + significance_colors |> + pull(!!as.symbol(factor_of_interest)) |> + intersect( + data_proportion |> + pull(!!as.symbol(factor_of_interest)) + ) |> + length() |> + equals(0) + ) { + + my_scatterplot= + my_scatterplot + + + # Add smoothed line without significance colors + geom_smooth( + aes(!!as.symbol(factor_of_interest), proportion, fill = NULL), + data = + data_proportion , + lwd=0.5, + color = "black", + span = 1 + ) + } else { + my_scatterplot= + my_scatterplot + + + # Add smoothed line with significance colors + geom_smooth( + aes(!!as.symbol(factor_of_interest), proportion, fill = name), + data = data_proportion , + linewidth = 0.5, + lwd=0.5, + color = "black", + span = 1 + ) + } + + my_scatterplot + + + # Add jittered points for individual data + geom_point( + aes(!!as.symbol(factor_of_interest), proportion, shape=outlier, color=outlier), + data = data_proportion, + position=position_jitterdodge(jitter.height = 0, jitter.width = 0.2), + size = 0.5 + ) + + + # Facet wrap by cell group + facet_wrap( + vars(!!.cell_group), + scales = "free_y", + nrow = 4 + ) + + scale_color_manual(values = c("black", "#e11f28")) + + scale_y_continuous(trans=S_sqrt_trans(), labels = dropLeadingZero) + + scale_fill_discrete(na.value = "white") + + xlab("Biological condition") + + ylab("Cell-group proportion") + + guides(color="none", alpha="none", size="none") + + labs(fill="Significant difference") + + ggtitle("Note: Be careful judging significance (or outliers) visually for lowly abundant cell groups. \nVisualising proportion hides the uncertainty characteristic of count data, that a count-based statistical model can estimate.") + + my_theme + + theme(axis.text.x = element_text(angle=20, hjust = 1), title = element_text(size = 3)) +} diff --git a/R/sccomp-package.R b/R/sccomp-package.R index 31493fbb..ab6af3fd 100644 --- a/R/sccomp-package.R +++ b/R/sccomp-package.R @@ -19,7 +19,7 @@ #' \item \code{\link{sccomp_predict}} - Predict cell type proportions for new samples #' \item \code{\link{sccomp_remove_unwanted_effects}} - Remove unwanted variation from data #' \item \code{\link{sccomp_proportional_fold_change}} - Calculate proportional fold changes -#' \item Plotting functions: \code{\link{plot.sccomp_tbl}}, \code{\link{sccomp_boxplot}}, \code{\link{plot_1D_intervals}}, \code{\link{plot_2D_intervals}} +#' \item Plotting functions: \code{\link{plot.sccomp_tbl}}, \code{\link{sccomp_boxplot}}, \code{\link{sccomp_plot_intervals_1D}}, \code{\link{sccomp_plot_intervals_2D}} #' } #' #' For detailed information on usage, see the package vignettes: diff --git a/R/sccomp_boxplot.R b/R/sccomp_boxplot.R index df1b645b..600b968b 100644 --- a/R/sccomp_boxplot.R +++ b/R/sccomp_boxplot.R @@ -71,11 +71,13 @@ sccomp_boxplot = function( if(.data |> select(any_of(c(paste0("c_", selected_statistic), paste0("v_", selected_statistic)))) |> ncol() |> equals(0)) stop("sccomp says: to produce plots, you need to run the function sccomp_test() on your estimates.") + .data_filtered = + subset_results_by_factor(.data, factor, keep_intercept = FALSE) + data_proportion = - .data |> + .data_filtered |> # Otherwise does not work - filter(factor == !!factor) |> select(-`factor`) data_proportion = @@ -281,7 +283,7 @@ plot_boxplot = function( aes(!!as.symbol(factor_of_interest), generated_proportions), fun.data = calc_boxplot_stat, geom = "boxplot", - fatten = 0.5, + median.linewidth = 0.5, lwd = 0.2, data = simulated_proportion %>% inner_join(data_proportion %>% distinct(!!as.symbol(factor_of_interest), !!.cell_group)), @@ -305,7 +307,7 @@ plot_boxplot = function( data = data_proportion |> mutate(!!as.symbol(factor_of_interest) := as.character(!!as.symbol(factor_of_interest))) , - fatten = 0.5, + median.linewidth = 0.5, lwd=0.5 ) } else { @@ -319,7 +321,7 @@ plot_boxplot = function( data_proportion |> mutate(!!as.symbol(factor_of_interest) := as.character(!!as.symbol(factor_of_interest))) %>% left_join(significance_colors, by = c(quo_name(.cell_group), factor_of_interest)), - fatten = 0.5, + median.linewidth = 0.5, lwd=0.5 ) } diff --git a/R/sccomp_estimate.R b/R/sccomp_estimate.R index 9a716d5e..1a326688 100644 --- a/R/sccomp_estimate.R +++ b/R/sccomp_estimate.R @@ -34,7 +34,12 @@ #' @param cores Number of cores to use for parallel calculations. #' @param bimodal_mean_variability_association Logical, whether to model mean-variability as bimodal. #' @param prior_mean A list specifying prior knowledge about the mean distribution, including intercept and coefficients. -#' @param prior_overdispersion_mean_association A list specifying prior knowledge about mean/variability association. +#' @param prior_overdispersion_mean_association A named list with numeric length-2 vectors +#' `intercept`, `slope`, and `standard_deviation` passed to the Stan Student-t / Normal +#' hyperpriors on `prec_intercept`, `prec_slope`, and `log_prec_sd`. Use `NULL` for +#' package defaults. A scalar logical such as `FALSE` is not meaningful here and is treated +#' as `NULL` after a message; to disable abundance dependence in the variability prior, +#' use `exclude_mean_variability_association = TRUE` instead. #' @param percent_false_positive A real number between 0 and 100 for outlier identification. #' @param inference_method Character string specifying the inference method to use ('pathfinder', 'hmc', or 'variational'). Replaces the deprecated `approximate_posterior_inference` and `variational_inference`. #' @param .sample_cell_group_pairs_to_exclude A column name indicating sample/cell-group pairs to exclude. @@ -42,7 +47,11 @@ #' @param verbose Logical, whether to print progression details. #' @param enable_loo Logical, whether to enable model comparison using the LOO package. #' @param noise_model A character string specifying the noise model (e.g., 'multi_beta_binomial'). -#' @param exclude_priors Logical, whether to run a prior-free model. +#' @param exclude_mean_variability_association Logical. When `TRUE`, the prior +#' on the variability parameters does not depend on the abundance: the +#' mean-variability regression is reduced to an intercept-only Normal (or a +#' two-component mixture when `bimodal_mean_variability_association = TRUE`) +#' while the rest of the hierarchical prior structure stays unchanged. #' @param use_data Logical, whether to run the model data-free. #' @param mcmc_seed An integer seed for MCMC reproducibility. #' @param max_sampling_iterations Integer to limit the maximum number of iterations for large datasets. @@ -51,30 +60,23 @@ #' @param cache_stan_model A character string specifying the cache directory for compiled Stan models. #' The sccomp version will be automatically appended to ensure version isolation. #' Default is `sccomp_stan_models_cache_dir` which points to `~/.sccomp_models`. -#' @param portable Logical, whether to keep the result **self-contained** after fitting. Default `TRUE`. -#' \describe{ -#' \item{`TRUE`}{After fitting, sccomp calls `incorporate_parameters_into_sccomp_object()` to pull needed draws into the fit, -#' then **deletes** the Stan CSV files under `output_directory` on **this** machine (only files that still exist are removed).} -#' \item{`FALSE`}{sccomp **does not** delete draw files; CSVs stay on disk. The fit object stores **absolute paths** to those files. -#' You must keep that directory alongside the object (or stay on the same filesystem) if you later need variables that -#' were never loaded in this session.} -#' } -#' **Efficiency:** At `sccomp_estimate()` time sccomp only asks cmdstanr for **summaries** of `beta`, `alpha_normalised`, -#' and random-effect blocks needed for the printed table, not full `draws()` tensors for every Stan parameter. CmdStanR still -#' reads output files to compute those summaries (how much is cached in RAM is implementation-dependent). -#' **Relocating the object:** If you save the tibble and open it elsewhere without the CSV directory (or paths point to another -#' host), `attr(..., "fit")$summary()` can fail for variables that were never touched during estimate—for example the -#' table uses **`alpha_normalised`**, not raw `alpha`, so `summary("alpha[1,1]")` may try to re-read CSVs whereas -#' `summary("beta[1,1]")` may still work if that block was already loaded. Use `portable = TRUE` before archiving, or copy -#' the whole draw folder so paths resolve. -#' If you delete CSVs manually with `portable = FALSE`, call `sccomp_test()` before deletion, run `incorporate_parameters_into_sccomp_object()` first, -#' or switch to `portable = TRUE`; otherwise `sccomp_test()` errors when recorded output paths are missing unless draws were incorporated. +#' @param portable Logical, whether to keep the result portable by caching required draws in memory and removing Stan draw CSV files after fitting. +#' Default is TRUE to save disk space and move needed values into memory. Set to FALSE to keep draw CSV files on disk. +#' With `portable = FALSE`, CSVs remain for you to inspect or archive, but cmdstanr typically **still holds posterior draws in RAM** +#' after fitting and summarisation (`fit$summary()`). The printed estimate table only calls `fit$summary()` on composition +#' (`beta`, …) and variability (`alpha`, …), not on every saved parameter (e.g. `prec_sd`), yet cmdstanr still +#' exposes all saved parameters from memory once output has been read, so `fit$draws(variables = "prec_sd")` can work +#' after CSV deletion just like `beta`. Deleting CSVs does **not** reliably invalidate the fit in the same R session. +#' Call `sccomp_test()` before deleting draw files, use `portable = TRUE` (draws cached then files removed), or run +#' `incorporate_parameters_into_sccomp_object()` before deletion if you remove files manually. `sccomp_test()` stops with a +#' clear error when recorded Stan output paths are missing unless draws were incorporated for portability as above. #' @param .count **DEPRECATED**. Use `abundance` instead. #' @param approximate_posterior_inference **DEPRECATED**. Use `inference_method` instead. #' @param variational_inference **DEPRECATED**. Use `inference_method` instead. #' @param .sample **DEPRECATED**. Use `sample` instead. #' @param .cell_group **DEPRECATED**. Use `cell_group` instead. #' @param .abundance **DEPRECATED**. Use `abundance` instead. +#' @param exclude_priors **DEPRECATED**. Use `exclude_mean_variability_association` instead. #' @param ... Additional arguments passed to the `cmdstanr::sample` function. #' #' @return A tibble (`tbl`), with the following columns: @@ -162,16 +164,16 @@ sccomp_estimate <- function(.data, inference_method = "pathfinder", prior_mean = list(intercept = c(0, 1), coefficients = c(0, 1)), prior_overdispersion_mean_association = list( - intercept = c(5, 2), - slope = c(0, 0.6), - standard_deviation = c(10, 20) + intercept = c(4, 2), + slope = c(0, 2), + standard_deviation = c(1, 0.5) ), .sample_cell_group_pairs_to_exclude = NULL, output_directory = "sccomp_draws_files", verbose = TRUE, enable_loo = FALSE, noise_model = "multi_beta_binomial", - exclude_priors = FALSE, + exclude_mean_variability_association = FALSE, use_data = TRUE, mcmc_seed = sample_seed(), max_sampling_iterations = 20000, @@ -187,7 +189,8 @@ sccomp_estimate <- function(.data, variational_inference = NULL, .sample = NULL, .cell_group = NULL, - .abundance = NULL) { + .abundance = NULL, + exclude_priors = NULL) { # Run the function check_and_install_cmdstanr() @@ -255,16 +258,16 @@ sccomp_estimate.Seurat <- function(.data, inference_method = "pathfinder", prior_mean = list(intercept = c(0, 1), coefficients = c(0, 1)), prior_overdispersion_mean_association = list( - intercept = c(5, 2), - slope = c(0, 0.6), - standard_deviation = c(10, 20) + intercept = c(4, 2), + slope = c(0, 2), + standard_deviation = c(1, 0.5) ), .sample_cell_group_pairs_to_exclude = NULL, output_directory = "sccomp_draws_files", verbose = TRUE, enable_loo = FALSE, noise_model = "multi_beta_binomial", - exclude_priors = FALSE, + exclude_mean_variability_association = FALSE, use_data = TRUE, mcmc_seed = sample_seed(), max_sampling_iterations = 20000, @@ -280,7 +283,8 @@ sccomp_estimate.Seurat <- function(.data, variational_inference = NULL, .sample = NULL, .cell_group = NULL, - .abundance = NULL) { + .abundance = NULL, + exclude_priors = NULL) { .count <- enquo(.count) .sample <- enquo(.sample) @@ -318,7 +322,7 @@ sccomp_estimate.Seurat <- function(.data, verbose = verbose, enable_loo = enable_loo, noise_model = noise_model, - exclude_priors = exclude_priors, + exclude_mean_variability_association = exclude_mean_variability_association, use_data = use_data, mcmc_seed = mcmc_seed, max_sampling_iterations = max_sampling_iterations, @@ -332,7 +336,8 @@ sccomp_estimate.Seurat <- function(.data, variational_inference = variational_inference, .sample = !!.sample, .cell_group = !!.cell_group, - .abundance = !!.abundance + .abundance = !!.abundance, + exclude_priors = exclude_priors ) } @@ -352,16 +357,16 @@ sccomp_estimate.SingleCellExperiment <- function(.data, inference_method = "pathfinder", prior_mean = list(intercept = c(0, 1), coefficients = c(0, 1)), prior_overdispersion_mean_association = list( - intercept = c(5, 2), - slope = c(0, 0.6), - standard_deviation = c(10, 20) + intercept = c(4, 2), + slope = c(0, 2), + standard_deviation = c(1, 0.5) ), .sample_cell_group_pairs_to_exclude = NULL, output_directory = "sccomp_draws_files", verbose = TRUE, enable_loo = FALSE, noise_model = "multi_beta_binomial", - exclude_priors = FALSE, + exclude_mean_variability_association = FALSE, use_data = TRUE, mcmc_seed = sample_seed(), max_sampling_iterations = 20000, @@ -377,7 +382,8 @@ sccomp_estimate.SingleCellExperiment <- function(.data, variational_inference = NULL, .sample = NULL, .cell_group = NULL, - .abundance = NULL) { + .abundance = NULL, + exclude_priors = NULL) { .count <- enquo(.count) @@ -416,7 +422,7 @@ sccomp_estimate.SingleCellExperiment <- function(.data, verbose = verbose, enable_loo = enable_loo, noise_model = noise_model, - exclude_priors = exclude_priors, + exclude_mean_variability_association = exclude_mean_variability_association, use_data = use_data, mcmc_seed = mcmc_seed, max_sampling_iterations = max_sampling_iterations, @@ -430,7 +436,8 @@ sccomp_estimate.SingleCellExperiment <- function(.data, variational_inference = variational_inference, .sample = !!.sample, .cell_group = !!.cell_group, - .abundance = !!.abundance + .abundance = !!.abundance, + exclude_priors = exclude_priors ) } @@ -450,16 +457,16 @@ sccomp_estimate.DFrame <- function(.data, inference_method = "pathfinder", prior_mean = list(intercept = c(0, 1), coefficients = c(0, 1)), prior_overdispersion_mean_association = list( - intercept = c(5, 2), - slope = c(0, 0.6), - standard_deviation = c(10, 20) + intercept = c(4, 2), + slope = c(0, 2), + standard_deviation = c(1, 0.5) ), .sample_cell_group_pairs_to_exclude = NULL, output_directory = "sccomp_draws_files", verbose = TRUE, enable_loo = FALSE, noise_model = "multi_beta_binomial", - exclude_priors = FALSE, + exclude_mean_variability_association = FALSE, use_data = TRUE, mcmc_seed = sample_seed(), max_sampling_iterations = 20000, @@ -475,7 +482,8 @@ sccomp_estimate.DFrame <- function(.data, variational_inference = NULL, .sample = NULL, .cell_group = NULL, - .abundance = NULL) { + .abundance = NULL, + exclude_priors = NULL) { .count <- enquo(.count) .sample <- enquo(.sample) @@ -505,7 +513,7 @@ sccomp_estimate.DFrame <- function(.data, verbose = verbose, enable_loo = enable_loo, noise_model = noise_model, - exclude_priors = exclude_priors, + exclude_mean_variability_association = exclude_mean_variability_association, use_data = use_data, mcmc_seed = mcmc_seed, max_sampling_iterations = max_sampling_iterations, @@ -519,7 +527,8 @@ sccomp_estimate.DFrame <- function(.data, variational_inference = variational_inference, .sample = !!.sample, .cell_group = !!.cell_group, - .abundance = !!.abundance + .abundance = !!.abundance, + exclude_priors = exclude_priors ) } @@ -542,16 +551,16 @@ sccomp_estimate.data.frame <- function(.data, inference_method = "pathfinder", prior_mean = list(intercept = c(0, 1), coefficients = c(0, 1)), prior_overdispersion_mean_association = list( - intercept = c(5, 2), - slope = c(0, 0.6), - standard_deviation = c(10, 20) + intercept = c(4, 2), + slope = c(0, 2), + standard_deviation = c(1, 0.5) ), .sample_cell_group_pairs_to_exclude = NULL, output_directory = "sccomp_draws_files", verbose = TRUE, enable_loo = FALSE, noise_model = "multi_beta_binomial", - exclude_priors = FALSE, + exclude_mean_variability_association = FALSE, use_data = TRUE, mcmc_seed = sample_seed(), max_sampling_iterations = 20000, @@ -567,7 +576,8 @@ sccomp_estimate.data.frame <- function(.data, variational_inference = NULL, .sample = NULL, .cell_group = NULL, - .abundance = NULL) { + .abundance = NULL, + exclude_priors = NULL) { .count <- enquo(.count) @@ -590,6 +600,18 @@ sccomp_estimate.data.frame <- function(.data, details = "The argument variational_inference is now deprecated. Please use inference_method. By default, inference_method value is inferred from variational_inference") inference_method <- ifelse(variational_inference, "variational", "hmc") } + + # `exclude_priors` was renamed because it never actually excluded all priors — + # it disabled the abundance dependence in the variability prior. Forward the + # value to the new argument when supplied so old callers keep working. + if (lifecycle::is_present(exclude_priors) && !is.null(exclude_priors)) { + lifecycle::deprecate_soft( + "2.1.32", + "sccomp::sccomp_estimate(exclude_priors = )", + "sccomp::sccomp_estimate(exclude_mean_variability_association = )" + ) + exclude_mean_variability_association <- exclude_priors + } # Handle deprecated column arguments if (lifecycle::is_present(.sample) && @@ -672,7 +694,7 @@ sccomp_estimate.data.frame <- function(.data, output_directory = output_directory, verbose = verbose, enable_loo = enable_loo, - exclude_priors = exclude_priors, + exclude_mean_variability_association = exclude_mean_variability_association, use_data = use_data, mcmc_seed = mcmc_seed, max_sampling_iterations = max_sampling_iterations, @@ -703,7 +725,7 @@ sccomp_estimate.data.frame <- function(.data, output_directory = output_directory, verbose = verbose, enable_loo = enable_loo, - exclude_priors = exclude_priors, + exclude_mean_variability_association = exclude_mean_variability_association, use_data = use_data, mcmc_seed = mcmc_seed, max_sampling_iterations = max_sampling_iterations, @@ -756,7 +778,7 @@ sccomp_glm_data_frame_raw = function(.data, # Secondary arguments contrasts = NULL, prior_mean = list(intercept = c(0,1), coefficients = c(0,1)), - prior_overdispersion_mean_association = list(intercept = c(5, 2), slope = c(0, 0.6), standard_deviation = c(20, 40)), + prior_overdispersion_mean_association = list(intercept = c(4, 2), slope = c(0, 2), standard_deviation = c(1, 0.5)), percent_false_positive = 5, check_outliers = TRUE, variational_inference = NULL, @@ -764,7 +786,7 @@ sccomp_glm_data_frame_raw = function(.data, test_composition_above_logit_fold_change = 0.1, .sample_cell_group_pairs_to_exclude = NULL, output_directory = "sccomp_draws_files", verbose = FALSE, - exclude_priors = FALSE, + exclude_mean_variability_association = FALSE, bimodal_mean_variability_association = FALSE, enable_loo = FALSE, use_data = TRUE, @@ -841,7 +863,7 @@ sccomp_glm_data_frame_raw = function(.data, percent_false_positive = percent_false_positive, check_outliers = check_outliers, inference_method = inference_method, - exclude_priors = exclude_priors, + exclude_mean_variability_association = exclude_mean_variability_association, bimodal_mean_variability_association = bimodal_mean_variability_association, enable_loo = enable_loo, use_data = use_data, @@ -858,8 +880,6 @@ sccomp_glm_data_frame_raw = function(.data, ) } - - sccomp_glm_data_frame_counts = function(.data, formula_composition = ~ 1 , formula_variability = ~ 1, @@ -872,7 +892,7 @@ sccomp_glm_data_frame_counts = function(.data, contrasts = NULL, #.grouping_for_random_effect = NULL, prior_mean = list(intercept = c(0,1), coefficients = c(0,1)), - prior_overdispersion_mean_association = list(intercept = c(5, 2), slope = c(0, 0.6), standard_deviation = c(20, 40)), + prior_overdispersion_mean_association = list(intercept = c(4, 2), slope = c(0, 2), standard_deviation = c(1, 0.5)), percent_false_positive = 5, check_outliers = TRUE, variational_inference = NULL, @@ -880,7 +900,7 @@ sccomp_glm_data_frame_counts = function(.data, test_composition_above_logit_fold_change = 0.1, .sample_cell_group_pairs_to_exclude = NULL, output_directory = "sccomp_draws_files", verbose = FALSE, - exclude_priors = FALSE, + exclude_mean_variability_association = FALSE, bimodal_mean_variability_association = FALSE, enable_loo = FALSE, use_data = TRUE, @@ -1041,7 +1061,7 @@ sccomp_glm_data_frame_counts = function(.data, data_for_model$prior_prec_sd = prior_overdispersion_mean_association$standard_deviation data_for_model$prior_mean_intercept = prior_mean$intercept data_for_model$prior_mean_coefficients = prior_mean$coefficients - data_for_model$exclude_priors = exclude_priors + data_for_model$exclude_mean_variability_association = exclude_mean_variability_association data_for_model$enable_loo = enable_loo # # Check that design matrix is not too big @@ -1062,9 +1082,11 @@ sccomp_glm_data_frame_counts = function(.data, seed = mcmc_seed, max_sampling_iterations = max_sampling_iterations, pars = c( - "beta", "alpha", "prec_coeff","prec_sd", "alpha_normalised", - "random_effect", "random_effect_2", - "random_effect_sigma", "random_effect_sigma_2", + "beta", "alpha", + "prec_intercept_1", "prec_slope_1", "prec_intercept_2", "prec_slope_2", + "prec_sd", + # Random effect outputs - one per slot (1..4) + "random_effect_1", "random_effect_2", "random_effect_3", "random_effect_4", "log_lik" ), sig_figs = sig_figs, diff --git a/R/sccomp_plot_intervals.R b/R/sccomp_plot_intervals.R new file mode 100644 index 00000000..3ddfd948 --- /dev/null +++ b/R/sccomp_plot_intervals.R @@ -0,0 +1,828 @@ + +#' Plot 1D Intervals for Cell-group Effects +#' +#' This function creates a series of 1D interval plots for cell-group effects, highlighting significant differences based on a given significance threshold. +#' +#' @param .data Data frame containing the main data. +#' @param significance_threshold Numeric value specifying the significance threshold for highlighting differences. +#' @param test_composition_above_logit_fold_change A positive integer. It is the effect threshold used for the hypothesis test. A value of 0.2 correspond to a change in cell proportion of 10% for a cell type with baseline proportion of 50%. That is, a cell type goes from 45% to 50%. When the baseline proportion is closer to 0 or 1 this effect thrshold has consistent value in the logit uncontrained scale. +#' @param show_fdr_message Logical. Whether to show the Bayesian FDR interpretation message on the plot. Default is TRUE. +#' @param significance_statistic Character vector indicating which statistic to highlight. Default is "pH0". +#' @param factor Optional character string selecting one model factor to plot. If provided, plots are restricted to that factor plus `(Intercept)`. +#' @param sort_by Character vector indicating how to sort taxa. Options are "none" (default), "effect" (by effect size), "significance" (by FDR/pH0), or "alphabetical". +#' @importFrom patchwork wrap_plots +#' @importFrom forcats fct_reorder fct_inorder +#' @importFrom tidyr drop_na +#' +#' @export +#' +#' @return A combined plot of 1D interval plots. +#' @examples +#' +#' print("cmdstanr is needed to run this example.") +#' +#' \donttest{ +#' if (instantiate::stan_cmdstan_exists()) { +#' data("counts_obj") +#' +#' estimate <- sccomp_estimate( +#' counts_obj, +#' ~ type, +#' ~1, +#' "sample", +#' "cell_group", +#' "count", +#' cores = 1 +#' ) |> +#' sccomp_test() +#' +#' # Example usage: +#' my_plot = sccomp_plot_intervals_1D(estimate, sort_by = "effect") +#' +#' } +#' } +#' +#' +sccomp_plot_intervals_1D = function( + .data, + factor = NULL, + significance_threshold = 0.05, + test_composition_above_logit_fold_change = .data |> attr("test_composition_above_logit_fold_change"), + show_fdr_message = TRUE, + significance_statistic = c("pH0", "FDR"), + sort_by = c("none", "effect", "significance", "alphabetical") +) { + significance_statistic <- match.arg(significance_statistic) + sort_by <- match.arg(sort_by) + + # Define the variables as NULL to avoid CRAN NOTES + parameter <- NULL + estimate <- NULL + value <- NULL + pH0 <- NULL + FDR <- NULL + effect <- NULL + + .cell_group = attr(.data, ".cell_group") + + # Check if test have been done + if(.data |> select(ends_with("FDR")) |> ncol() |> equals(0)) + stop("sccomp says: to produce plots, you need to run the function sccomp_test() on your estimates.") + + .data <- subset_results_by_factor(.data, factor, keep_intercept = TRUE) + + plot_list = + .data |> + + # Reshape data + select(-contains("n_eff"), -contains("R_k_hat"), -contains("_ess"), -contains("_rhat")) |> + pivot_longer(c(contains("c_"), contains("v_")), names_sep = "_", names_to = c("which", "estimate")) |> + pivot_wider(names_from = estimate, values_from = value) |> + + # Nest data by parameter and which + nest(data = -c(parameter, which)) |> + mutate(plot = pmap( + list(data, which, parameter), + function(plot_data, plot_which, plot_param) { + # Check if there are any statistics to plot + if(plot_data |> filter(!is.na(effect)) |> nrow() |> equals(0)) + return(NA) + + # Choose color variable and legend + if (significance_statistic == "FDR") { + color_aes <- aes(xmin = lower, xmax = upper, color = FDR < significance_threshold) + color_scale <- scale_color_manual(values = c("grey40", "red")) + legend_title <- "FDR < significance_threshold" + } else { + color_aes <- aes(xmin = lower, xmax = upper, color = pH0 < significance_threshold) + color_scale <- scale_color_manual(values = c("grey40", "red")) + legend_title <- "pH0 < significance_threshold" + } + + # Determine y-axis variable based on sort_by + # Use string-based column selection instead of NSE + if (sort_by == "none") { + # No sorting + plot_data$y_var <- plot_data[[.cell_group]] + } else if (sort_by == "effect") { + # Sort by absolute effect size + plot_data$y_var <- fct_reorder(plot_data[[.cell_group]], abs(plot_data$effect)) + } else if (sort_by == "significance") { + # Sort by significance + if (significance_statistic == "FDR") { + plot_data$y_var <- fct_reorder(plot_data[[.cell_group]], plot_data$FDR, .desc = TRUE) + } else { + plot_data$y_var <- fct_reorder(plot_data[[.cell_group]], plot_data$pH0, .desc = TRUE) + } + } else if (sort_by == "alphabetical") { + # Alphabetical sorting + plot_data <- plot_data %>% arrange(.data[[.cell_group]]) + plot_data$y_var <- fct_inorder(plot_data[[.cell_group]]) + } + + ggplot(plot_data, aes(x = effect, y = y_var)) + + geom_vline(xintercept = test_composition_above_logit_fold_change, colour = "grey") + + geom_vline(xintercept = -test_composition_above_logit_fold_change, colour = "grey") + + geom_errorbar(color_aes) + + geom_point() + + color_scale + + xlab("Credible interval of the effect") + + ylab("Cell group") + + ggtitle(sprintf("%s %s", plot_which, plot_param)) + + sccomp_theme() + + theme(legend.position = "bottom") + + guides(color = guide_legend(title = legend_title)) + } + )) %>% + + # Filter out NA plots + filter(!is.na(plot)) |> + pull(plot) + + # Combine all individual plots into one plot + combined_plot <- plot_list |> + wrap_plots(ncol = plot_list |> length() |> sqrt() |> ceiling()) + + # Only show the FDR message if significance_statistic == "FDR" and show_fdr_message is TRUE + if (significance_statistic == "FDR" && show_fdr_message) { + combined_plot <- combined_plot + theme(plot.caption = ggplot2::element_text(hjust = 0)) + combined_plot <- combined_plot + patchwork::plot_annotation( + caption = paste( + "Bayesian FDR: Stephens' method (doi: 10.1093/biostatistics/kxw041)", + "\nFDR-significant populations may cross fold change thresholds because Bayesian FDR considers posterior probabilities rather than p-values.", + "\nThe method sorts null hypothesis probabilities in ascending order and calculates cumulative averages for robust false discovery control.", + sep = "" + ) + ) + } + combined_plot +} + +#' Plot 2D Intervals for Mean-Variance Association +#' +#' This function creates a 2D interval plot for mean-variance association, handling both single and bimodal models. +#' It highlights significant differences based on a given significance threshold. +#' +#' @param .data Data frame containing the main data. +#' @param significance_threshold Numeric value specifying the significance threshold for highlighting differences. Default is 0.05. +#' @param test_composition_above_logit_fold_change A positive integer. It is the effect threshold used for the hypothesis test. +#' @param show_fdr_message Logical. Whether to show the Bayesian FDR interpretation message on the plot. Default is TRUE. +#' @param significance_statistic Character vector indicating which statistic to highlight. Default is "pH0". +#' @param factor Optional character string selecting one model factor to plot. If provided, plots are restricted to that factor plus `(Intercept)`. +#' @param add_marginal_density Logical. Whether to add marginal density plots on adjusted panels. Default is TRUE. +#' @param omit_ci Logical. Whether to omit credible interval error bars. Default is FALSE. +#' +#' @importFrom dplyr filter arrange mutate if_else row_number bind_rows distinct slice pull with_groups +#' @importFrom ggplot2 ggplot geom_vline geom_hline geom_errorbar geom_point geom_line geom_blank aes xlab ylab facet_wrap theme_bw theme labs guides guide_legend scale_color_manual scale_alpha_manual scale_fill_manual element_rect element_blank element_text +#' @importFrom ggrepel geom_text_repel +#' @importFrom stringr str_detect +#' @importFrom ggside geom_ysidedensity theme_ggside_void scale_ysidex_continuous +#' +#' @export +#' +#' @return A ggplot object representing the 2D interval plot. +#' +#' @examples +#' +#' print("cmdstanr is needed to run this example.") +#' +#' \donttest{ +#' if (instantiate::stan_cmdstan_exists()) { +#' data("counts_obj") +#' +#' estimate <- sccomp_estimate( +#' counts_obj, +#' ~ type, +#' ~type, +#' "sample", +#' "cell_group", +#' "count", +#' cores = 1, +#' bimodal_mean_variability_association = TRUE +#' ) |> +#' sccomp_test() +#' +#' # Example usage: +#' my_plot = sccomp_plot_intervals_2D(estimate) +#' +#' } +#' } +#' +sccomp_plot_intervals_2D <- function( + .data, + factor = NULL, + significance_threshold = 0.05, + test_composition_above_logit_fold_change = + .data |> attr("test_composition_above_logit_fold_change"), + show_fdr_message = TRUE, + significance_statistic = c("pH0", "FDR"), + add_marginal_density = TRUE, + omit_ci = FALSE +) { + + significance_statistic <- match.arg(significance_statistic) + + # Locals declared as NULL to silence R CMD check NOTEs about "no visible binding" + # for tidyverse NSE references (dplyr/ggplot evaluate these as column names). + v_effect <- NULL + parameter <- NULL + c_effect <- NULL + c_lower <- NULL + c_upper <- NULL + c_FDR <- NULL + v_lower <- NULL + v_upper <- NULL + v_FDR <- NULL + cell_type_label <- NULL + pH0 <- NULL + FDR <- NULL + c_pH0 <- NULL + v_pH0 <- NULL + component <- NULL + assigned_component <- NULL + v_value <- NULL + + .cell_group <- attr(.data, ".cell_group") + + # The plot relies on FDR / pH0 columns that only exist after sccomp_test(); + # fail fast with a user-meaningful message rather than later with a cryptic NSE error. + if(.data |> select(ends_with("FDR")) |> ncol() == 0) + stop("sccomp says: you need to run sccomp_test() first.") + + # Intercept is kept even when subsetting to a single factor because the Intercept + # panel acts as the visual baseline against which factor-specific effects are read. + .data <- subset_results_by_factor(.data, factor, keep_intercept = TRUE) + + fit <- attr(.data, "fit") + + # Determine model topology *before* pulling Stan summaries: `prec_*_2` variables + # are only declared in the bimodal Stan program, so blindly requesting them + # would error on single-component fits. + bimodal_flag <- attr(.data, "model_input")$bimodal_mean_variability_association + # no check needed; bimodal_flag should always be present + bimodal_flag <- isTRUE(as.logical(bimodal_flag)) + + # Posterior summaries of the mean-variability regression coefficients + # v_pred = -(prec_intercept + prec_slope * c_effect) + # These are pulled once and reused per-parameter via `param_idx`. + prec_intercept_1_summary <- fit$summary("prec_intercept_1") + prec_slope_1_summary <- fit$summary("prec_slope_1") + + # Only parameters that actually have a v_effect (i.e. are estimated under the + # variability design Xa) can be plotted on the 2D space. + param_names <- .data |> + filter(!is.na(v_effect)) |> + distinct(parameter) |> + pull("parameter") + + # Map each parameter name to its column index in the variability design matrix + # `Xa` because Stan stores `prec_*` indexed by Xa column position, not by name. + param_idx <- match(param_names, colnames(attr(.data, "model_input")$Xa)) + if (any(is.na(param_idx))) { + stop("sccomp says: could not map selected parameters to model coefficients.") + } + + # `params_list` flattens the Stan posterior summaries into one entry per parameter, + # carrying only the scalars (`intercept`, `slope`, or per-component variants) + # required by downstream geometry, so the heavy `fit$summary` tibbles are not + # re-scanned inside per-row mutate() / lapply() calls. + if (!bimodal_flag) { + params_list <- lapply(seq_along(param_names), function(a) { + param_name <- param_names[a] + idx <- param_idx[a] + list( + parameter = param_name, + intercept = prec_intercept_1_summary$mean[idx], + slope = prec_slope_1_summary$mean[idx] + ) + }) + + # Print the fitted lines so users can sanity-check the regression underlying + # the visual adjustment (especially useful when comparing across fits). + message("=== Single Model Parameters ===") + for(i in 1:length(params_list)) { + p <- params_list[[i]] + message(sprintf("\n%s:", p$parameter)) + message(sprintf(" v = -(%.3f + %.3f \u00d7 c)", p$intercept, p$slope)) + } + message("") + + } else { + prec_intercept_2_summary <- fit$summary("prec_intercept_2") + prec_slope_2_summary <- fit$summary("prec_slope_2") + # `mix_p` is the posterior mixing weight of component 1; reported in the + # caption so users can judge which component dominates the population. + mix_p <- fit$summary("mix_p") |> pull(mean) + + params_list <- lapply(seq_along(param_names), function(a) { + param_name <- param_names[a] + idx <- param_idx[a] + + list( + parameter = param_name, + intercept_1 = prec_intercept_1_summary$mean[idx], + slope_1 = prec_slope_1_summary$mean[idx], + slope_2 = prec_slope_2_summary$mean[idx], + intercept_2 = prec_intercept_2_summary$mean[idx] + ) + }) + + message("=== Bimodal Model Parameters ===") + for(i in 1:length(params_list)) { + p <- params_list[[i]] + message(sprintf("\n%s:", p$parameter)) + message(sprintf(" Component 1: v = -(%.3f + %.3f \u00d7 c)", p$intercept_1, p$slope_1)) + message(sprintf(" Component 2: v = -(%.3f + %.3f \u00d7 c)", p$intercept_2, p$slope_2)) + } + message("") + } + + # The 2D plot juxtaposes two views of v_effect for each parameter: + # - "raw": v_effect with the mean-variability association reintroduced, + # i.e. the data as it would look *without* the model's correction; + # the fitted regression line should pass through this cloud. + # - "adjusted": v_effect with the association removed (what sccomp uses for + # inference). Cells off-axis here are the genuine outliers. + # Since `.data$v_effect` is already the adjusted form (it comes from the + # `alpha_normalised` draws), we *reverse the adjustment* to construct "raw" + # and we use v_effect verbatim for "adjusted". + if (!bimodal_flag) { + .data_raw_list <- lapply(params_list, function(params) { + .data %>% + filter(parameter == params$parameter) %>% + mutate( + v_effect = v_effect - params$slope * c_effect, + v_lower = v_lower - params$slope * c_effect, + v_upper = v_upper - params$slope * c_effect, + parameter = paste0(params$parameter, ", raw") + ) + }) + } else { + # Bimodal case: each cell type's "raw" position is generated by *one* of the + # two mixture components, but the latent assignment is not directly observed. + # We hard-assign by nearest predicted v_effect (smallest residual against + # each component's regression line), then reverse the adjustment using only + # that component's slope. This produces a visually coherent "raw" scatter + # in which each point lies near its parent regression line rather than at + # an arbitrary average of the two. + .data_raw_list <- lapply(params_list, function(params) { + .data %>% + filter(parameter == params$parameter) %>% + rowwise() %>% + mutate( + raw_v_comp1 = v_effect - params$slope_1 * c_effect, + raw_v_comp2 = v_effect - params$slope_2 * c_effect, + pred_comp1 = -(params$intercept_1 + params$slope_1 * c_effect), + pred_comp2 = -(params$intercept_2 + params$slope_2 * c_effect), + assigned_component = if_else( + abs(raw_v_comp1 - pred_comp1) < abs(raw_v_comp2 - pred_comp2), 1, 2 + ), + slope_to_use = if_else(assigned_component == 1, params$slope_1, params$slope_2), + v_effect = v_effect - slope_to_use * c_effect, + v_lower = v_lower - slope_to_use * c_effect, + v_upper = v_upper - slope_to_use * c_effect, + parameter = paste0(params$parameter, ", raw") + ) %>% + ungroup() %>% + select(-raw_v_comp1, -raw_v_comp2, -pred_comp1, -pred_comp2, -slope_to_use) + }) + } + + .data_raw <- bind_rows(.data_raw_list) + + # "adjusted" rows are the original v_effect, only the parameter label changes + # so facet_wrap can place them in a separate panel. + .data_adjusted_list <- lapply(params_list, function(params) { + .data %>% + filter(parameter == params$parameter) %>% + mutate(parameter = paste0(params$parameter, ", adjusted")) + }) + .data_adjusted <- bind_rows(.data_adjusted_list) + + .data_plot <- bind_rows(.data_raw, .data_adjusted) + + # Interleave "raw" before "adjusted" for each parameter so the facets read + # left-to-right as "before -> after the model's adjustment". param_order is + # also reused by the helper to ensure stable panel order across patchwork and + # facet_wrap callers. + param_order <- c() + for(p in params_list) { + param_order <- c(param_order, paste0(p$parameter, ", raw"), paste0(p$parameter, ", adjusted")) + } + + .data_plot$parameter <- factor(.data_plot$parameter, levels = param_order) + + # Label only the top-3 cell groups by significance per panel, and only in + # ", adjusted" panels because the "raw" panel is meant to show the underlying + # association (labels there would clutter without adding analytical value). + # Two sequential passes — first by c_FDR (abundance), then by v_FDR + # (variability) — guarantee a cell is labeled if it is significant on + # *either* axis, while preserving the abundance label when both apply. + .data_plot <- .data_plot %>% + filter(!is.na(v_effect)) %>% + with_groups( + parameter, + ~ .x %>% + arrange(c_FDR) %>% + mutate( + cell_type_label = if_else( + row_number() <= 3 & + c_FDR < significance_threshold & + str_detect(parameter, ", adjusted$"), + !!.cell_group, + "" + ) + ) + ) %>% + with_groups( + parameter, + ~ .x %>% + arrange(v_FDR) %>% + mutate( + cell_type_label = if_else( + row_number() <= 3 & + v_FDR < significance_threshold & + str_detect(parameter, ", adjusted$") & + cell_type_label == "", + !!.cell_group, + cell_type_label + ) + ) + ) + + # Encode significance in the errorbar aesthetics so the plot is readable + # without consulting the underlying table. Significance is only meaningful on + # ", adjusted" panels — raw panels are unconditionally desaturated to convey + # "descriptive, not inferential". FDR (Stephens) uses red, pH0 (posterior + # tail probability) uses blue, mirroring their distinct meanings. + if (significance_statistic == "FDR") { + color_c_aes <- aes( + xmin = c_lower, xmax = c_upper, + color = c_FDR < significance_threshold & str_detect(parameter, ", adjusted$"), + alpha = c_FDR < significance_threshold & str_detect(parameter, ", adjusted$") + ) + color_v_aes <- aes( + ymin = v_lower, ymax = v_upper, + color = v_FDR < significance_threshold & str_detect(parameter, ", adjusted$"), + alpha = v_FDR < significance_threshold & str_detect(parameter, ", adjusted$") + ) + color_scale <- scale_color_manual(values = c("#D3D3D3", "#E41A1C")) + alpha_scale <- scale_alpha_manual(values = c(0.4, 1)) + legend_title <- "FDR < significance_threshold" + } else { + color_c_aes <- aes( + xmin = c_lower, xmax = c_upper, + color = c_pH0 < significance_threshold & str_detect(parameter, ", adjusted$"), + alpha = c_pH0 < significance_threshold & str_detect(parameter, ", adjusted$") + ) + color_v_aes <- aes( + ymin = v_lower, ymax = v_upper, + color = v_pH0 < significance_threshold & str_detect(parameter, ", adjusted$"), + alpha = v_pH0 < significance_threshold & str_detect(parameter, ", adjusted$") + ) + color_scale <- scale_color_manual(values = c("#D3D3D3", "#377EB8")) + alpha_scale <- scale_alpha_manual(values = c(0.4, 1)) + legend_title <- "pH0 < significance_threshold" + } + + # Prepare regression line data based on model type + if (!bimodal_flag) { + regression_data_all <- lapply(params_list, function(params) { + raw_param <- paste0(params$parameter, ", raw") + param_data <- .data_plot %>% filter(parameter == raw_param) + if(nrow(param_data) == 0) return(NULL) + + c_range <- range(param_data$c_effect, na.rm = TRUE) + c_seq <- seq(c_range[1], c_range[2], length.out = 100) + v_pred <- -(params$intercept + params$slope * c_seq) + + data.frame( + c_effect = c_seq, + v_effect = v_pred, + parameter = raw_param, + stringsAsFactors = FALSE + ) + }) %>% bind_rows() + + adjusted_lines_all <- lapply(params_list, function(params) { + adj_param <- paste0(params$parameter, ", adjusted") + adj_data <- .data_plot %>% filter(parameter == adj_param) + if(nrow(adj_data) == 0) return(NULL) + + c_range_adj <- range(adj_data$c_effect, na.rm = TRUE) + mean_adj <- mean(adj_data$v_effect, na.rm = TRUE) + + data.frame( + c_effect = c_range_adj, + v_effect = rep(mean_adj, 2), + parameter = adj_param, + stringsAsFactors = FALSE + ) + }) %>% bind_rows() + + } else { + # Bimodal: two regression lines and two horizontal references per panel. + # The `component` column lets the plot helper colour them distinctly. + regression_data_all <- lapply(params_list, function(params) { + raw_param <- paste0(params$parameter, ", raw") + param_data <- .data_plot %>% filter(parameter == raw_param) + if(nrow(param_data) == 0) return(NULL) + + c_range <- range(param_data$c_effect, na.rm = TRUE) + c_seq <- seq(c_range[1], c_range[2], length.out = 100) + + v_pred_1 <- -(params$intercept_1 + params$slope_1 * c_seq) + v_pred_2 <- -(params$intercept_2 + params$slope_2 * c_seq) + + bind_rows( + data.frame( + c_effect = c_seq, v_effect = v_pred_1, + parameter = raw_param, component = "Component 1" + ), + data.frame( + c_effect = c_seq, v_effect = v_pred_2, + parameter = raw_param, component = "Component 2" + ) + ) + }) %>% bind_rows() + + adjusted_lines_all <- lapply(params_list, function(params) { + adj_param <- paste0(params$parameter, ", adjusted") + c_range_adj <- range(.data_plot$c_effect[.data_plot$parameter == adj_param], na.rm = TRUE) + + bind_rows( + data.frame( + c_effect = c_range_adj, v_effect = rep(-params$intercept_1, 2), + parameter = adj_param, component = "Component 1" + ), + data.frame( + c_effect = c_range_adj, v_effect = rep(-params$intercept_2, 2), + parameter = adj_param, component = "Component 2" + ) + ) + }) %>% bind_rows() + } + + # Caption is reserved for FDR mode: users that opt into FDR-based inference + # benefit from the explicit citation/interpretation reminder; pH0 users are + # assumed to already understand the posterior-tail interpretation. + if (significance_statistic == "FDR" && show_fdr_message) { + if (!bimodal_flag) { + caption_text <- paste( + "Single model with parameter-specific slopes", + "\nBlue line shows mean-variability relationship for each parameter", + "\n'Adjusted' panels show variability after removing parameter-specific mean-variability association", + "\nMarginal density shows posterior distribution of variability intercept parameters", + "\nBayesian FDR: Stephens' method (doi: 10.1093/biostatistics/kxw041)", + sep = "" + ) + } else { + caption_text <- paste( + sprintf("Bimodal model: mix_p = %.3f (Component 1 weight)", mix_p), + "\nSolid blue line = Component 1 | Dashed orange line = Component 2", + "\n'Adjusted' panels show variability after removing mean-variability association", + "\nMarginal density shows posterior distribution of variability intercept parameters", + "\nBayesian FDR: Stephens' method (doi: 10.1093/biostatistics/kxw041)", + sep = "" + ) + } + } else { + caption_text <- NULL + } + + # Both modes (with and without marginal density) produce a single faceted plot; the only difference is + # whether ggside's y-side area layer is added on top. We rely on ggside + # rather than patchwork-of-per-panel-plots so axis titles are not duplicated + # and panel widths stay grid-aligned. + build_2d_interval_plot <- function(plot_data, regression_data, adjusted_lines, density_data = NULL) { + + # Force the canonical factor order on every input: the data preparation + # above does it on `.data_plot`, but `regression_data` / `adjusted_lines` / + # `density_data` come from separate pipelines and must agree on level order + # for facet_wrap to keep panels in sync. + plot_data <- plot_data %>% + mutate(parameter = factor(as.character(parameter), levels = param_order)) + + if (!is.null(regression_data) && nrow(regression_data) > 0) { + regression_data <- regression_data %>% + mutate(parameter = factor(as.character(parameter), levels = param_order)) + } + + if (!is.null(adjusted_lines) && nrow(adjusted_lines) > 0) { + adjusted_lines <- adjusted_lines %>% + mutate(parameter = factor(as.character(parameter), levels = param_order)) + } + + if (!is.null(density_data) && nrow(density_data) > 0) { + density_data <- density_data %>% + mutate(parameter = factor(as.character(parameter), levels = param_order)) + } + + # Decision boundary at ± test_composition_above_logit_fold_change: cells + # outside this box are the only ones a hypothesis test could call significant. + p <- ggplot(plot_data, aes(c_effect, v_effect)) + + geom_vline( + xintercept = c(-test_composition_above_logit_fold_change, test_composition_above_logit_fold_change), + colour = "grey", linetype = "dashed", linewidth = 0.3 + ) + + geom_hline( + yintercept = c(-test_composition_above_logit_fold_change, test_composition_above_logit_fold_change), + colour = "grey", linetype = "dashed", linewidth = 0.3 + ) + + # Regression / reference lines. + # `inherit.aes = FALSE` shields these geoms from the top-level + # aes(c_effect, v_effect) mapping so the lines use their own data verbatim + # — important because their tibbles have a different row schema (no CI + # columns, no cell_group, etc.) than `plot_data`. + # In bimodal mode the two components get distinct visual treatment so the + # eye can separate them without consulting the caption. + if (!bimodal_flag) { + if(!is.null(regression_data) && nrow(regression_data) > 0) { + p <- p + geom_line(data = regression_data, mapping = aes(c_effect, v_effect), + color = "#0072B2", linewidth = 0.5, alpha = 0.8, inherit.aes = FALSE) + } + if(!is.null(adjusted_lines) && nrow(adjusted_lines) > 0) { + p <- p + geom_line(data = adjusted_lines, mapping = aes(c_effect, v_effect), + color = "#0072B2", linewidth = 0.5, alpha = 0.8, inherit.aes = FALSE) + } + } else { + if(!is.null(regression_data) && nrow(regression_data) > 0) { + p <- p + + geom_line(data = regression_data %>% filter(component == "Component 1"), + mapping = aes(c_effect, v_effect), color = "#0072B2", + linewidth = 0.5, alpha = 0.8, inherit.aes = FALSE) + + geom_line(data = regression_data %>% filter(component == "Component 2"), + mapping = aes(c_effect, v_effect), color = "#D55E00", + linewidth = 0.5, alpha = 0.8, linetype = "dashed", inherit.aes = FALSE) + } + if(!is.null(adjusted_lines) && nrow(adjusted_lines) > 0) { + p <- p + + geom_line(data = adjusted_lines %>% filter(component == "Component 1"), + mapping = aes(c_effect, v_effect), color = "#0072B2", + linewidth = 0.5, alpha = 0.8, inherit.aes = FALSE) + + geom_line(data = adjusted_lines %>% filter(component == "Component 2"), + mapping = aes(c_effect, v_effect), color = "#D55E00", + linewidth = 0.5, alpha = 0.8, linetype = "dashed", inherit.aes = FALSE) + } + } + + # Credible intervals. When omitted, we substitute invisible `geom_blank` + # layers at the CI bounds so the plot's coordinate system is still trained + # by the same data range — keeping axes identical whether errorbars are + # drawn or not (otherwise the no-CI plot would zoom in onto just the points). + if (!omit_ci) { + p <- p + + geom_errorbar(color_c_aes, linewidth = 0.2) + + geom_errorbar(color_v_aes, linewidth = 0.2) + + color_scale + + alpha_scale + + guides(color = guide_legend(title = legend_title), alpha = "none") + } else { + p <- p + + geom_blank(aes(x = c_lower, y = v_lower)) + + geom_blank(aes(x = c_upper, y = v_upper)) + } + + # Points are drawn *after* the errorbars/lines so they sit on top and + # remain readable when CIs overlap; labels go last for highest z-order. + # Note `geom_text_repel` uses `-v_effect` (y-axis flip) because labels are + # placed relative to the inverted visual reading where higher v means + # higher dispersion — this matches users' mental model of "outlier upward". + p <- p + + geom_point(size = 0.2) + + geom_text_repel( + aes(c_effect, -v_effect, label = cell_type_label), + size = 2.5, + data = plot_data %>% filter(cell_type_label != ""), + max.overlaps = 20 + ) + + xlab("c_effect (Abundance effect)") + + ylab("v_effect (Variability effect)") + + theme_bw() + + theme( + legend.position = "bottom", + strip.background = element_rect(fill = "white"), + panel.grid.minor = element_blank() + ) + + # Marginal posterior densities, drawn via ggside so they participate in the + # same facet system as the main scatter (no patchwork → no duplicated axis + # titles, no panel-width misalignment). We pass raw posterior draws (one + # row per draw, with `parameter` set to the panel they belong to) and let + # `geom_ysidedensity` compute the kernel density per facet. ggside lays + # the result on the y-side, mapping the data's `y` aesthetic to the main + # panel's y-axis. `scales = "free"` propagates so each panel's density + # rescales to fit its own y-range. + if (!is.null(density_data) && nrow(density_data) > 0) { + density_aes <- if (bimodal_flag) { + aes(y = v_value, fill = component) + } else { + aes(y = v_value) + } + p <- p + + geom_ysidedensity( + data = density_data, + mapping = density_aes, + alpha = 0.5, + position = "identity", + inherit.aes = FALSE + ) + + # Drop ggside's own axis decorations: the density value axis is + # uninformative on its own (relative scale per panel) and would + # otherwise clutter the strip. + theme_ggside_void() + + scale_ysidex_continuous(expand = c(0, 0)) + + if (bimodal_flag) { + p <- p + scale_fill_manual(values = c("Component 1" = "#0072B2", "Component 2" = "#D55E00")) + } + } + + # `scales = "free"` is intentional: raw and adjusted panels live on + # different natural scales (raw v_effect can span much wider than adjusted + # residuals), and the same applies across parameters. + p + facet_wrap(~ parameter, scales = "free", ncol = 2) + } + + # Posterior draws of `-prec_intercept_*[param_idx]` for every parameter, + # replicated across both that parameter's facets (", raw" and ", adjusted"). + # The intercept density is a property of the parameter, not of the visual + # mode, so showing it on every facet is both honest and what ggside requires + # to attach a side layer per panel (omitting a facet's density would leave + # that side panel empty, breaking visual rhythm). + # The sign flip mirrors the parameterisation v = -(prec_intercept + slope·c) + # so the density lies on the same axis orientation as v_effect. We pass raw + # draws (not pre-computed densities) because ggside's density geom computes + # the kernel per facet itself. + density_data_all <- if (add_marginal_density) { + bind_rows(lapply(seq_along(param_names), function(i) { + param_name <- param_names[i] + idx <- param_idx[i] + + component_vars <- if (bimodal_flag) { + c("Component 1" = paste0("prec_intercept_1[", idx, "]"), + "Component 2" = paste0("prec_intercept_2[", idx, "]")) + } else { + c("Component 1" = paste0("prec_intercept_1[", idx, "]")) + } + + bind_rows(lapply(seq_along(component_vars), function(j) { + var_name <- component_vars[[j]] + values <- as.vector(fit$draws(variables = var_name, format = "draws_df")[[var_name]]) + draws_per_param <- tibble( + component = names(component_vars)[j], + v_value = -values + ) + bind_rows( + draws_per_param %>% mutate(parameter = paste0(param_name, ", raw")), + draws_per_param %>% mutate(parameter = paste0(param_name, ", adjusted")) + ) + })) + })) + } else { + NULL + } + + p <- build_2d_interval_plot( + .data_plot, + regression_data_all, + adjusted_lines_all, + density_data = density_data_all + ) + + if (!is.null(caption_text)) { + p <- p + + theme(plot.caption = element_text(hjust = 0, size = 9)) + + labs(caption = caption_text) + } + + p +} + +#' Deprecated interval-plot aliases +#' +#' These soft-deprecated aliases forward to [sccomp_plot_intervals_1D()] and +#' [sccomp_plot_intervals_2D()]. Use those functions directly instead. +#' +#' @param ... Arguments passed on to [sccomp_plot_intervals_1D()] or +#' [sccomp_plot_intervals_2D()]. +#' @return A `ggplot` object, as returned by the target function. +#' @keywords internal +#' @name deprecated-interval-plots +NULL + +#' @rdname deprecated-interval-plots +#' @importFrom lifecycle deprecate_soft +#' @export +plot_1D_intervals <- function(...) { + deprecate_soft("2.1.29", "plot_1D_intervals()", "sccomp_plot_intervals_1D()") + sccomp_plot_intervals_1D(...) +} + +#' @rdname deprecated-interval-plots +#' @export +plot_2D_intervals <- function(...) { + deprecate_soft("2.1.29", "plot_2D_intervals()", "sccomp_plot_intervals_2D()") + sccomp_plot_intervals_2D(...) +} \ No newline at end of file diff --git a/R/sccomp_remove_outliers.R b/R/sccomp_remove_outliers.R index 1ba3ba9f..9b101747 100644 --- a/R/sccomp_remove_outliers.R +++ b/R/sccomp_remove_outliers.R @@ -196,35 +196,39 @@ sccomp_remove_outliers.sccomp_tbl = function(.estimate, # This is for the new data generation with selected factors to do adjustment data = - .estimate |> - attr("model_input") |> - c(list( - - # Add subset of coefficients - X_original = data_for_model$X, - N_original = data_for_model$N, - length_X_which = ncol(data_for_model$X), - length_XA_which = ncol(data_for_model$XA), - X_which = seq_len(ncol(data_for_model$X)) |> as.array(), - XA_which = seq_len(ncol(data_for_model$Xa)) |> as.array(), - - # Random intercept common variable between grouping 1 and 2 - ncol_X_random_eff_new = ncol(data_for_model$X_random_effect) |> c(ncol(data_for_model$X_random_effect_2) ), # I could put this in the intial data - length_X_random_effect_which = ncol(data_for_model$X_random_effect) |> c(ncol(data_for_model$X_random_effect_2)), - - # Grouping 1 - X_random_effect_which = seq_len(ncol(data_for_model$X_random_effect)) |> as.array(), - - # Grouping 2 - Random intercept DUPLICATED - X_random_effect_which_2 = seq_len(ncol(data_for_model$X_random_effect_2)) |> as.array(), - - # Initialize unseen random effect variables - ncol_X_random_eff_unseen = c(0, 0), - X_random_effect_unseen = matrix(0, nrow = nrow(data_for_model$X), ncol = 0), - X_random_effect_2_unseen = matrix(0, nrow = nrow(data_for_model$X), ncol = 0), - - create_intercept = FALSE - )), + data_for_model |> + c( + list( + # Add subset of coefficients + X_original = data_for_model$X, + N_original = data_for_model$N, + length_X_which = ncol(data_for_model$X), + length_XA_which = ncol(data_for_model$XA), + X_which = seq_len(ncol(data_for_model$X)) |> as.array(), + XA_which = seq_len(ncol(data_for_model$Xa)) |> as.array(), + + # Per-slot random-effect pass-throughs (the outlier model uses the + # same design, so X_random_effect_which_k is just an identity over + # the slot's columns; unseen matrices are empty). + ncol_X_random_eff_new = data_for_model$ncol_X_random_eff, + length_X_random_effect_which = data_for_model$ncol_X_random_eff, + ncol_X_random_eff_unseen = rep(0L, 4L), + + create_intercept = FALSE + ), + # Identity which-indices per slot, generated programmatically + setNames( + lapply(seq_len(4L), function(k) + seq_len(data_for_model$ncol_X_random_eff[k]) |> as.array()), + paste0("X_random_effect_which_", seq_len(4L)) + ), + # Empty unseen design matrices per slot + setNames( + replicate(4L, matrix(0, nrow = nrow(data_for_model$X), ncol = 0), + simplify = FALSE), + paste0("X_random_effect_", seq_len(4L), "_unseen") + ) + ), parallel_chains = ifelse( inference_method %in% c("variational", "pathfinder") | @@ -233,8 +237,10 @@ sccomp_remove_outliers.sccomp_tbl = function(.estimate, attr(.estimate , "fit")$num_chains() ), threads_per_chain = cores, - seed = mcmc_seed, - sig_figs = sig_figs + sig_figs = sig_figs, + # GQ uses RNG (beta_binomial_rng). Seed with the posterior's sampling run so the + # predictive path is reproducible for a given `sccomp_estimate` draw set. + seed = attr(.estimate, "fit")$metadata()$seed ) @@ -330,7 +336,11 @@ sccomp_remove_outliers.sccomp_tbl = function(.estimate, verbose = verbose, seed = mcmc_seed, max_sampling_iterations = max_sampling_iterations, - pars = c("beta", "alpha", "prec_coeff", "prec_sd", "alpha_normalised", "random_effect", "random_effect_2"), + pars = c( + "beta", "alpha", + "prec_intercept_1", "prec_slope_1", "prec_intercept_2", "prec_slope_2", "prec_sd", + "random_effect_1", "random_effect_2", "random_effect_3", "random_effect_4" + ), sig_figs = sig_figs, cache_stan_model = cache_stan_model, ... @@ -342,39 +352,39 @@ sccomp_remove_outliers.sccomp_tbl = function(.estimate, fit2$draws(format = "matrix"), # This is for the new data generation with selected factors to do adjustment - data = data_for_model |> c(list( - - # Add subset of coefficients - X_original = data_for_model$X, - N_original = data_for_model$N, - length_X_which = ncol(data_for_model$X), - length_XA_which = ncol(data_for_model$XA), - X_which = seq_len(ncol(data_for_model$X)) |> as.array(), - XA_which = seq_len(ncol(data_for_model$Xa)) |> as.array(), - - # Random intercept common variable between grouping 1 and 2 - ncol_X_random_eff_new = ncol(data_for_model$X_random_effect) |> c(ncol(data_for_model$X_random_effect_2) ), # I could put this in the intial data - length_X_random_effect_which = ncol(data_for_model$X_random_effect) |> c(ncol(data_for_model$X_random_effect_2)), - - # Grouping 1 - X_random_effect_which = seq_len(ncol(data_for_model$X_random_effect)) |> as.array(), - - # Grouping 2 - Random intercept DUPLICATED - X_random_effect_which_2 = seq_len(ncol(data_for_model$X_random_effect_2)) |> as.array(), - - # Initialize unseen random effect variables - ncol_X_random_eff_unseen = c(0, 0), - X_random_effect_unseen = matrix(0, nrow = nrow(data_for_model$X), ncol = 0), - X_random_effect_2_unseen = matrix(0, nrow = nrow(data_for_model$X), ncol = 0), - - create_intercept = FALSE - - )), + data = data_for_model |> c( + list( + # Add subset of coefficients + X_original = data_for_model$X, + N_original = data_for_model$N, + length_X_which = ncol(data_for_model$X), + length_XA_which = ncol(data_for_model$XA), + X_which = seq_len(ncol(data_for_model$X)) |> as.array(), + XA_which = seq_len(ncol(data_for_model$Xa)) |> as.array(), + + # Per-slot random-effect pass-throughs (see notes in the first call site) + ncol_X_random_eff_new = data_for_model$ncol_X_random_eff, + length_X_random_effect_which = data_for_model$ncol_X_random_eff, + ncol_X_random_eff_unseen = rep(0L, 4L), + + create_intercept = FALSE + ), + setNames( + lapply(seq_len(4L), function(k) + seq_len(data_for_model$ncol_X_random_eff[k]) |> as.array()), + paste0("X_random_effect_which_", seq_len(4L)) + ), + setNames( + replicate(4L, matrix(0, nrow = nrow(data_for_model$X), ncol = 0), + simplify = FALSE), + paste0("X_random_effect_", seq_len(4L), "_unseen") + ) + ), parallel_chains = ifelse(inference_method %in% c("variational", "pathfinder"), 1, fit2$num_chains()), threads_per_chain = cores, - seed = mcmc_seed, - sig_figs = sig_figs + sig_figs = sig_figs, + seed = fit2$metadata()$seed ) @@ -463,7 +473,11 @@ sccomp_remove_outliers.sccomp_tbl = function(.estimate, verbose = verbose, seed = mcmc_seed, max_sampling_iterations = max_sampling_iterations, - pars = c("beta", "alpha", "prec_coeff","prec_sd", "alpha_normalised", "random_effect", "random_effect_2", "log_lik"), + pars = c( + "beta", "alpha", + "prec_intercept_1", "prec_slope_1", "prec_intercept_2", "prec_slope_2", "prec_sd", + "random_effect_1", "random_effect_2", "random_effect_3", "random_effect_4", "log_lik" + ), cache_stan_model = cache_stan_model, ... ) diff --git a/R/sccomp_replicate.R b/R/sccomp_replicate.R index 0c785feb..941e9b3f 100644 --- a/R/sccomp_replicate.R +++ b/R/sccomp_replicate.R @@ -126,8 +126,8 @@ sccomp_replicate.sccomp_tbl = function(fit, #' @param Xa Original variability design matrix #' @param N Original number of samples #' @param intercept_in_design Whether intercept is in design -#' @param X_random_effect Original random effect design matrix -#' @param X_random_effect_2 Original second random effect design matrix +#' @param X_random_effect_slots Length-4 list of original random-effect design +#' matrices (one per slot). Empty slots are zero-column matrices. #' @param .sample Quosure for the sample identifier column #' @param .cell_group Quosure for the cell group column #' @param .count Quosure for the count column @@ -135,13 +135,22 @@ sccomp_replicate.sccomp_tbl = function(fit, #' @param formula_variability Formula for the variability model #' @param new_data New data to generate predictions for. If NULL, uses the original data #' @param original_count_data Original count data from the model +#' @param smooth_results Optional named list of smooth-term metadata captured +#' at fit time. In the standard call path this is read from +#' `get_smooth_results(.data)` (an attribute on `model_input`, not Stan data). +#' Contains +#' `smooth_specs`, the `mgcv::smoothCon()` objects used to evaluate bases at +#' `new_data`; `smooth_re_objs`, the parallel `mgcv::smooth2random()` results +#' used to recover the fitted `Xf` / `Xr` parameterisation; and +#' `smooth_labels`, the original term labels used to name generated smooth +#' columns so they match the fit-time design. `NULL` when the composition +#' formula has no smooths. #' #' @return A list containing: #' - model_input: The prepared model input data #' - X_which: Indices for the composition design matrix #' - XA_which: Indices for the variability design matrix -#' - X_random_effect_which: Indices for the first random effect design matrix -#' - X_random_effect_which_2: Indices for the second random effect design matrix +#' - X_random_effect_which_1..4: per-slot indices into the original RE design matrix #' - create_intercept: Boolean indicating if intercept should be created #' #' @noRd @@ -149,8 +158,7 @@ prepare_replicate_data = function(X, Xa, N, intercept_in_design, - X_random_effect, - X_random_effect_2, + X_random_effect_slots, .sample, .cell_group, .count, @@ -158,7 +166,8 @@ prepare_replicate_data = function(X, original_formula_composition, formula_variability, new_data = NULL, - original_count_data) { + original_count_data, + smooth_results = NULL) { .sample = enquo(.sample) @@ -228,13 +237,17 @@ prepare_replicate_data = function(X, new_data = old_data |> bind_rows( new_data ) + # Smooth columns are evaluated separately via PredictMat below; keep the + # random-effect clauses so the later RE parsing still sees the same formula. + formula_composition = strip_smooth_terms(formula_composition) + new_X = new_data |> get_design_matrix( # Drop random intercept formula_composition |> + strip_random_effect_terms() |> as.character() |> - str_remove_all("\\+ ?\\(.+\\|.+\\)") |> paste(collapse="") |> as.formula(), !!.sample, @@ -244,6 +257,17 @@ prepare_replicate_data = function(X, # Remove columns that are not in the original design matrix .[,colnames(.) %in% colnames(X), drop=FALSE] + # Evaluate any smooth bases on the replicate rows and merge the resulting + # design pieces (unpenalised columns appended to `new_X`, one RE slot per + # penalised block) using the fit-time `smoothCon` / `smooth2random` objects. + smooth_design = build_smooth_replicate_design( + parametric_X = new_X, + new_data_tail = new_data |> tail(nrow_new_data), + smooth_results = smooth_results + ) + new_X = smooth_design$new_X + smooth_replicate_slots = smooth_design$smooth_replicate_slots + # Check that all effect combination were present when the model was fitted check_missing_parameters( new_X |> colnames(), @@ -265,8 +289,8 @@ prepare_replicate_data = function(X, get_design_matrix( # Drop random intercept formula_variability |> + strip_random_effect_terms() |> as.character() |> - str_remove_all("\\+ ?\\(.+\\|.+\\)") |> paste(collapse="") |> as.formula(), !!.sample, @@ -291,137 +315,130 @@ prepare_replicate_data = function(X, "(Intercept)" %in% colnames(new_X) if(create_intercept) warning("sccomp says: your estimated model is intercept free, while your desired replicated data do have an intercept term. The intercept estimate will be calculated averaging your first factor in your formula ~ 0 + . If you don't know the meaning of this warning, this is likely undesired, and please reconsider your formula for replicate_data()") - # Original grouping - original_grouping_names = original_formula_composition |> formula_to_random_effect_formulae() |> pull(grouping) + # Original grouping (one entry per RE clause in the original formula) + original_grouping_names = original_formula_composition |> + formula_to_random_effect_formulae() |> + pull(grouping) - # Random intercept + # Parse the (possibly new) composition formula to find RE clauses in new_data random_effect_elements = parse_formula_random_effect(formula_composition) - - # Initialize unseen random effect variables - new_X_random_effect_unseen = matrix(rep(0, nrow_new_data))[,0, drop=FALSE] - new_X_random_effect_2_unseen = matrix(rep(0, nrow_new_data))[,0, drop=FALSE] - - # Set default random intercept - X_random_effect_which = array()[0] - new_X_random_effect = matrix(rep(0, nrow_new_data))[,0, drop=FALSE] - - # setup default unknown_grouping variable for generated quantities - unknown_grouping = c(FALSE, FALSE) - - #check_random_effect_design(.data_spread, any_of(factor_names), random_effect_elements, formula, X) - random_effect_grouping = + random_effect_grouping = formula_composition |> formula_to_random_effect_formulae() |> mutate(design = map2( formula, grouping, - ~ get_random_effect_design3(new_data, .x, .y, !!.sample, - accept_NA_as_average_effect = TRUE ) + ~ get_random_effect_design3(new_data, .x, .y, !!.sample, + accept_NA_as_average_effect = TRUE) )) - - - if((random_effect_grouping$grouping %in% original_grouping_names[1]) |> any() && !unknown_grouping[1]) { - new_X_random_effect = - random_effect_grouping |> - filter(grouping==original_grouping_names[1]) |> - mutate(design_matrix = map( - design, - ~ ..1 |> - select(!!.sample, group___label, value) |> - - # Some combinations might not have been present in a specific group so the parameter does not exist - filter(group___label %in% colnames(X_random_effect)) |> - - pivot_wider(names_from = group___label, values_from = value) |> - mutate(across(everything(), ~ .x |> replace_na(0))) - )) |> - # Merge - pull(design_matrix) |> - _[[1]] |> - column_to_rownames(quo_name(.sample)) |> - tail(nrow_new_data) + # ---------------------------------------------------------------------- + # Build the per-slot replicate design matrices. + # + # For each of the 4 slots: if the slot was active in the original fit + # (original_grouping_names[k] exists) and the new formula references it, + # build a new design matrix restricted to the columns the model saw, and + # an index vector mapping new columns back to those of the original matrix. + # Otherwise emit an empty placeholder. + # ---------------------------------------------------------------------- + empty_mat_new = matrix(rep(0, nrow_new_data))[, 0, drop = FALSE] + empty_which_new = array()[0] + + build_replicate_slot = function(slot_idx) { + grouping_for_slot = original_grouping_names[slot_idx] # NA if slot wasn't used originally + X_original_slot = X_random_effect_slots[[slot_idx]] - # Separate NA group column into new_X_random_effect_unseen - new_X_random_effect_unseen = new_X_random_effect[, colnames(new_X_random_effect) |> str_detect("___NA$"), drop = FALSE] - new_X_random_effect = new_X_random_effect[, !colnames(new_X_random_effect) |> str_detect("___NA$"), drop = FALSE] + slot_is_active = !is.na(grouping_for_slot) && + any(random_effect_grouping$grouping %in% grouping_for_slot) - # Check that all effect combination were present when the model was fitted - check_missing_parameters( - new_X_random_effect |> colnames(), - X_random_effect |> colnames() - ) + if (!slot_is_active) { + return(list( + X = empty_mat_new, + X_unseen = empty_mat_new, + which = empty_which_new + )) + } - X_random_effect_which = - colnames(new_X_random_effect) |> - match( - X_random_effect %>% - colnames() - ) |> - as.array() - } - - # Set default X random intercept - X_random_effect_which_2 = array()[0] - new_X_random_effect_2 = matrix(rep(0, nrow_new_data))[,0, drop=FALSE] - - if((random_effect_grouping$grouping %in% original_grouping_names[2]) |> any()){ - new_X_random_effect_2 = - random_effect_grouping |> - filter(grouping==original_grouping_names[2]) |> + X_new = random_effect_grouping |> + filter(grouping == grouping_for_slot) |> mutate(design_matrix = map( design, ~ ..1 |> select(!!.sample, group___label, value) |> - - # Some combinations might not have been present in a specific group so the parameter does not exist - filter(group___label %in% colnames(X_random_effect_2)) |> - + # Combinations not present in the original fit have no parameters + filter(group___label %in% colnames(X_original_slot)) |> pivot_wider(names_from = group___label, values_from = value) |> mutate(across(everything(), ~ .x |> replace_na(0))) )) |> - # Merge - pull(design_matrix) |> - _[[1]] |> - column_to_rownames(quo_name(.sample)) |> + pull(design_matrix) |> + _[[1]] |> + column_to_rownames(quo_name(.sample)) |> tail(nrow_new_data) - # Separate NA group column into new_X_random_effect_2_unseen - new_X_random_effect_2_unseen = new_X_random_effect_2[, colnames(new_X_random_effect_2) |> str_detect("___NA$"), drop = FALSE] - new_X_random_effect_2 = new_X_random_effect_2[, !colnames(new_X_random_effect_2) |> str_detect("___NA$"), drop = FALSE] + is_NA_col = str_detect(colnames(X_new), "___NA$") + X_new_unseen = X_new[, is_NA_col, drop = FALSE] + X_new = X_new[, !is_NA_col, drop = FALSE] - # Check that all effect combination were present when the model was fitted - check_missing_parameters( - new_X_random_effect_2 |> colnames(), - X_random_effect_2 |> colnames() - ) + check_missing_parameters(colnames(X_new), colnames(X_original_slot)) - X_random_effect_which_2 = - colnames(new_X_random_effect_2) |> - match( - X_random_effect_2 %>% - colnames() - ) |> + which_idx = colnames(X_new) |> + match(colnames(X_original_slot)) |> as.array() + + list(X = X_new, X_unseen = X_new_unseen, which = which_idx) } - # Prepare the list of inputs to the model + replicate_slots = map(seq_len(4L), build_replicate_slot) + + # Append smooth-derived replicate slots (one per smooth term in the + # composition formula). They occupy whichever slots come after the + # explicit RE clauses, mirroring the slot ordering used at fit time. + if (length(smooth_replicate_slots) > 0) { + n_explicit_re = length(original_grouping_names) + n_smooth = length(smooth_replicate_slots) + n_used = n_explicit_re + n_smooth + if (n_used > 4L) { + stop(sprintf( + "sccomp says: the replicate model needs %d RE slot(s) but only 4 are available.", + n_used + )) + } + # Replace placeholder slots `[n_explicit_re + 1 .. n_used]` with smooths. + for (k in seq_len(n_smooth)) { + replicate_slots[[n_explicit_re + k]] = smooth_replicate_slots[[k]] + } + } + + # setup default unknown_grouping variable for generated quantities + unknown_grouping = rep(0L, 4L) + list( - X = new_X, - Xa = new_Xa, - N = nrow_new_data, + X = new_X, + Xa = new_Xa, + N = nrow_new_data, exposure = new_exposure, - X_random_effect = new_X_random_effect, - X_random_effect_2 = new_X_random_effect_2, - X_random_effect_unseen = new_X_random_effect_unseen, - X_random_effect_2_unseen = new_X_random_effect_2_unseen, - ncol_X_random_eff_new = c(ncol(new_X_random_effect), ncol(new_X_random_effect_2)), - unknown_grouping = unknown_grouping, - ncol_X_random_eff_unseen = c(ncol(new_X_random_effect_unseen), ncol(new_X_random_effect_2_unseen)), - X_which = X_which, - XA_which = XA_which, - X_random_effect_which = X_random_effect_which, - X_random_effect_which_2 = X_random_effect_which_2, + # Per-slot design + unseen-column matrices + index vectors + X_random_effect_1 = replicate_slots[[1]]$X, + X_random_effect_2 = replicate_slots[[2]]$X, + X_random_effect_3 = replicate_slots[[3]]$X, + X_random_effect_4 = replicate_slots[[4]]$X, + + X_random_effect_1_unseen = replicate_slots[[1]]$X_unseen, + X_random_effect_2_unseen = replicate_slots[[2]]$X_unseen, + X_random_effect_3_unseen = replicate_slots[[3]]$X_unseen, + X_random_effect_4_unseen = replicate_slots[[4]]$X_unseen, + + X_random_effect_which_1 = replicate_slots[[1]]$which, + X_random_effect_which_2 = replicate_slots[[2]]$which, + X_random_effect_which_3 = replicate_slots[[3]]$which, + X_random_effect_which_4 = replicate_slots[[4]]$which, + + ncol_X_random_eff_new = map_int(replicate_slots, ~ ncol(.x$X)), + ncol_X_random_eff_unseen = map_int(replicate_slots, ~ ncol(.x$X_unseen)), + unknown_grouping = unknown_grouping, + + X_which = X_which, + XA_which = XA_which, create_intercept = create_intercept ) @@ -474,14 +491,14 @@ replicate_data = function(.data, # create model input model_input = attr(.data, "model_input") - # Prepare data + # Prepare data - pass the per-slot original RE design matrices as a list prepared_data = prepare_replicate_data( X = model_input$X, Xa = model_input$Xa, N = model_input$N, intercept_in_design = model_input$intercept_in_design, - X_random_effect = model_input$X_random_effect, - X_random_effect_2 = model_input$X_random_effect_2, + X_random_effect_slots = lapply(seq_len(4L), function(k) + model_input[[paste0("X_random_effect_", k)]]), .sample = !!.sample, .cell_group = !!.cell_group, .count = !!.count, @@ -492,40 +509,39 @@ replicate_data = function(.data, original_count_data = .data |> attr("count_data") |> - .subset(!!.sample) + .subset(!!.sample), + smooth_results = get_smooth_results(.data) ) - # Original input + # Original input model_input$X_original = model_input$X model_input$N_original = model_input$N # New input - model_input$X = prepared_data$X - model_input$Xa = prepared_data$Xa - model_input$N = prepared_data$N + model_input$X = prepared_data$X + model_input$Xa = prepared_data$Xa + model_input$N = prepared_data$N model_input$exposure = prepared_data$exposure - model_input$X_random_effect = prepared_data$X_random_effect - model_input$X_random_effect_2 = prepared_data$X_random_effect_2 - model_input$X_random_effect_unseen = prepared_data$X_random_effect_unseen - model_input$X_random_effect_2_unseen = prepared_data$X_random_effect_2_unseen - model_input$ncol_X_random_eff_new = prepared_data$ncol_X_random_eff_new - model_input$unknown_grouping = prepared_data$unknown_grouping + + # Per-slot RE design + unseen + which-indices (4 slots) + for (k in seq_len(4L)) { + model_input[[paste0("X_random_effect_", k)]] = prepared_data[[paste0("X_random_effect_", k)]] + model_input[[paste0("X_random_effect_", k, "_unseen")]] = prepared_data[[paste0("X_random_effect_", k, "_unseen")]] + model_input[[paste0("X_random_effect_which_", k)]] = prepared_data[[paste0("X_random_effect_which_", k)]] + } + model_input$ncol_X_random_eff_new = prepared_data$ncol_X_random_eff_new model_input$ncol_X_random_eff_unseen = prepared_data$ncol_X_random_eff_unseen + model_input$unknown_grouping = prepared_data$unknown_grouping - # Add subset of coefficients - # Add subset of coefficients - model_input$length_X_which = length(prepared_data$X_which) + # Subset of coefficients + model_input$length_X_which = length(prepared_data$X_which) model_input$length_XA_which = length(prepared_data$XA_which) - model_input$X_which = prepared_data$X_which - model_input$XA_which = prepared_data$XA_which - - # Add random effect coefficients - model_input$X_random_effect_which = prepared_data$X_random_effect_which - model_input$X_random_effect_which_2 = prepared_data$X_random_effect_which_2 - model_input$length_X_random_effect_which = c( - length(prepared_data$X_random_effect_which), - length(prepared_data$X_random_effect_which_2) - ) + model_input$X_which = prepared_data$X_which + model_input$XA_which = prepared_data$XA_which + + # Length-4 vector of which-index lengths for the random-effect slots + model_input$length_X_random_effect_which = + map_int(seq_len(4L), ~ length(prepared_data[[paste0("X_random_effect_which_", .x)]])) # Should I create an intercept for generate quantities? model_input$create_intercept = prepared_data$create_intercept diff --git a/R/sccomp_test.R b/R/sccomp_test.R index 6da53c8a..f9208b84 100644 --- a/R/sccomp_test.R +++ b/R/sccomp_test.R @@ -87,63 +87,46 @@ sccomp_test.sccomp_tbl = function(.data, truncation_df2 = .data |> attr("truncation_df2") inference_method = .data |> attr("inference_method") - identified <- sccomp_identify_covariate_contrasts(contrasts, model_input) - mapped_covariates <- if (is.null(identified)) NULL else unique(identified$contrast_mapping$design_param) - contrasts_parameters <- NULL - if (!is.null(contrasts)) { - contrasts_parameters <- contrasts_to_parameter_list(contrasts) - } - # sccomp_test always computes pH0/FDR from posterior draws. - abundance_CI <- get_abundance_contrast_draws(.data, contrasts, mapped_covariates, contrasts_parameters) - - if ("parameter" %in% colnames(abundance_CI)) { - abundance_CI <- - abundance_CI |> + result <- + get_abundance_contrast_draws(.data, contrasts) |> draws_to_statistics( percent_false_positive / 100, test_composition_above_logit_fold_change, !!.cell_group, "c_" ) - } - - variability_CI <- get_variability_contrast_draws(.data, contrasts, mapped_covariates, contrasts_parameters) + + variability_draws <- get_variability_contrast_draws(.data, contrasts) - if ("parameter" %in% colnames(variability_CI)) { - variability_CI <- - variability_CI |> - draws_to_statistics( - percent_false_positive / 100, - test_composition_above_logit_fold_change, - !!.cell_group, - "v_" - ) + # If I have variability draws for those contrasts, compute the variability CI + if ("parameter" %in% colnames(variability_draws)) { + result = + result |> + left_join( + variability_draws |> + draws_to_statistics( + percent_false_positive / 100, + test_composition_above_logit_fold_change, + !!.cell_group, + "v_" + ) + ) } + - # If I don't have factors (~1) - if (!"factor" %in% colnames(model_input$factor_parameter_dictionary)) - factor_parameter_dictionary = tibble(`factor` = character(), design_matrix_col = character()) - else - factor_parameter_dictionary = + factor_parameter_dictionary = model_input$factor_parameter_dictionary |> select(`factor`, design_matrix_col) # Merge and parse result = - abundance_CI |> - - # Add ALPHA - left_join(variability_CI) |> - suppressMessages() |> - - # Add easy to understand factor labels - left_join(factor_parameter_dictionary, - by = c("parameter" = "design_matrix_col")) |> - select(parameter, `factor`, everything()) |> + result |> - select(!!.cell_group, everything(),-M) - + # Keep factor labels in the result schema; plotting/subsetting helpers rely on this. + left_join(factor_parameter_dictionary, by = c("parameter" = "design_matrix_col")) |> + select(!!.cell_group, parameter, `factor`, everything(), -M) + if(pass_fit) result = result |> @@ -189,50 +172,37 @@ sccomp_test.sccomp_tbl = function(.data, #' @noRd summarise_stan_matrix_for_estimate <- function( fit, - par_name, - design_colnames, - cell_group_levels, + model_input, + stan_parameter, + parameter_names, probs, - cell_group_colname, - prefix, - full_design_colnames = NULL, - full_block_parameter = NULL) { - if (is.null(full_design_colnames)) { - full_design_colnames <- design_colnames - } - # Translate design-matrix column names to Stan matrix column indices (C). - C_idx <- match(design_colnames, full_design_colnames) - if (anyNA(C_idx)) { - stop( - "sccomp says: each name in design_colnames must appear in full_design_colnames.", - call. = FALSE - ) - } - if (is.null(par_name)) { - summ <- summary_to_tibble(fit, full_block_parameter, "C", "M", probs = probs) - } else { - # Subset request uses explicit indexed Stan variable names for selected C across all M. - n_M <- length(cell_group_levels) - g <- expand.grid(C = C_idx, M = seq_len(n_M), stringsAsFactors = FALSE) - variable_names <- sprintf("%s[%d,%d]", par_name, g$C, g$M) - summ <- summary_to_tibble(fit, variable_names, "C", "M", probs = probs) - } - qcols <- names(summ)[grepl("%$", names(summ))] - q_bounds <- qcols |> + prefix) { + # Map Stan matrix column index M back to user-facing cell-group labels. + cell_group_levels <- colnames(model_input$y) + + # Query Stan summaries for this parameter family (beta/random effects). + summ <- summary_to_tibble(fit, stan_parameter, "C", "M", probs = probs) + + # Identify lower/upper quantile column names from the summary output. + quantile_columns <- names(summ)[grepl("%$", names(summ))] + ordered_quantile_columns <- quantile_columns |> tibble::enframe(value = "qcol") |> dplyr::mutate(pct = as.numeric(gsub("%$", "", qcol, perl = TRUE))) |> dplyr::arrange(pct) |> dplyr::pull(qcol) - qlo <- q_bounds[[1]] - qhi <- q_bounds[[2]] + lower_quantile_column <- ordered_quantile_columns[[1]] + upper_quantile_column <- ordered_quantile_columns[[2]] + # Provide a stable mapping from C indices to design-matrix parameter names. par_by_C <- - tibble::tibble(C = C_idx, parameter = design_colnames) |> + tibble::tibble(C = seq_along(parameter_names), parameter = parameter_names) |> dplyr::distinct() + # Attach parameter labels to each C/M summary row. joined <- summ |> dplyr::left_join(par_by_C, by = "C") + # Keep output column names configurable for c_/v_ style reuse. lower_col <- paste0(prefix, "lower") effect_col <- paste0(prefix, "effect") upper_col <- paste0(prefix, "upper") @@ -240,14 +210,119 @@ summarise_stan_matrix_for_estimate <- function( ess_bulk_col <- paste0(prefix, "ess_bulk") ess_tail_col <- paste0(prefix, "ess_tail") + # Return compact, user-facing summary table with diagnostics. joined |> dplyr::transmute( - !!rlang::sym(cell_group_colname) := cell_group_levels[M], + cell_group = cell_group_levels[M], M, parameter, - !!lower_col := .data[[qlo]], + !!lower_col := .data[[lower_quantile_column]], !!effect_col := mean, - !!upper_col := .data[[qhi]], + !!upper_col := .data[[upper_quantile_column]], + !!rhat_col := rhat, + !!ess_bulk_col := ess_bulk, + !!ess_tail_col := ess_tail + ) +} + +#' Build variability summaries from R-side alpha normalisation draws. +#' +#' This mirrors `summarise_stan_matrix_for_estimate()` but computes quantiles for +#' alpha_normalised from derived draws, while borrowing convergence diagnostics +#' (rhat/ESS) from the corresponding `alpha` elements. +#' +#' @keywords internal +#' @noRd +get_variability_draws_for_test <- function(fit, model_input, alpha_variable_subset = NULL) { + + if (isTRUE(as.logical(model_input$exclude_mean_variability_association))) { + alpha_var <- if (is.null(alpha_variable_subset)) "alpha" else alpha_variable_subset + return( + draws_to_tibble_x_y(fit, alpha_var, "C", "M") |> + dplyr::mutate(.variable = "alpha_normalised") + ) + } + + compute_alpha_normalised_draws( + fit = fit, + model_input = model_input, + alpha_variable_subset = alpha_variable_subset + ) +} + +#' +#' @keywords internal +#' @noRd +summarise_alpha_normalised_for_estimate <- function( + fit, + model_input, + design_colnames, + probs, + cell_group_colname, + prefix) { + # Map Stan matrix column index M back to user-facing cell-group labels. + cell_group_levels <- colnames(model_input$y) + + # Match variability design columns to Stan C indices. + C_idx <- match(design_colnames, colnames(model_input$XA)) + if (anyNA(C_idx)) { + stop( + "sccomp says: each variability design column must appear in model_input$XA.", + call. = FALSE + ) + } + + n_M <- length(cell_group_levels) + g <- expand.grid(C = C_idx, M = seq_len(n_M), stringsAsFactors = FALSE) + alpha_subset <- sprintf("alpha[%d,%d]", g$C, g$M) + + # Compute variability summaries from derived R-side draws. + # If mean-variability association is excluded, this returns raw alpha draws + # (relabelled for downstream schema parity); otherwise alpha_normalised draws. + draws_summary <- get_variability_draws_for_test( + fit = fit, + model_input = model_input, + alpha_variable_subset = alpha_subset + ) |> + dplyr::group_by(C, M) |> + dplyr::summarise( + effect = mean(.value), + lower_quantile = stats::quantile(.value, probs = probs[[1]], na.rm = TRUE), + upper_quantile = stats::quantile(.value, probs = probs[[2]], na.rm = TRUE), + .groups = "drop" + ) + + # Pull diagnostics from the corresponding base alpha parameters in Stan. + alpha_diagnostics <- summary_to_tibble(fit, alpha_subset, "C", "M", probs = probs) |> + dplyr::select(C, M, rhat, ess_bulk, ess_tail) |> + dplyr::distinct() + + # Combine derived summaries, diagnostics, and design parameter labels. + joined <- draws_summary |> + dplyr::left_join(alpha_diagnostics, by = c("C", "M")) |> + dplyr::left_join( + tibble::tibble(C = C_idx, parameter = design_colnames) |> + dplyr::distinct(), + by = "C" + ) + + # Keep output column names configurable for v_ style summaries. + lower_col <- paste0(prefix, "lower") + effect_col <- paste0(prefix, "effect") + upper_col <- paste0(prefix, "upper") + rhat_col <- paste0(prefix, "rhat") + ess_bulk_col <- paste0(prefix, "ess_bulk") + ess_tail_col <- paste0(prefix, "ess_tail") + + # Return compact, user-facing summary table with diagnostics. + joined |> + dplyr::transmute( + !!rlang::sym(cell_group_colname) := cell_group_levels[M], + M, + parameter, + !!lower_col := lower_quantile, + !!effect_col := effect, + !!upper_col := upper_quantile, !!rhat_col := rhat, !!ess_bulk_col := ess_bulk, !!ess_tail_col := ess_tail @@ -267,38 +342,43 @@ sccomp_summarise_posterior_for_estimate <- function( fp <- percent_false_positive / 100 probs <- c(fp / 2, 1 - fp / 2) - cell_group_levels <- model_input$y |> colnames() - abundance_parts <- list( summarise_stan_matrix_for_estimate( - fit, NULL, colnames(model_input$X), cell_group_levels, - probs, cg, "c_", full_block_parameter = "beta" + fit = fit, + model_input = model_input, + stan_parameter = "beta", + parameter_names = colnames(model_input$X), + probs = probs, + prefix = "c_" ) ) - if (model_input$n_random_eff > 0) { + # Random effect blocks: append a summary for each non-empty slot (1..4). + for (k in seq_len(4L)) { + if (model_input$ncol_X_random_eff[k] == 0) next + X_slot <- model_input[[paste0("X_random_effect_", k)]] abundance_parts <- c( abundance_parts, list(summarise_stan_matrix_for_estimate( - fit, NULL, colnames(model_input$X_random_effect), cell_group_levels, - probs, cg, "c_", full_block_parameter = "random_effect" - )) - ) - } - if (model_input$n_random_eff > 1) { - abundance_parts <- c( - abundance_parts, - list(summarise_stan_matrix_for_estimate( - fit, NULL, colnames(model_input$X_random_effect_2), cell_group_levels, - probs, cg, "c_", full_block_parameter = "random_effect_2" + fit = fit, + model_input = model_input, + stan_parameter = paste0("random_effect_", k), + parameter_names = colnames(X_slot), + probs = probs, + prefix = "c_" )) ) } - abundance <- dplyr::bind_rows(abundance_parts) + abundance <- dplyr::bind_rows(abundance_parts) |> + dplyr::rename(!!cg := cell_group) - variability <- summarise_stan_matrix_for_estimate( - fit, NULL, colnames(model_input$XA), cell_group_levels, - probs, cg, "v_", full_block_parameter = "alpha_normalised" + variability <- summarise_alpha_normalised_for_estimate( + fit = fit, + model_input = model_input, + design_colnames = colnames(model_input$XA), + probs = probs, + cell_group_colname = cg, + prefix = "v_" ) |> dplyr::mutate( v_lower = -v_lower, @@ -307,19 +387,13 @@ sccomp_summarise_posterior_for_estimate <- function( ) |> dplyr::rename(v_lower = v_upper, v_upper = v_lower) - if (!"factor" %in% colnames(model_input$factor_parameter_dictionary)) { - factor_parameter_dictionary <- - tibble::tibble(`factor` = character(), design_matrix_col = character()) - } else { - factor_parameter_dictionary <- - model_input$factor_parameter_dictionary |> - dplyr::select(`factor`, design_matrix_col) - } + factor_parameter_dictionary <- + model_input$factor_parameter_dictionary |> + dplyr::select(`factor`, design_matrix_col) result <- abundance |> dplyr::left_join(variability, by = c(cg, "M", "parameter")) |> - suppressMessages() |> dplyr::left_join( factor_parameter_dictionary, by = c("parameter" = "design_matrix_col") @@ -373,8 +447,117 @@ sccomp_identify_covariate_contrasts <- function(contrasts, model_input) { ) } +#' Build Stan indexed-variable subset from contrast terms. +#' +#' @param contrasts Character vector of contrast expressions. Can be `NULL`. +#' @param design_columns Character vector of design-matrix column names for the target parameter block. +#' @param stan_parameter Character scalar Stan parameter name (for example, `"beta"` or `"alpha"`). +#' @param model_input Model input list containing `y`, used to infer the `M` index size. +#' +#' @keywords internal +#' @noRd +build_stan_parameter_subset <- function(contrasts, design_columns, stan_parameter, model_input) { + design_columns <- as.character(design_columns) + + if (is.null(contrasts)) { + return( + tibble::tibble( + parameter = design_columns, + variable = rep(stan_parameter, length(design_columns)) + ) + ) + } + + candidate_terms <- contrasts_to_parameter_list(contrasts) + if (is.null(candidate_terms) || length(candidate_terms) == 0) { + return(tibble::tibble(parameter = character(), variable = character())) + } + + candidate_terms <- unique(candidate_terms) + candidate_terms <- candidate_terms[!is.na(candidate_terms) & candidate_terms != ""] + if (length(candidate_terms) == 0) { + return(tibble::tibble(parameter = character(), variable = character())) + } + + matched <- intersect(candidate_terms, design_columns) + if (length(matched) == 0) { + return(tibble::tibble(parameter = character(), variable = character())) + } + + n_M <- ncol(model_input$y) + C_idx <- match(matched, design_columns) + g <- expand.grid(C = C_idx, M = seq_len(n_M), stringsAsFactors = FALSE) + + tibble::tibble( + parameter = matched[g$C], + variable = sprintf("%s[%d,%d]", stan_parameter, g$C, g$M) + ) +} + +# ---------------------------------------------------------------------- +# Random effect draws: extract one slot at a time (1..4) and left-join into +# `draws`. Per-slot logic is identical, so we loop over a helper instead of +# duplicating the block once per slot. +# ---------------------------------------------------------------------- +# This is still inefficient as it is drawing all random effects regardless of the ciontrasts +add_random_effect_draws = function(draws, contrasts, model_input, slot_idx, fit) { + + re_covariates = model_input[[paste0("X_random_effect_", slot_idx)]] |> colnames() + re_subset = build_stan_parameter_subset( + contrasts = contrasts, + design_columns = re_covariates, + stan_parameter = paste0("random_effect_", slot_idx), + model_input = model_input + ) + re_parameters = re_subset |> dplyr::pull("parameter") |> unique() + re_variables = re_subset |> dplyr::pull("variable") |> unique() + + # No-op when contrasts filter out everything in this slot + if (!is.null(contrasts) && length(re_parameters) == 0) + return(list(draws = draws, covariates = re_covariates)) + + re_draws = + fit |> + draws_to_tibble_x_y(re_variables, "C", "M") + + # Reconstruct the omitted last cell-group (sum-to-zero closure) + other_group_re = + re_draws |> + with_groups(c(C, .chain, .iteration, .draw, .variable), + ~ .x |> summarise(.value = sum(.value))) |> + mutate(.value = -.value, M = re_draws |> pull(M) |> max() + 1) + + re_draws = bind_rows(re_draws, other_group_re) + + if (!is.null(contrasts)) { + re_draws = re_draws |> + left_join( + re_covariates |> enframe(name = "C", value = "parameters_name"), + by = "C" + ) |> + filter(parameters_name %in% re_parameters) |> + select(-C) |> + pivot_wider(names_from = parameters_name, values_from = .value) + } else { + re_draws = re_draws |> + pivot_wider(names_from = C, values_from = .value) %>% + setNames(colnames(.)[1:5] |> c(re_covariates)) + } + + new_draws = + if (nrow(draws) == 0) + select(re_draws, -.variable) + else + draws |> + left_join(select(re_draws, -.variable), + by = c("M", ".chain", ".iteration", ".draw")) + + list(draws = new_draws, covariates = re_covariates) +} + + # this can be helpful if we want to draw PCA with uncertainty -get_abundance_contrast_draws = function(.data, contrasts, design_param_subset = NULL, contrasts_parameters = NULL){ +get_abundance_contrast_draws = function(.data, contrasts = NULL){ # Define the variables as NULL to avoid CRAN NOTES X <- NULL @@ -390,307 +573,80 @@ get_abundance_contrast_draws = function(.data, contrasts, design_param_subset = .cell_group = .data |> attr(".cell_group") + model_input <- .data |> attr("model_input") + cell_index_map <- + model_input %$% + y %>% + colnames() |> + enframe(name = "M", value = quo_name(.cell_group)) # Beta - beta_factor_of_interest = .data |> attr("model_input") %$% X |> colnames() - # beta = - # .data |> - # attr("fit") %>% - # draws_to_tibble_x_y("beta", "C", "M") |> - # pivot_wider(names_from = C, values_from = .value) %>% - # setNames(colnames(.)[1:5] |> c(beta_factor_of_interest)) - beta_variable_subset <- NULL - if (!is.null(design_param_subset)) { - needed <- intersect(design_param_subset, beta_factor_of_interest) - if (length(needed) > 0) { - n_M <- ncol(.data |> attr("model_input") %$% y) - C_idx <- match(needed, beta_factor_of_interest) - g <- expand.grid(C = C_idx, M = seq_len(n_M), stringsAsFactors = FALSE) - beta_variable_subset <- sprintf("beta[%d,%d]", g$C, g$M) - } - } - - if(contrasts |> is.null()) - draws = - .data |> - attr("fit") %>% - draws_to_tibble_x_y("beta", "C", "M") |> - pivot_wider(names_from = C, values_from = .value) %>% - setNames(colnames(.)[1:5] |> c(beta_factor_of_interest)) |> - select(-.variable) - - else if((beta_factor_of_interest %in% contrasts_parameters) |> which() |> length() > 0) + beta_covariates = model_input %$% X |> colnames() + beta_subset <- build_stan_parameter_subset( + contrasts = contrasts, + design_columns = beta_covariates, + stan_parameter = "beta", + model_input = model_input + ) + beta_parameters <- beta_subset |> dplyr::pull("parameter") |> unique() + beta_variable_subset <- beta_subset |> dplyr::pull("variable") |> unique() + + if(length(beta_variable_subset) == 0 ) + warning(sprintf("sccomp says: These components of your contrasts are not present in the model as parameters: %s. Factors including special characters, e.g. \"(Intercept)\" require backquotes e.g. \"`(Intercept)`\" ", paste(contrasts, sep = ", "))) + draws = .data |> attr("fit") %>% - draws_to_tibble_x_y( - if (is.null(beta_variable_subset)) "beta" else beta_variable_subset, - "C", - "M" - ) |> + draws_to_tibble_x_y( beta_variable_subset, "C", "M") |> left_join( - beta_factor_of_interest |> enframe(name = "C", value = "parameters_name"), + beta_covariates |> enframe(name = "C", value = "parameters_name"), by = "C" ) |> - filter(parameters_name %in% contrasts_parameters) |> select(-C) |> pivot_wider(names_from = parameters_name, values_from = .value) |> select(-.variable) - else - draws = tibble() - - - # Random effect - - beta_random_effect_factor_of_interest = .data |> attr("model_input") %$% X_random_effect |> colnames() - - if( - .data |> attr("model_input") %$% n_random_eff > 0 && - ( - contrasts |> is.null() || - (beta_random_effect_factor_of_interest %in% contrasts_parameters) |> which() |> length() > 0 - ) - ){ - - - beta_random_effect = - .data |> - attr("fit") %>% - draws_to_tibble_x_y("random_effect", "C", "M") - - # Add last component - other_group_random_effect = - beta_random_effect |> - with_groups(c(C, .chain, .iteration, .draw, .variable ), ~ .x |> summarise(.value = sum(.value))) |> - mutate(.value = -.value, M = beta_random_effect |> pull(M) |> max() + 1) - - - beta_random_effect = - beta_random_effect |> - bind_rows( other_group_random_effect ) - - - # Reshape - # Speed up if I have contrasts - if(!contrasts |> is.null()) - beta_random_effect = - beta_random_effect |> - left_join( - beta_random_effect_factor_of_interest |> enframe(name = "C", value = "parameters_name"), - by = "C" - ) |> - filter(parameters_name %in% contrasts_parameters) |> - select(-C) |> - pivot_wider(names_from = parameters_name, values_from = .value) - - else - beta_random_effect = - beta_random_effect |> - pivot_wider(names_from = C, values_from = .value) %>% - setNames(colnames(.)[1:5] |> c(beta_random_effect_factor_of_interest)) - - # If I don't have fix nor 1st level random effect - if(draws |> nrow() == 0) - draws = select(beta_random_effect, -.variable) - else - draws = draws |> - left_join(select(beta_random_effect, -.variable), - by = c("M", ".chain", ".iteration", ".draw") - ) - - } else { - beta_random_effect_factor_of_interest = "" - } - - # Second random effect. IN THE FUTURE THIS WILL BE VECTORISED TO ARBUTRARY GRI+OUING - beta_random_effect_factor_of_interest_2 = .data |> attr("model_input") %$% X_random_effect_2 |> colnames() - - if( - .data |> attr("model_input") %$% n_random_eff > 1 && - ( - contrasts |> is.null() || - (beta_random_effect_factor_of_interest_2 %in% contrasts_parameters) |> which() |> length() > 0 - ) - ){ - - beta_random_effect_2 = - .data |> - attr("fit") %>% - draws_to_tibble_x_y("random_effect_2", "C", "M") - - # Add last component - other_group_random_effect = - beta_random_effect_2 |> - with_groups(c(C, .chain, .iteration, .draw, .variable ), ~ .x |> summarise(.value = sum(.value))) |> - mutate(.value = -.value, M = beta_random_effect_2 |> pull(M) |> max() + 1) - - beta_random_effect_2 = - beta_random_effect_2 |> - bind_rows( other_group_random_effect ) - - # Reshape - # Speed up if I have contrasts - if(!contrasts |> is.null()) - beta_random_effect_2 = - beta_random_effect_2 |> - left_join( - beta_random_effect_factor_of_interest_2 |> enframe(name = "C", value = "parameters_name"), - by = "C" - ) |> - filter(parameters_name %in% contrasts_parameters) |> - select(-C) |> - pivot_wider(names_from = parameters_name, values_from = .value) - - else - beta_random_effect_2 = - beta_random_effect_2 |> - pivot_wider(names_from = C, values_from = .value) %>% - setNames(colnames(.)[1:5] |> c(beta_random_effect_factor_of_interest_2)) - - # If I don't have fix nor 1st level random effect - if(draws |> nrow() == 0) - draws = select(beta_random_effect_2, -.variable) - else - draws = draws |> - left_join(select(beta_random_effect_2, -.variable), - by = c("M", ".chain", ".iteration", ".draw") - ) - } else { - beta_random_effect_factor_of_interest_2 = "" - } - - - # If I have constrasts calculate - if(!is.null(contrasts)) - draws = - draws |> - mutate_from_expr_list(contrasts, ignore_errors = FALSE) |> - select(- any_of(c(beta_factor_of_interest, beta_random_effect_factor_of_interest) |> setdiff(contrasts)) ) - - # Add cell name - draws = draws |> - left_join( - .data |> - attr("model_input") %$% - y %>% - colnames() |> - enframe(name = "M", value = quo_name(.cell_group)), - by = "M" - ) %>% - select(!!.cell_group, everything()) - - - # If no contrasts of interest just return an empty data frame - if(ncol(draws)==5) return(draws |> distinct(M, !!.cell_group)) - - # Get convergence for fixed effects - convergence_df = - .data |> - attr("fit") |> - summary_to_tibble("beta", "C", "M") |> - - # Add cell name - left_join( - .data |> - attr("model_input") %$% - y %>% - colnames() |> - enframe(name = "M", value = quo_name(.cell_group)), - by = "M" - ) |> - - # factor names - left_join( - beta_factor_of_interest |> - enframe(name = "C", value = "parameter"), - by = "C" - ) - - # Get convergence for random effects if they exist - if(.data |> attr("model_input") %$% n_random_eff > 0) { - convergence_df_random = - .data |> - attr("fit") |> - summary_to_tibble("random_effect", "C", "M") |> - - # Add cell name - left_join( - .data |> - attr("model_input") %$% - y %>% - colnames() |> - enframe(name = "M", value = quo_name(.cell_group)), - by = "M" - ) |> - - # factor names - left_join( - beta_random_effect_factor_of_interest |> - enframe(name = "C", value = "parameter"), - by = "C" - ) + # RANDOM EFFECTS - # Combine fixed and random effects convergence - convergence_df = bind_rows(convergence_df, convergence_df_random) + random_effect_covariates_all = character(0) + + for (k in seq_len(4L)) { + if (model_input$ncol_X_random_eff[k] == 0) next + res <- add_random_effect_draws(draws, contrasts, model_input, k, attr(.data, "fit")) + draws <- res$draws + random_effect_covariates_all <- c(random_effect_covariates_all, res$covariates) } - # Get convergence for second random effect if it exists - if(.data |> attr("model_input") %$% n_random_eff > 1) { - convergence_df_random_2 = - .data |> - attr("fit") |> - summary_to_tibble("random_effect_2", "C", "M") |> - - # Add cell name - left_join( - .data |> - attr("model_input") %$% - y %>% - colnames() |> - enframe(name = "M", value = quo_name(.cell_group)), - by = "M" - ) |> - - # factor names - left_join( - beta_random_effect_factor_of_interest_2 |> - enframe(name = "C", value = "parameter"), - by = "C" - ) + # CALCULATE CONTRAST + if (!is.null(contrasts)){ + draws = + draws |> + mutate_from_expr_list(contrasts, ignore_errors = FALSE) |> + select(- any_of(c(beta_covariates, random_effect_covariates_all) |> setdiff(contrasts))) - # Combine with existing convergence data - convergence_df = bind_rows(convergence_df, convergence_df_random_2) } - - # if ("Rhat" %in% colnames(convergence_df)) { - # convergence_df <- rename(convergence_df, R_k_hat = Rhat) - # } else if ("khat" %in% colnames(convergence_df)) { - # convergence_df <- rename(convergence_df, R_k_hat = khat) - # } - - - convergence_df = - convergence_df |> - select(!!.cell_group, parameter, any_of(c("n_eff", "R_k_hat", "rhat", "ess_bulk", "ess_tail"))) |> - suppressWarnings() - - draws |> - pivot_longer(-c(1:5), names_to = "parameter", values_to = ".value") |> - - # Attach convergence if I have no contrasts - left_join(convergence_df, by = c(quo_name(.cell_group), "parameter")) |> + + # RESHAPE + draws = draws |> + pivot_longer(-c(1:4), names_to = "parameter", values_to = ".value") |> # Reorder because pivot long is bad - mutate(parameter = parameter |> fct_relevel(colnames(draws)[-c(1:5)])) |> + mutate(parameter = parameter |> fct_relevel(colnames(draws)[-c(1:4)])) |> arrange(parameter) + + # Add cell name + draws = draws |> + left_join(cell_index_map, by = "M" ) %>% + select(!!.cell_group, everything()) + + draws } #' @importFrom forcats fct_relevel #' @noRd -get_variability_contrast_draws = function(.data, contrasts, design_param_subset = NULL, contrasts_parameters = NULL){ +get_variability_contrast_draws = function(.data, contrasts){ # Define the variables as NULL to avoid CRAN NOTES XA <- NULL @@ -703,33 +659,34 @@ get_variability_contrast_draws = function(.data, contrasts, design_param_subset R_k_hat <- NULL .cell_group = .data |> attr(".cell_group") - - variability_factor_of_interest = .data |> attr("model_input") %$% XA |> colnames() - alpha_variable_subset <- NULL - if (!is.null(design_param_subset)) { - needed <- intersect(design_param_subset, variability_factor_of_interest) - if (length(needed) == 0 && !is.null(contrasts)) { - return( - .data |> - attr("model_input") %$% - y %>% - colnames() |> - enframe(name = "M", value = quo_name(.cell_group)) |> - dplyr::select(!!.cell_group, M) - ) - } - n_M <- ncol(.data |> attr("model_input") %$% y) - C_idx <- match(needed, variability_factor_of_interest) - g <- expand.grid(C = C_idx, M = seq_len(n_M), stringsAsFactors = FALSE) - alpha_variable_subset <- sprintf("alpha_normalised[%d,%d]", g$C, g$M) + model_input <- .data |> attr("model_input") + cell_index_map <- + model_input %$% + y %>% + colnames() |> + enframe(name = "M", value = quo_name(.cell_group)) + + variability_covariates = model_input %$% XA |> colnames() + alpha_subset <- build_stan_parameter_subset( + contrasts = contrasts, + design_columns = variability_covariates, + stan_parameter = "alpha", + model_input = model_input + ) + alpha_parameters <- alpha_subset |> dplyr::pull("parameter") |> unique() + alpha_variable_subset <- alpha_subset |> dplyr::pull("variable") |> unique() + if (length(alpha_variable_subset) == 0) alpha_variable_subset <- NULL + if (!is.null(contrasts) && length(alpha_parameters) == 0) { + return(cell_index_map |> dplyr::select(!!.cell_group, M)) } draws = - - .data |> - attr("fit") %>% - draws_to_tibble_x_y(if (is.null(alpha_variable_subset)) "alpha_normalised" else alpha_variable_subset, "C", "M") |> - + get_variability_draws_for_test( + fit = .data |> attr("fit"), + model_input = model_input, + alpha_variable_subset = alpha_variable_subset + ) |> + # We want variability, not concentration mutate(.value = -.value) @@ -739,10 +696,9 @@ get_variability_contrast_draws = function(.data, contrasts, design_param_subset draws = draws |> left_join( - variability_factor_of_interest |> enframe(name = "C", value = "parameters_name"), + variability_covariates |> enframe(name = "C", value = "parameters_name"), by = "C" ) |> - filter(parameters_name %in% contrasts_parameters) |> select(-C) |> pivot_wider(names_from = parameters_name, values_from = .value) |> select( -.variable) @@ -751,71 +707,46 @@ get_variability_contrast_draws = function(.data, contrasts, design_param_subset draws = draws |> pivot_wider(names_from = C, values_from = .value) %>% - setNames(colnames(.)[1:5] |> c(variability_factor_of_interest)) |> + setNames(colnames(.)[1:5] |> c(variability_covariates)) |> select( -.variable) # If I have constrasts calculate if (!is.null(contrasts)) draws <- mutate_from_expr_list(draws, contrasts, ignore_errors = TRUE) - draws = draws |> - - # Add cell name - left_join( - .data |> attr("model_input") %$% - y %>% - colnames() |> - enframe(name = "M", value = quo_name(.cell_group)), - by = "M" - ) %>% - select(!!.cell_group, everything()) - # If no contrasts of interest just return an empty data frame - if(ncol(draws)==5) return(draws |> distinct(M, !!.cell_group)) + if(ncol(draws) <= 4) return(cell_index_map |> dplyr::select(!!.cell_group, M) |> distinct()) - # Get convergence - convergence_df = - .data |> - attr("fit") |> - summary_to_tibble("alpha_normalised", "C", "M") |> - - # Add cell name - left_join( - .data |> - attr("model_input") %$% - y %>% - colnames() |> - enframe(name = "M", value = quo_name(.cell_group)), - by = "M" - ) |> - - # factor names - left_join( - variability_factor_of_interest |> - enframe(name = "C", value = "parameter"), - by = "C" - ) - - convergence_df = - convergence_df |> - select(!!.cell_group, parameter, any_of(c("n_eff", "R_k_hat", "rhat", "ess_bulk", "ess_tail"))) |> - suppressWarnings() - - - draws |> - pivot_longer(-c(1:5), names_to = "parameter", values_to = ".value") |> - - # Attach convergence if I have no contrasts - left_join(convergence_df, by = c(quo_name(.cell_group), "parameter")) |> - + draws = draws |> + pivot_longer(-c(1:4), names_to = "parameter", values_to = ".value") |> + # Reorder because pivot long is bad - mutate(parameter = parameter |> fct_relevel(colnames(draws)[-c(1:5)])) |> + mutate(parameter = parameter |> fct_relevel(colnames(draws)[-c(1:4)])) |> arrange(parameter) + + draws = draws |> + left_join(cell_index_map, by = "M") %>% + select(!!.cell_group, everything()) + + draws } #' Mutate Data Frame Based on Expression List #' +#' @noRd +add_missing_contrast_names = function(formula_expr){ + contrast_names = names(formula_expr) + if (is.null(contrast_names)) { + contrast_names = formula_expr + } else { + missing_names = is.na(contrast_names) | contrast_names == "" + contrast_names[missing_names] = formula_expr[missing_names] + } + + make.unique(contrast_names) +} + #' @description #' `mutate_from_expr_list` takes a data frame and a list of formula expressions, #' and mutates the data frame based on these expressions. It allows for ignoring @@ -843,8 +774,7 @@ get_variability_contrast_draws = function(.data, contrasts, design_param_subset #' mutate_from_expr_list = function(x, formula_expr, ignore_errors = TRUE){ - if(formula_expr |> names() |> is.null()) - names(formula_expr) = formula_expr + names(formula_expr) = add_missing_contrast_names(formula_expr) # Creating a named vector where the names are the strings to be replaced # and the values are empty strings @@ -852,7 +782,7 @@ mutate_from_expr_list = function(x, formula_expr, ignore_errors = TRUE){ # Check if all elements of contrasts are in the parameter parameter_names = x |> colnames() - + # Check is backquoted are not used require_back_quotes = !contrasts_elements |> str_remove_all("`") |> contains_only_valid_chars_for_column() has_left_back_quotes = contrasts_elements |> str_detect("^`") @@ -873,6 +803,7 @@ mutate_from_expr_list = function(x, formula_expr, ignore_errors = TRUE){ if(length(contrasts_not_in_the_model) > 0 & !ignore_errors) warning(sprintf("sccomp says: These components of your contrasts are not present in the model as parameters: %s. Factors including special characters, e.g. \"(Intercept)\" require backquotes e.g. \"`(Intercept)`\" ", paste(contrasts_not_in_the_model, sep = ", "))) + # Calculate if(ignore_errors) my_mutate = mutate_ignore_error else my_mutate = mutate @@ -891,10 +822,17 @@ mutate_from_expr_list = function(x, formula_expr, ignore_errors = TRUE){ } +#' @importFrom dplyr pull +#' @importFrom posterior as_draws_df summarise_draws rhat ess_bulk ess_tail +#' @noRd draws_to_statistics = function(draws, false_positive_rate, test_composition_above_logit_fold_change, .cell_group, prefix = ""){ # Define the variables as NULL to avoid CRAN NOTES M <- NULL + .chain <- NULL + .iteration <- NULL + .draw <- NULL + .value <- NULL parameter <- NULL bigger_zero <- NULL smaller_zero <- NULL @@ -907,27 +845,55 @@ draws_to_statistics = function(draws, false_positive_rate, test_composition_abov R_k_hat <- NULL .cell_group = enquo(.cell_group) - - draws = + lower_prob = false_positive_rate / 2 + upper_prob = 1 - lower_prob + + draw_summaries = + draws |> + group_by(!!.cell_group, M, parameter) |> + group_modify( + ~ { + draws_df = + .x |> + select(.chain, .iteration, .draw, value = .value) |> + as_draws_df() |> + summarise_draws( + lower = function(.x) stats::quantile(.x, probs = lower_prob, na.rm = TRUE), + effect = function(.x) mean(.x, na.rm = TRUE), + upper = function(.x) stats::quantile(.x, probs = upper_prob, na.rm = TRUE), + rhat, + ess_bulk, + ess_tail + ) |> + select(-variable) + + # Standardise CI column names when posterior returns percentage names (e.g. `5%`, `97.5%`). + quantile_cols = names(draws_df)[grepl("%$", names(draws_df))] + names(draws_df)[names(draws_df) %in% quantile_cols] = c("lower", "upper") + + draws_df + } + ) |> + ungroup() + + probability_stats = draws %>% - group_by(!!.cell_group, M, parameter, rhat, ess_bulk, ess_tail) %>% + group_by(!!.cell_group, M, parameter) %>% summarise( - lower = quantile(.value, false_positive_rate / 2), - effect = mean(.value), - upper = quantile(.value, 1 - (false_positive_rate / 2)), - bigger_zero = sum(.value > test_composition_above_logit_fold_change), - smaller_zero = sum(.value < -test_composition_above_logit_fold_change), - # R_k_hat = unique(R_k_hat), - # n_eff = unique(n_eff), + bigger_zero = sum(.value > test_composition_above_logit_fold_change, na.rm = TRUE), + smaller_zero = sum(.value < -test_composition_above_logit_fold_change, na.rm = TRUE), n = n(), - .groups = "drop" # To ungroup the output if needed - ) |> - + .groups = "drop" + ) + + draws = + draw_summaries |> + left_join(probability_stats, by = c(quo_name(.cell_group), "M", "parameter")) |> # Calculate probability non 0 mutate(pH0 = (1 - (pmax(bigger_zero, smaller_zero) / n))) |> with_groups(parameter, ~ mutate(.x, FDR = get_FDR(pH0))) |> - select(!!.cell_group, M, parameter, lower, effect, upper, pH0, FDR, any_of(c("n_eff", "R_k_hat", "rhat", "ess_bulk", "ess_tail"))) |> + select(!!.cell_group, M, parameter, lower, effect, upper, pH0, FDR, any_of(c("rhat", "ess_bulk", "ess_tail"))) |> suppressWarnings() # Setting up names separately because |> is not flexible enough diff --git a/R/smooths.R b/R/smooths.R new file mode 100644 index 00000000..75de99a9 --- /dev/null +++ b/R/smooths.R @@ -0,0 +1,520 @@ +# Smooth-term utilities for sccomp formulas (brms-style splines) +# +# This file implements the R-side machinery that turns smooth terms +# (`s()`, `t2()`) in a sccomp formula into design-matrix columns that the +# existing Stan model can consume without modification. +# +# The decomposition is the standard mgcv / brms one: +# * `Xf` — unpenalised null-space columns of the smooth (linear part). +# These are appended to the fixed-effect design matrix `X` and +# estimated as ordinary `beta` coefficients. +# * `Xr` — penalised "wiggly" basis columns. These occupy a sccomp +# random-effect slot with `n_factors = 1` and `n_groups = ncol(Xr)`, +# so the slot's per-cell-group SD plays the role of brms's `sds_*`. +# +# Built via: +# 1. `mgcv::smoothCon(s(...), data, absorb.cons = TRUE, +# diagonal.penalty = TRUE)[[1]]` — basis + penalty. +# 2. `mgcv::smooth2random(sm, vnames = "", type = 2)` — re-parameterise +# into `Xf` + standardised `rand$Xr` (so the wiggly coefficients have +# a `N(0, sds * I)` prior). +# +# At prediction time, `predict_smooth_at_newdata()` evaluates the basis at +# new covariate values via `mgcv::PredictMat()` and re-applies the same +# `trans.D` (diagonal rescaling) and `trans.U` (orthonormal rotation) +# captured at fit time, so the new columns line up with the columns the +# model was trained on. + +#' Detect smooth specials in a formula +#' +#' @param fm A one-sided formula. +#' @return TRUE if `fm` contains any `s(...)` or `t2(...)` term. +#' @keywords internal +#' @noRd +has_smooth_terms <- function(fm) { + if (is.null(fm)) return(FALSE) + trm <- stats::terms(fm, specials = c("s", "t2")) + length(unlist(attr(trm, "specials"))) > 0 +} + + +#' Strip smooth specials from a formula +#' +#' Returns a new formula with `s()` / `t2()` term labels removed. All other +#' terms, including random-effect clauses (`(... | g)`), are preserved. The +#' intercept setting (`+ 0` / `- 1`) is preserved. If nothing remains, returns +#' `~ 1` (or `~ 0`). +#' +#' @param fm A one-sided formula. +#' @return A one-sided formula without smooth term labels. +#' @keywords internal +#' @noRd +strip_smooth_terms <- function(fm) { + if (is.null(fm)) return(fm) + trm <- stats::terms(fm, specials = c("s", "t2")) + smooth_idx <- unlist(attr(trm, "specials")) + + vars <- attr(trm, "variables") + smooth_labels <- if (length(smooth_idx) > 0) { + vapply( + smooth_idx + 1L, + function(i) deparse(vars[[i]], width.cutoff = 500L), + character(1) + ) + } else { + character(0) + } + + all_labels <- attr(trm, "term.labels") + keep_labels <- setdiff(all_labels, smooth_labels) + keep_labels <- ifelse(grepl("\\|", keep_labels), paste0("(", keep_labels, ")"), keep_labels) + + has_intercept <- attr(trm, "intercept") == 1L + + if (length(keep_labels) == 0) { + return(if (has_intercept) ~ 1 else ~ 0) + } + + rhs <- paste(keep_labels, collapse = " + ") + if (!has_intercept) rhs <- paste0(rhs, " + 0") + stats::as.formula(paste("~", rhs), env = environment(fm)) +} + + +#' Strip random-effect clauses from a formula +#' +#' Returns a formula with `(... | g)` term labels removed, preserving all +#' ordinary fixed-effect terms and the intercept setting. This is used just +#' before calling `model.matrix()`, which cannot evaluate sccomp/brms-style +#' random-effect clauses. +#' +#' @param fm A one-sided formula. +#' @return A one-sided formula without random-effect term labels. +#' @keywords internal +#' @noRd +strip_random_effect_terms <- function(fm) { + if (is.null(fm)) return(fm) + trm <- stats::terms(fm) + + keep_labels <- attr(trm, "term.labels") + keep_labels <- keep_labels[!grepl("\\|", keep_labels)] + + has_intercept <- attr(trm, "intercept") == 1L + + if (length(keep_labels) == 0) { + return(if (has_intercept) ~ 1 else ~ 0) + } + + rhs <- paste(keep_labels, collapse = " + ") + if (!has_intercept) rhs <- paste0(rhs, " + 0") + stats::as.formula(paste("~", rhs), env = environment(fm)) +} + + +#' Parse smooth terms from a sccomp formula +#' +#' For each `s(...)` / `t2(...)` term, builds an mgcv smooth via +#' `smoothCon(..., absorb.cons = TRUE, diagonal.penalty = TRUE)` and the +#' brms-style re-parameterisation via `smooth2random(..., type = 2)`. Single- +#' penalty smooths (e.g. plain `s(x)`) yield one penalised block; multi- +#' penalty smooths (e.g. `t2(x, y)`, `s(x, z, bs = "fs")`) yield 2+ blocks, +#' one per penalty. Each block is mapped to its own sccomp random-effect +#' slot, with its own per-cell-group smoothing SD — matching the way brms +#' generates one `sds_*` parameter per penalty. +#' +#' @param fm A one-sided formula, possibly containing `s()` / `t2()` terms. +#' @param data A data frame containing every variable referenced inside the +#' smooth terms (in the same row order as the design matrices that will be +#' built downstream). +#' @return A list with components: +#' * `parametric_formula` — `fm` with smooth specials stripped. +#' * `smooth_labels` — character vector of original smooth labels +#' (one per `s()` / `t2()` term). +#' * `Xf_list` — list of `N × Ks_k` matrices (unpenalised cols), +#' one per smooth term. +#' * `Xr_list` — **flat** list of `N × k_b` matrices, one per +#' penalty block. A single-penalty smooth contributes 1 entry; a multi- +#' penalty smooth contributes 2+. Each consumes one RE slot. +#' * `Xr_to_smooth` — integer vector parallel to `Xr_list` mapping +#' each block to its parent smooth index. +#' * `Xr_slot_labels` — character vector parallel to `Xr_list` giving +#' each block its sccomp RE-slot label: just the smooth's label for +#' single-penalty smooths, or `