Add option to visualizes samplings with plot_model()

bodkan · Feb 21, 2024 · d72ac59 · d72ac59
1 parent 8707b9c
commit d72ac59
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 6 deletions.
diff --git a/R/visualization.R b/R/visualization.R
@@ -265,6 +265,7 @@ sort_splits <- function(model) {
 #'   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 samples Sampling schedule to be visualized over the model
 #' @param ... Optional argument which will be passed to \code{ggsave}
 #'
 #' @return A ggplot2 object with the visualized slendr model
@@ -284,7 +285,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, file = NULL, ...) {
+                       order = NULL, file = NULL, samples = NULL, ...) {
   populations <- model$populations
 
   log10_ydelta <- 0.001
@@ -487,11 +488,12 @@ plot_model <- function(model, sizes = TRUE, proportions = FALSE, gene_flow = TRU
   # labels with population names
   # except for outgroups and populations with two or more daughter populations, all
   # population labels will be plotted at the bottom of the figure
-  end_labels <- model$splits[
-    model$splits$parent != "__pop_is_ancestor" &
-    vapply(model$splits$pop, function(x) sum(x == model$splits$parent), integer(1)) < 2
-  , ]$pop
-  centers[centers$pop %in% end_labels, ]$time[end_labels] <- end_times[end_labels] + log10_ydelta
+  # (not currently done because of issues with overplotting with sampling labels)
+  # end_labels <- model$splits[
+  #   model$splits$parent != "__pop_is_ancestor" &
+  #   vapply(model$splits$pop, function(x) sum(x == model$splits$parent), integer(1)) < 2
+  # , ]$pop
+  # centers[centers$pop %in% end_labels, ]$time[end_labels] <- end_times[end_labels] + log10_ydelta
   p <- p + geom_label(data = centers, size = 3,
                       aes(label = pop, x = center, y = time, fill = pop), fontface = "bold")
 
@@ -521,6 +523,14 @@ plot_model <- function(model, sizes = TRUE, proportions = FALSE, gene_flow = TRU
 
   p <- p + scale_y_continuous(trans = trans)
 
+  # if specified, overlay sampling points over the model
+  if (!is.null(samples)) {
+    sampling_points <- dplyr::select(centers, pop, center) %>%
+      dplyr::inner_join(samples, by = "pop")
+    p <- p + geom_label(data = sampling_points, aes(label = n, x = center, y = time),
+                   fontface = "bold", alpha = 0.5)
+  }
+
   if (!is.null(file))
     ggplot2::ggsave(file, plot = p, ...)
   else

diff --git a/vignettes/vignette-04-nonspatial-models.Rmd b/vignettes/vignette-04-nonspatial-models.Rmd
@@ -75,6 +75,22 @@ Using the `plot_map()` function doesn't make sense, as there are no spatial maps
 plot_model(model)
 ```
 
+Let's say we also want to schedule specific sampling events at specific times (which record only specified individuals in a tree sequence). We can use `schedule_sampling()` to do just that:
+
+```{r}
+samples <- schedule_sampling(
+  model,
+  times = c(0, 5000, 12000, 20000, 35000, 39000, 43000),
+  list(eur, 3), list(ehg, 1), list(yam, 1), list(ana, 3), list(ooa, 1), list(afr, 1)
+)
+```
+
+Because it's not just the model itself that's useful to visually verify (which is the main purpose of `plot_model()`) but also the sampling scheme, _slendr_ makes it possible to overlay the numbers of individuals scheduled for tree-sequence recording from each lineage at each timepoint. Here we drop the gene-flow arrows to make the figure a bit less cluttered:
+
+```{r, non-spatial-graph_sampling, fig.width=7, fig.height=4}
+plot_model(model, samples = samples, gene_flow = FALSE)
+```
+
 Even the final step---execution of the model in SLiM---is the same, using the built-in `slim()` function:
 
 ```{r, eval = FALSE}