R/plotting.R

In XSLmodels: Run and Optimize Cross-Situational Word Learning Models

#' Create Animated Plot of Co-Occurrence Matrix from Training Trials
#'
#' This function takes a sequence of training trials, each consisting of words
#' and objects, and creates an animated plot of the word-object co-occurrence
#' matrix. The plot evolves over each trial, visually representing the
#' accumulation of co-occurrences. The animation is saved to a file if a
#' filename is provided.
#'
#' @param train A list representing the training data, where each element is a
#'   trial containing lists of words (`words`) and objects (`objects`).
#' @param filename A string specifying the filename where the animation should
#'   be saved. If NULL, no file will be saved.
#'
#' @return A long-format data frame where each row corresponds to a trial and
#'   includes columns for trial number, word, object, and the count of
#'   co-occurrences up to that trial.
#' @export
#'
#' @examples
#' \donttest{
#' plot_training_trials(xsl_datasets[[1]]$train)
#' }
plot_training_trials <- function(train, filename = NULL) {
  if (!requireNamespace("viridis", quietly = TRUE) ||
        !requireNamespace("gganimate", quietly = TRUE)) {
    stop("Packages \"viridis\" and \"gganimate\" are required for ",
         "plot_training_trials(). Install them with ",
         "install.packages(c(\"viridis\", \"gganimate\")).", call. = FALSE)
  }

  words <- sort(unique(unlist(train$words)))
  objects <- sort(unique(unlist(train$objects)))

  matrix_to_long <- function(m, t) {
    m |>
      as_tibble() |>
      mutate(word = rownames(m), .before = everything()) |>
      pivot_longer(-.data$word, names_to = "object", values_to = "coocs") |>
      mutate(trial = t)
  }

  m <- matrix(0, nrow = length(words), ncol = length(objects),
              dimnames = list(words, objects))
  df <- matrix_to_long(m, 0)

  for (t in seq_along(train$words)) {
    tr_w <- unlist(train$words[[t]])
    tr_o <- unlist(train$objects[[t]])
    tr_w <- tr_w[!is.na(tr_w)]
    tr_o <- tr_o[!is.na(tr_o)]
    if (length(tr_w) == 0 || length(tr_o) == 0) {
      df <- bind_rows(df, matrix_to_long(m, t))
      next
    }
    tr_w <- as.character(tr_w)  # index by label, not position (#14)
    tr_o <- as.character(tr_o)
    m[tr_w, tr_o] <- m[tr_w, tr_o] + 1
    df_t <- matrix_to_long(m, t)
    df <- bind_rows(df, df_t)
  }

  p <- ggplot(df, aes(x = .data$object, y = .data$word, fill = .data$coocs)) +
    geom_tile(color = "white", linewidth = 0.1) +
    viridis::scale_fill_viridis(name = "Count", option = "magma") +
    coord_equal() +
    labs(x = "Object", y = "Word") +
    gganimate::transition_states(.data$trial)

  # gganimate's renderer falls back to writing raw PNG frames into the
  # current working directory when none of gifski/magick/av are installed;
  # render inside a scratch dir so that can never happen regardless of which
  # backend ends up in use.
  old_wd <- setwd(tempdir())
  on.exit(setwd(old_wd), add = TRUE)

  if (is.null(filename)) {
    print(p)
  } else {
    gganimate::animate(p, format = "mp4",
                       renderer = gganimate::ffmpeg_renderer(format = "mp4"))
    setwd(old_wd)
    gganimate::anim_save(filename)
  }

  as_tibble(df)
}