Dispersion estimation

DESCRIPTION

@@ -83,6 +83,7 @@ Suggests:
     pbapply,
     pbmcapply,
     lme4,
+    reformulas,
     glmmTMB,
     MASS,
     pkgconfig,
@@ -97,7 +98,7 @@ Biarch: true
 biocViews: AssayDomain, Infrastructure, RNASeq, DifferentialExpression, GeneExpression, Normalization, Clustering, QualityControl, Sequencing, Transcription, Transcriptomics
 Encoding: UTF-8
 LazyData: true
-RoxygenNote: 7.3.3
 LazyDataCompression: xz
 URL: https://github.com/stemangiola/tidybulk
 BugReports: https://github.com/stemangiola/tidybulk/issues
+Config/roxygen2/version: 8.0.0

R/estimate_dispersion.R

@@ -0,0 +1,69 @@
+# Internal. edgeR tagwise (or trended if n >= 1000) written to rowData.
+# formula_abundance must be a fixed-effects formula; the user chooses it.
+
+estimate_dispersion <- function(.data, formula_abundance, abundance = "counts") {
+  check_and_install_packages("edgeR")
+  check_formula(formula_abundance)
+  if (grepl("|", paste(deparse(formula_abundance), collapse = ""), fixed = TRUE)) {
+    stop(
+      "tidybulk says: estimate_dispersion() uses edgeR, which has no random effects. ",
+      "Pass a fixed-effects formula (e.g. ~ treatment + donor), not ",
+      paste(deparse(formula_abundance), collapse = ""),
+      ".",
+      call. = FALSE
+    )
+  }
+
+  n_sample <- ncol(.data)
+  design <- stats::model.matrix(
+    formula_abundance,
+    data = droplevels(as.data.frame(colData(.data)))
+  )
+  if (n_sample <= ncol(design)) {
+    stop(
+      "tidybulk says: The dispersion design has ", ncol(design),
+      " coefficients for ", n_sample,
+      " samples, leaving no residual degrees of freedom.",
+      call. = FALSE
+    )
+  }
+
+  formula_txt <- paste(deparse(formula_abundance), collapse = "")
+  if (n_sample < 1000L) {
+    message(
+      "tidybulk says: calculating tagwise (shrinked) and trended dispersion ",
+      "with edgeR::estimateDisp() using ", formula_txt, "."
+    )
+    counts <- assay(.data, abundance)
+    fit <- edgeR::estimateDisp(counts, design = design)
+    disp <- fit$tagwise.dispersion
+    trended <- fit$trended.dispersion
+    d_eff <- (n_sample - ncol(design)) + fit$prior.df
+  } else {
+    sampled <- sample(seq_len(n_sample), size = min(n_sample, 2000L))
+    se_sub <- .data[, sampled, drop = FALSE]
+    design <- stats::model.matrix(
+      formula_abundance,
+      data = droplevels(as.data.frame(colData(se_sub)))
+    )
+    message(
+      "tidybulk says: n >= 1000; calculating trended dispersion ",
+      "with edgeR::estimateTrendedDisp()."
+    )
+    fit <- edgeR::estimateTrendedDisp(
+      assay(se_sub, abundance),
+      design = design,
+      subset = 1000,
+      rowsum.filter = 10
+    )
+    disp <- trended <- fit
+    d_eff <- ncol(se_sub) - ncol(design)
+  }
+
+  rowData(.data)[["dispersion_shrinked"]] <- as.numeric(disp)
+  rowData(.data)[["dispersion_trended"]] <- as.numeric(trended)
+  rowData(.data)[["dispersion_degrees_freedom"]] <-
+    rep_len(as.numeric(d_eff), nrow(.data))
+  .data <- attach_to_metadata(.data, fit, "estimateDisp")
+  memorise_methods_used(.data, "edger")
+}

R/glmmSeq.R

