From cb96f3245ee1ef561d7a49e55516b632c7c143f8 Mon Sep 17 00:00:00 2001
From: Stefano Mangiola <mangiolastefano@gmail.com>
Date: Wed, 30 Aug 2023 00:00:46 +1000
Subject: [PATCH] add custom dispersion for tibble input

---
 R/functions.R                                 |  6 ++-
 tests/testthat/test-bulk_methods.R            | 38 +++++++++++++++++--
 .../test-bulk_methods_SummarizedExperiment.R  | 13 ++++---
 3 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/R/functions.R b/R/functions.R
index 271483d6..82f9fc2a 100755
--- a/R/functions.R
+++ b/R/functions.R
@@ -685,13 +685,15 @@ get_differential_transcript_abundance_glmmSeq <- function(.data,
                                                                omit_contrast_in_colnames = FALSE,
                                                                prefix = "",
                                                                .sample_total_read_count = NULL,
+                                                          .dispersion = NULL,
                                                           ...) {
   # Get column names
   .sample = enquo(.sample)
   .transcript = enquo(.transcript)
   .abundance = enquo(.abundance)
   .sample_total_read_count = enquo(.sample_total_read_count)
-
+  .dispersion = enquo(.dispersion)
+  
   # Check if omit_contrast_in_colnames is correctly setup
   if(omit_contrast_in_colnames & length(.contrasts) > 1){
     warning("tidybulk says: you can omit contrasts in column names only when maximum one contrast is present")
@@ -757,7 +759,7 @@ get_differential_transcript_abundance_glmmSeq <- function(.data,
   counts = counts[,rownames(metadata),drop=FALSE]
   
   if(quo_is_symbolic(.dispersion))
-    dispersion = .data |> pivot_transcript() |> select(!!feature__$symbol, !!.dispersion) |> deframe()
+    dispersion = .data |> pivot_transcript(!!.transcript) |> select(!!.transcript, !!.dispersion) |> deframe()
   else
     dispersion = setNames(edgeR::estimateDisp(counts)$tagwise.dispersion, rownames(counts))
   
diff --git a/tests/testthat/test-bulk_methods.R b/tests/testthat/test-bulk_methods.R
index 6db6d34f..417c5827 100755
--- a/tests/testthat/test-bulk_methods.R
+++ b/tests/testthat/test-bulk_methods.R
@@ -835,18 +835,22 @@ test_that("DESeq2 differential trancript abundance - no object",{
 
 test_that("differential trancript abundance - random effects",{
 
-  input_df |>
+  my_input = 
+    input_df |>
     identify_abundant(a, b, c, factor_of_interest = condition) |>
     mutate(time = time |> stringr::str_replace_all(" ", "_")) |>
-
-    filter(b %in% c("ABCB4" , "ABCB9" , "ACAP1",  "ACHE",   "ACP5" ,  "ADAM28"))|>
+    
+    filter(b %in% c("ABCB4" , "ABCB9" , "ACAP1",  "ACHE",   "ACP5" ,  "ADAM28"))
+  
+  my_input |>
     test_differential_abundance(
       ~ condition + (1 + condition | time),
       .sample = a,
       .transcript = b,
       .abundance = c,
       method = "glmmseq_lme4",
-      action="only"
+      action="only", 
+      cores = 1
     ) |>
     pull(P_condition_adjusted) |>
     head(4) |>
@@ -855,6 +859,32 @@ test_that("differential trancript abundance - random effects",{
       tolerance=1e-3
     )
 
+  # Custom dispersion
+  my_input = 
+    my_input |>
+    left_join(
+      my_input |> pivot_transcript(b) |> mutate(disp_ = 2 ), 
+      by = join_by(b, entrez, .abundant)
+    )
+   
+
+    my_input |>
+    test_differential_abundance(
+      ~ condition + (1 + condition | time),
+      .sample = a,
+      .transcript = b,
+      .abundance = c,
+      method = "glmmseq_lme4",
+      action="only", 
+      cores = 1, 
+      .dispersion = disp_
+    ) |>
+    pull(P_condition_adjusted) |>
+    head(4) |>
+    expect_equal(
+      c(0.1081176, 0.1303558, 0.1303558, 0.1693276),
+      tolerance=1e-3
+    )
 
 })
 
diff --git a/tests/testthat/test-bulk_methods_SummarizedExperiment.R b/tests/testthat/test-bulk_methods_SummarizedExperiment.R
index 453f0496..aa67e9d5 100755
--- a/tests/testthat/test-bulk_methods_SummarizedExperiment.R
+++ b/tests/testthat/test-bulk_methods_SummarizedExperiment.R
@@ -465,7 +465,8 @@ test_that("differential trancript abundance - random effects SE",{
     #mutate(time = time |> stringr::str_replace_all(" ", "_")) |>
     test_differential_abundance(
       ~ condition + (1 + condition | time),
-      method = "glmmseq_lme4"
+      method = "glmmseq_lme4", 
+      cores = 1
     )
 
  rowData(res)[,"P_condition_adjusted"] |>
@@ -480,22 +481,22 @@ test_that("differential trancript abundance - random effects SE",{
    se_mini |>
    identify_abundant(factor_of_interest = condition)
  
- rowData(se_mini)$disp_ = rnorm(nrow(se_mini))
+ rowData(se_mini)$disp_ = rep(2,nrow(se_mini))
  
  res = 
-   se_mini |>
-   identify_abundant(factor_of_interest = condition) |> 
+   se_mini[1:10,] |>
    #mutate(time = time |> stringr::str_replace_all(" ", "_")) |> 
    test_differential_abundance(
      ~ condition + (1 + condition | time),
      method = "glmmseq_lme4", 
-     dispersion = disp_
+     .dispersion = disp_, 
+     cores = 1
    ) 
  
  rowData(res)[,"P_condition_adjusted"] |> 
    head(4) |> 
    expect_equal(
-     c(0.1441371, 0.1066183, 0.1370748, NA),
+     c(0.1153254, 0.1668555, 0.1668555 ,       NA),
      tolerance=1e-3
    )