R/plotting.R

In alxsimon/rfss: R package to deal with Fisher Stepping Stone output

#' Combined plot of mean fitness and allele frequencies for each deme
#'
#' @param res A dataset produced by \code{\link{load_h5}}.
#' @param gen The generation at which you whish to plot (default 1).
#' @param message Option to remove the message indicating the closest real
#' generation available that has been plotted (default TRUE).
#'
#' @return A plot generated by \code{\link{gridExtra::grid.arrange}}
#'
#' @examples
#' data(result_sim)
#' plot_timestep(result_sim, 11)
#' plot_timestep(result_sim, 10) # produces the same plot
#'
#' @import ggplot2
#' @export

plot_timestep <- function(res, gen = 1, message = T) {
  # find index of generation or closest one
  if (gen %in% res$savedGen) {
    igen <- which(res$savedGen == gen)
  } else {
    igen <- which.min(abs(res$savedGen - gen))
    gen <- res$savedGen[igen]
    if (message)
      cat("Plotting the closest generation available: ", gen)
  }

  allele_freq <- reshape2::melt(res$alleleFreq[,,igen],
                      varnames = c("deme","locus"), value.name = "freq")
  allele_freq$locus <- as.factor(allele_freq$locus)

  cols <- colorspace::diverge_hcl(dim(res$alleleFreq)[2])

  f <- ggplot(allele_freq, aes(x = deme, y = freq, colour = locus)) +
    geom_line() +
    scale_colour_manual(values = cols) +
    ylim(c(0,1)) +
    labs(x = "demes", y = "allele frequency") +
    theme_classic() +
    theme(legend.position = "none")


  w <- qplot(seq(1, dim(res$w)[1]), res$w[,igen], geom = "line") +
    ylim(c(0, 1)) +
    labs(x = "demes", y = "mean w") +
    theme_classic()

  gridExtra::grid.arrange(w, f, nrow = 4, ncol = 1,
                          layout_matrix = matrix(c(1,2,2,2)),
                          top = paste("generation", gen))
}

#' Produce an animation combining all generations
#'
#' @description Produce an animation combining all generations as an html file.
#'
#' @param res A dataset produced by \code{\link{load_h5}}.
#'
#' @return html file, images/ folder, css/ folder and js/ folder.
#'
#' @examples
#' data(result_sim)
#' plot_animation(result_sim)
#'
#' @export

plot_animation <- function(res) {
  oopt <- animation::ani.options(interval = 0.1)

  anim_fun <- function() {
    lapply(res$savedGen, function(g) {
      plot_timestep(res, gen = g)
      animation::ani.pause()
    })
  }

  outdir <- paste0("anim_",
                   "d", res$parameters$d, "_",
                   "N", res$parameters$N, "_",
                   "mig", res$parameters$mig, "_",
                   "b", res$parameters$b, "_",
                   "n", res$parameters$n, "_",
                   "m", res$parameters$m, "_",
                   "sig", res$parameters$sig, "_",
                   "diff", res$parameters$diff, "_",
		   "a", res$parameters$a, "_",
                   "Q", res$parameters$Q, "_",
                   "U", res$parameters$U, "_",
                   "nbS", res$parameters$nbS, "_",
                   "L", res$parameters$L, "_",
                   "Ts", res$parameters$T1, "-",
                   res$parameters$T2, "-",
                   res$parameters$T3, "-",
                   res$parameters$no)
  dir.create(outdir, showWarnings = F)
  saved_wd <- getwd()
  setwd(outdir)

  animation::saveHTML(anim_fun(), autoplay = FALSE, loop = FALSE,
                      verbose = FALSE, htmlfile = "animation.html",
                      single.opts = "'controls': ['first', 'previous', 'play',
                      'next', 'last', 'loop', 'speed'], 'delayMin': 0")

  setwd(saved_wd)
}