From 84df5e8b26b6c96de11ab0d4bc448fd7a64b7321 Mon Sep 17 00:00:00 2001
From: Andrea Manica <am315@cam.ac.uk>
Date: Thu, 12 Dec 2024 17:08:29 +0000
Subject: [PATCH] fix plotting order

---
 R/q_matrix.R                       | 24 +++++++++++++++++++-----
 man/autoplot_q_matrix.Rd           | 12 +++++++++++-
 tests/testthat/test_gt_admixture.R |  9 ++++++---
 3 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/R/q_matrix.R b/R/q_matrix.R
index fe3856af..4cecaf93 100644
--- a/R/q_matrix.R
+++ b/R/q_matrix.R
@@ -248,13 +248,21 @@ augment.q_matrix <- function(x, data = NULL, ...) {
 #' generate the Q matrix
 #' @param annotate_group Boolean determining whether to annotate the plot with the
 #' group information
+#' @param reorder_within_groups Boolean determining whether to reorder the individuals within each group based
+#' on their ancestry proportion (note that this is not advised if you are making multiple plots, as you would get
+#' a different order for each plot!). If TRUE, `annotate_group` must also be TRUE.
 #' @param ... not currently used.
 #' @returns a barplot of individuals, coloured by ancestry proportion
 #' @name autoplot_q_matrix
 #' @export
-autoplot.q_matrix <- function(object, data = NULL, annotate_group = TRUE, ...){
+autoplot.q_matrix <- function(object, data = NULL, annotate_group = TRUE, reorder_within_groups = FALSE, ...){
 
   rlang::check_dots_empty()
+  # test that if reorder_within_groups is TRUE, annotate_group should also be TRUE
+  if (reorder_within_groups & !annotate_group){
+    stop("If reorder_within_groups is TRUE, annotate_group should also be TRUE")
+  }
+
   K <- ncol(object)
   # create dataset if we don't have a gen_tibble
   if (is.null(data)) {
@@ -269,17 +277,23 @@ autoplot.q_matrix <- function(object, data = NULL, annotate_group = TRUE, ...){
     if ("group" %in% names(attributes(object))){
       q_tbl$group <- rep(attr(object, "group"), each=nrow(q_tbl)/length(attr(object, "group")))
     }
-    browser()
   } else { # if we have the info from the gen_tibble
     q_tbl <- tidy(object, data)
   }
+  # if we have a grouping variable and we plan to use it, then reorder by it
+  q_tbl$id <- 1:nrow(q_tbl)
+  if (("group" %in% names(q_tbl)) && annotate_group){
+    q_tbl <- q_tbl %>%
+      dplyr::arrange(.data$group, .data$id)
+  }
+  # now reset the id to the new order
+  q_tbl$id <- 1:nrow(q_tbl)
   q_tbl <- q_tbl %>% tidyr::pivot_longer(cols = dplyr::starts_with(".Q"),
                                          names_to = "q", values_to = "percentage") %>%
     dplyr::mutate(percentage = as.numeric(.data$percentage))
 
-  # resort data if we have a grouping variable and we plan to use it
-  if (("group" %in% names(q_tbl))&&annotate_group){
-
+  # if we reorder within group
+  if (("group" %in% names(q_tbl)) && reorder_within_groups){
     q_tbl <- q_tbl %>%
       dplyr::group_by(.data$group, .data$id) %>%
       #dplyr::arrange(.data$group, .data$id) %>%
diff --git a/man/autoplot_q_matrix.Rd b/man/autoplot_q_matrix.Rd
index 3533d483..407dde61 100644
--- a/man/autoplot_q_matrix.Rd
+++ b/man/autoplot_q_matrix.Rd
@@ -5,7 +5,13 @@
 \alias{autoplot.q_matrix}
 \title{Autoplots for \code{q_matrix} objects}
 \usage{
-\method{autoplot}{q_matrix}(object, data = NULL, annotate_group = TRUE, ...)
+\method{autoplot}{q_matrix}(
+  object,
+  data = NULL,
+  annotate_group = TRUE,
+  reorder_within_groups = FALSE,
+  ...
+)
 }
 \arguments{
 \item{object}{A Q matrix object (as returned by \code{\link[=q_matrix]{q_matrix()}}).}
@@ -16,6 +22,10 @@ generate the Q matrix}
 \item{annotate_group}{Boolean determining whether to annotate the plot with the
 group information}
 
+\item{reorder_within_groups}{Boolean determining whether to reorder the individuals within each group based
+on their ancestry proportion (note that this is not advised if you are making multiple plots, as you would get
+a different order for each plot!). If TRUE, \code{annotate_group} must also be TRUE.}
+
 \item{...}{not currently used.}
 }
 \value{
diff --git a/tests/testthat/test_gt_admixture.R b/tests/testthat/test_gt_admixture.R
index 5921cb62..3bff9d62 100644
--- a/tests/testthat/test_gt_admixture.R
+++ b/tests/testthat/test_gt_admixture.R
@@ -63,7 +63,10 @@ test_that("run admixture as multiple runs", {
 
   # test the reorder of q matrices
   anole_adm_cv_reorder <- gt_admix_reorder_q(anole_adm_cv)
-  # TODO check how ordering is done automatically when plotting
-  autoplot(anole_adm_cv, type = "barplot", k=3,run = 1, annotate_group=FALSE)
-  autoplot(anole_adm_cv_reorder, type = "barplot", k=3,run = 1)
+  # check plot ordering
+  unord_plot <- autoplot(anole_adm_cv, type = "barplot", k=3,run = 1, annotate_group=FALSE)
+  ord_plot <- autoplot(anole_adm_cv_reorder, type = "barplot", k=3,run = 1, annotate_group=TRUE)
+  reord_plot <- autoplot(anole_adm_cv_reorder, type = "barplot", k=3,run = 1, annotate_group=TRUE)
+  expect_identical (ord_plot, reord_plot)
+  expect_false (identical (unord_plot, ord_plot))
 })