Add option to save a plot_model() figure to a file

bodkan · Jul 21, 2023 · e600784 · e600784
1 parent 5b677bd
commit e600784
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 4 deletions.
diff --git a/R/visualization.R b/R/visualization.R
@@ -264,6 +264,8 @@ sort_splits <- function(model) {
 #'   vector of population names. If \code{NULL} (the default), the default plotting
 #'   algorithm will be used, ordering populations from the most ancestral to the
 #'   most recent using an in-order tree traversal.
+#' @param file Output file for a figure saved via \code{ggsave}
+#' @param ... Optional argument which will be passed to \code{ggsave}
 #'
 #' @return A ggplot2 object with the visualized slendr model
 #'
@@ -282,7 +284,7 @@ sort_splits <- function(model) {
 #'   labs geom_segment arrow
 #' @export
 plot_model <- function(model, sizes = TRUE, proportions = FALSE, gene_flow = TRUE, log = FALSE,
-                       order = NULL) {
+                       order = NULL, file = NULL, ...) {
   populations <- model$populations
 
   log10_ydelta <- 0.001
@@ -519,7 +521,10 @@ plot_model <- function(model, sizes = TRUE, proportions = FALSE, gene_flow = TRU
 
   p <- p + scale_y_continuous(trans = trans)
 
-  p
+  if (!is.null(file))
+    ggplot2::ggsave(file, plot = p, ...)
+  else
+    p
 }
 
 #' Animate the simulated population dynamics

diff --git a/docs/reference/plot_model.html b/docs/reference/plot_model.html
diff --git a/man/plot_model.Rd b/man/plot_model.Rd