@@ -71,7 +71,7 @@ organiseStats = function (resultList, test.stat)
 
 hyp_matrix = function (fullFormula, metadata, LHS)
 {
-  reduced2 <- lme4::nobars(fullFormula)
+  reduced2 <- reformulas::nobars(fullFormula)
   fac <- attr(terms(reduced2), "factors")
   data2 <- metadata
   data2[, LHS] <- rep(0, nrow(data2))
@@ -256,10 +256,23 @@ lmer_to_confidence_intervals_random_effects = function(fit){
 glmmTMBcore = function (geneList, fullFormula, reduced, data, family, control,
           offset, modelData, designMatrix, hyp.matrix, ...)
 {
-  data[, "count"] <- geneList
-  fit <- try(suppressMessages(suppressWarnings(glmmTMB::glmmTMB(fullFormula,
-                                                       data, family, control = control, offset = offset, ...))),
-             silent = TRUE)
+  data[, "count"] <- geneList$y
+  # Plugin edgeR phi so glmmTMB does not estimate it per gene (same idea as
+  # glmer with MASS::negative.binomial(theta = 1/disp) vs glmer.nb).
+  # nbinom2 uses Var = μ + μ²/θ, so θ = 1/φ; betadisp is log(θ).
+  # map$betadisp = factor(NA) fixes that parameter at start.
+  disp <- geneList$dispersion
+  extra <- list(...)
+  if (is.finite(disp) && disp > 0) {
+    extra$start$betadisp <- log(1 / disp)
+    extra$map$betadisp <- factor(NA)
+  }
+  fit <- try(suppressMessages(suppressWarnings(
+    do.call(
+      glmmTMB::glmmTMB,
+      c(list(fullFormula, data, family, control = control, offset = offset), extra)
+    )
+  )), silent = TRUE)
   if (!inherits(fit, "try-error")) {
     singular <- conv <- NA
     stdErr <- suppressWarnings(coef(summary(fit))$cond[, 2])
@@ -272,17 +285,22 @@ glmmTMBcore = function (geneList, fullFormula, reduced, data, family, control,
                   message = "vcov. failed to calculate", tryErrors = fit[1]))
 
     fixedEffects <- glmmTMB::fixef(fit)$cond
-    disp <- glmmTMB::sigma(fit)
+    if (!is.finite(disp) || disp <= 0) {
+      disp <- glmmTMB::sigma(fit)
+    }
     msg <- fit$fit$message
     stats <- setNames(c(disp, AIC(fit), as.numeric(logLik(fit))),
                       c("Dispersion", "AIC", "logLik"))
     if (is.null(reduced)) {
       test <- lmer_wald(fixedEffects, hyp.matrix, vcov.)
     }
     else {
-      fit2 <- try(suppressMessages(suppressWarnings(glmmTMB::glmmTMB(reduced,
-                                                            data, family, control = control, offset = offset,
-                                                            ...))), silent = TRUE)
+      fit2 <- try(suppressMessages(suppressWarnings(
+        do.call(
+          glmmTMB::glmmTMB,
+          c(list(reduced, data, family, control = control, offset = offset), extra)
+        )
+      )), silent = TRUE)
       if (!inherits(fit2, "try-error")) {
         lrt <- anova(fit, fit2)
         test <- list(chisq = setNames(lrt$Chisq[2], "LRT"),
@@ -327,9 +345,15 @@ glmerCore = function (geneList, fullFormula, reduced, data, control, offset,
 {
   data[, "count"] <- geneList$y
   disp <- geneList$dispersion
-  fit <- try(suppressMessages(suppressWarnings(lme4::glmer(fullFormula,
-                                                           data = data, control = control, offset = offset, family = MASS::negative.binomial(theta = 1/disp),
-                                                           ...))), silent = TRUE)
+  plugin <- is.finite(disp) && disp > 0
+  fit <- try(suppressMessages(suppressWarnings(
+    if (plugin) {
+      lme4::glmer(fullFormula, data = data, control = control, offset = offset,
+                  family = MASS::negative.binomial(theta = 1 / disp), ...)
+    } else {
+      lme4::glmer.nb(fullFormula, data = data, control = control, offset = offset, ...)
+    }
+  )), silent = TRUE)
 
 
   # If errors return
@@ -347,16 +371,23 @@ glmerCore = function (geneList, fullFormula, reduced, data, control, offset,
   conv <- length(slot(fit, "optinfo")$conv$lme4$messages)
   vcov. <- suppressWarnings(as.matrix(vcov(fit, complete = FALSE)))
   fixedEffects <- lme4::fixef(fit)
+  if (!plugin) {
+    disp <- 1 / lme4::getME(fit, "glmer.nb.theta")
+  }
   stats <- setNames(c(disp, AIC(fit), as.numeric(logLik(fit))),
                     c("Dispersion", "AIC", "logLik"))
   if (is.null(reduced)) {
     test <- lmer_wald(fixedEffects, hyp.matrix, vcov.)
   }
   else {
-    fit2 <- try(suppressMessages(suppressWarnings(lme4::glmer(reduced,
-                                                              data = data, control = control, offset = offset,
-                                                              family = MASS::negative.binomial(theta = 1/disp),
-                                                              ...))), silent = TRUE)
+    fit2 <- try(suppressMessages(suppressWarnings(
+      if (plugin) {
+        lme4::glmer(reduced, data = data, control = control, offset = offset,
+                    family = MASS::negative.binomial(theta = 1 / disp), ...)
+      } else {
+        lme4::glmer.nb(reduced, data = data, control = control, offset = offset, ...)
+      }
+    )), silent = TRUE)
     if (!inherits(fit2, "try-error")) {
       lrt <- anova(fit, fit2)
       test <- list(chisq = setNames(lrt$Chisq[2], "LRT"),
@@ -459,13 +490,14 @@ glmmSeq = function (modelFormula, countdata, metadata, id = NULL, dispersion = N
                     returnList = FALSE, progress = FALSE, max_rows_for_matrix_multiplication = Inf, avoid_forking = FALSE, ...)
 {
   glmmcall <- match.call(expand.dots = TRUE)
-  method <- match.arg(method)
-  if (is.null(control)) {
-    control <- switch(method, lme4 = lme4::glmerControl(optimizer = "bobyqa"),
-                      glmmTMB = glmmTMB::glmmTMBControl())
+  method <- method[1]
+  if (identical(tolower(method), "glmmtmb")) {
+    method <- "glmmTMB"
   }
+  method <- match.arg(method, c("lme4", "glmmTMB"))
   countdata <- as.matrix(countdata)
-  if (length(lme4::findbars(modelFormula)) == 0) {
+  check_and_install_packages("reformulas")
+  if (length(reformulas::findbars(modelFormula)) == 0) {
     stop("No random effects terms specified in formula")
   }
   if (ncol(countdata) != nrow(metadata)) {
@@ -481,11 +513,11 @@ glmmSeq = function (modelFormula, countdata, metadata, id = NULL, dispersion = N
   if (zeroCount > 0)
     countdata[countdata == 0] <- zeroCount
   fullFormula <- update.formula(modelFormula, count ~ ., simplify = FALSE)
-  subFormula <- lme4::subbars(modelFormula)
+  subFormula <- reformulas::subbars(modelFormula)
   variables <- rownames(attr(terms(subFormula), "factors"))
   subsetMetadata <- metadata[, variables, drop=FALSE]
   if (is.null(id)) {
-    fb <- lme4::findbars(modelFormula)
+    fb <- reformulas::findbars(modelFormula)
     id <- sub(".*[|]", "", fb)
     id <- gsub(" ", "", id)
   }
@@ -504,7 +536,7 @@ glmmSeq = function (modelFormula, countdata, metadata, id = NULL, dispersion = N
   if (verbose)
     message(paste0("\nn = ", length(ids), " samples, ", length(unique(ids)),
                " individuals\n"))
-  FEformula <- lme4::nobars(modelFormula)
+  FEformula <- reformulas::nobars(modelFormula)
   if (is.null(modelData)) {
     reducedVars <- rownames(attr(terms(FEformula), "factors"))
     varLevels <- lapply(reducedVars, function(x) {
@@ -531,10 +563,10 @@ glmmSeq = function (modelFormula, countdata, metadata, id = NULL, dispersion = N
     hyp.matrix <- hyp_matrix(fullFormula, metadata, "count")
   }
   else {
-    if (length(lme4::findbars(reduced)) == 0) {
+    if (length(reformulas::findbars(reduced)) == 0) {
       stop("No random effects terms specified in reduced formula")
     }
-    subReduced <- lme4::subbars(reduced)
+    subReduced <- reformulas::subbars(reduced)
     redvars <- rownames(attr(terms(subReduced), "factors"))
     if (any(!redvars %in% variables)) {
       stop("Extra terms in reduced formula not found full formula")
@@ -556,15 +588,21 @@ glmmSeq = function (modelFormula, countdata, metadata, id = NULL, dispersion = N
 
 
   if (method == "lme4") {
+    check_and_install_packages("lme4")
+    if (is.null(control)) {
+      control <- lme4::glmerControl(optimizer = "bobyqa")
+    }
 
     # FASTER - Stefano
     control$calc.derivs = FALSE
 
-    if (!all(rownames(countdata) %in% names(dispersion))) {
-      stop("Some dispersion values are missing")
-    }
+    plugin <- is.numeric(dispersion) &&
+      all(rownames(countdata) %in% names(dispersion))
     fullList <- lapply(rownames(countdata), function(i) {
-      list(y = countdata[i, ], dispersion = dispersion[i])
+      list(
+        y = countdata[i, ],
+        dispersion = if (plugin) unname(dispersion[i]) else NA_real_
+      )
     })
     if(cores == 1){
       resultList <- lapply(fullList, function(geneList) {
@@ -654,12 +692,21 @@ glmmSeq = function (modelFormula, countdata, metadata, id = NULL, dispersion = N
     }
   }
   else {
+    check_and_install_packages("glmmTMB")
+    if (is.null(control)) {
+      control <- glmmTMB::glmmTMBControl()
+    }
 
     # FASTER - Stefano
     control$profile=TRUE
 
+    plugin <- is.numeric(dispersion) &&
+      all(rownames(countdata) %in% names(dispersion))
     fullList <- lapply(rownames(countdata), function(i) {
-      countdata[i, ]
+      list(
+        y = countdata[i, ],
+        dispersion = if (plugin) unname(dispersion[i]) else NA_real_
+      )
     })
 
     if(cores == 1){