R/plots.R

In keyATM: Keyword Assisted Topic Models

#' Save a figure
#'
#' @param x the keyATM_fig object.
#' @param filename file name to create on disk.
#' @param ... other arguments passed on to the [ggplot2::ggsave()][ggplot2::ggsave] function.
#' @seealso [visualize_keywords()], [plot_alpha()], [plot_modelfit()], [plot_pi()], [plot_timetrend()], [plot_topicprop()], [by_strata_DocTopic()], [values_fig()]
#' @export
save_fig <- function(x, filename, ...) {
  UseMethod("save_fig")
}


#' Get values used to create a figure
#'
#' @param x the keyATM_fig object.
#' @seealso [save_fig()], [visualize_keywords()], [plot_alpha()], [plot_modelfit()], [plot_pi()], [plot_timetrend()], [plot_topicprop()], [by_strata_DocTopic()]
#' @export
values_fig <- function(x) {
  UseMethod("values_fig")
}

#' @noRd
#' @export
values_fig.keyATM_fig <- function(x)
{
  return(x$values)
}


#' @noRd
#' @export
save_fig.keyATM_fig <- function(x, filename, ...)
{
  ggplot2::ggsave(filename = filename, plot = x$figure, ...)
}


#' @noRd
#' @export
print.keyATM_fig <- function(x, ...)
{
  print(x$figure)
}


#' Show a diagnosis plot of alpha
#'
#' @param x the output from a keyATM model (see [keyATM()]).
#' @param start integer. The start of slice iteration. Default is \code{0}.
#' @param show_topic a vector to specify topic indexes to show. Default is \code{NULL}.
#' @param scales character. Control the scale of y-axis (the parameter in [ggplot2::facet_wrap()][ggplot2::facet_wrap]): \code{free} adjusts y-axis for parameters. Default is \code{fixed}.
#' @return keyATM_fig object
#' @import ggplot2
#' @import magrittr
#' @importFrom rlang .data
#' @seealso [save_fig()]
#' @export
plot_alpha <- function(x, start = 0, show_topic = NULL, scales = "fixed")
{
  check_arg_type(x, "keyATM_output")
  modelname <- extract_full_model_name(x)

  if (modelname %in% c("lda", "ldacov", "ldahmm")) {
    cli::cli_abort(paste0("This is not a model with keywords."))  # only plot keywords later
  }
  if (!"alpha_iter" %in% names(x$values_iter)) {
    cli::cli_abort("`alpha` is not stored. Please check the options.\nNote that the covariate model does not have `alpha`.\nPlease check our paper for details.")
  }
  if (is.null(show_topic)) {
    show_topic <- 1:x$keyword_k
  } else {
    if (!all(show_topic %in% 1:x$keyword_k))
      cli::cli_abort("Topics specified in `show_topic` are not the keyword topics.")
  }
  if (!is.numeric(start) | length(start) != 1) {
    cli::cli_abort("`start` argument is invalid.")
  }

  tnames <- as.character(show_topic)
  names(tnames) <- c(names(x$keywords_raw))[show_topic]
  temp <- x$values_iter$alpha_iter %>%
    dplyr::filter(.data$Iteration >= start) %>%
    dplyr::filter(.data$Topic %in% (!!show_topic)) %>%
    tidyr::pivot_wider(names_from = "Topic", values_from = "alpha")

  if (modelname %in% c("base", "lda")) {
    res_alpha <- temp %>% dplyr::rename(tidyselect::all_of(tnames)) %>%
      tidyr::pivot_longer(-"Iteration", names_to = "Topic", values_to = "alpha")

    p <- ggplot(res_alpha, aes(x = .data$Iteration, y = .data$alpha, group = .data$Topic)) +
          geom_line() +
          geom_point(size = 0.3) +
          facet_wrap(~ .data$Topic, ncol = 2, scales = scales) +
          xlab("Iteration") + ylab("Value") +
          ggtitle("Estimated alpha") + theme_bw() +
          theme(plot.title = element_text(hjust = 0.5))
  } else if (modelname %in% c("hmm", "ldahmm")) {
    res_alpha <- temp %>% dplyr::rename(tidyselect::all_of(tnames)) %>%
      tidyr::pivot_longer(-c("Iteration", "State"), names_to = "Topic", values_to = "alpha")
    res_alpha$State <- factor(res_alpha$State, levels = 1:max(res_alpha$State))

    p <- ggplot(res_alpha, aes(x = .data$Iteration, y = .data$alpha, group = .data$State, colour = .data$State)) +
          geom_line() +
          geom_point(size = 0.3) +
          facet_wrap(~ .data$Topic, ncol = 2, scales = scales) +
          xlab("Iteration") + ylab("Value") +
          ggtitle("Estimated alpha") + theme_bw() +
          theme(plot.title = element_text(hjust = 0.5))
  }

  p <- list(figure = p, values = res_alpha)
  class(p) <- c("keyATM_fig", class(p))
  return(p)
}


#' Show a diagnosis plot of log-likelihood and perplexity
#'
#' @param x the output from a keyATM model (see [keyATM()]).
#' @param start integer. The starting value of iteration to use in plot. Default is \code{1}.
#' @return keyATM_fig object.
#' @import ggplot2
#' @importFrom stats as.formula
#' @importFrom rlang .data
#' @seealso [save_fig()]
#' @export
plot_modelfit <- function(x, start = 1)
{
  check_arg_type(x, "keyATM_output")
  modelfit <- x$model_fit

  if (!is.numeric(start) | length(start) != 1) {
    cli::cli_abort("`start` argument is invalid.")
  }

  if (!is.null(start)) {
    modelfit <- modelfit[ modelfit$Iteration >= start, ]
  }

  modelfit <- tidyr::gather(modelfit, key = "Measures", value = "value", -"Iteration")

  p <- ggplot(data = modelfit, aes(x = .data$Iteration, y = .data$value, group = .data$Measures, color = .data$Measures)) +
     geom_line(show.legend = FALSE) +
     geom_point(size = 0.3, show.legend = FALSE) +
     facet_wrap(~ .data$Measures, ncol = 2, scales = "free") +
     xlab("Iteration") + ylab("Value")
  p <- p + ggtitle("Model Fit") + theme_bw() + theme(plot.title = element_text(hjust = 0.5))

  p <- list(figure = p, values = modelfit)
  class(p) <- c("keyATM_fig", class(p))
  return(p)
}


#' Show a diagnosis plot of pi
#'
#' @param x the output from a keyATM model (see [keyATM()]).
#' @param show_topic an integer or a vector. Indicate topics to visualize. Default is \code{NULL}.
#' @param start integer. The starting value of iteration to use in the plot. Default is \code{0}.
#' @param ci value of the credible interval (between 0 and 1) to be estimated. Default is \code{0.9} (90%). This is an option when calculating credible intervals (you need to set \code{store_pi = TRUE} in [keyATM()]).
#' @param method method for computing the credible interval. The Highest Density Interval (\code{hdi}, default) or Equal-tailed Interval (\code{eti}). This is an option when calculating credible intervals (you need to set \code{store_pi = TRUE} in [keyATM()]).
#' @param point method for computing the point estimate. \code{mean} (default) or \code{median}. This is an option when calculating credible intervals (you need to set \code{store_pi = TRUE} in [keyATM()]).
#' @return keyATM_fig object.
#' @import ggplot2
#' @import magrittr
#' @importFrom rlang .data
#' @seealso [save_fig()]
#' @export
plot_pi <- function(x, show_topic = NULL, start = 0, ci = 0.9, method = c("hdi", "eti"), point = c("mean", "median"))
{
  method <- rlang::arg_match(method)
  point <- rlang::arg_match(point)
  check_arg_type(x, "keyATM_output")
  modelname <- extract_full_model_name(x)

  if (modelname %in% c("lda", "ldacov", "ldahmm")) {
    cli::cli_abort(paste0("This is not a model with keywords."))
  }

  if (is.null(show_topic)) {
    show_topic <- 1:x$keyword_k
  } else if (sum(!show_topic %in% 1:x$keyword_k) != 0) {
    cli::cli_abort("`plot_pi` only visualize keyword topics.")
  }

  if (!is.numeric(start) | length(start) != 1) {
    cli::cli_abort("`start` argument is invalid.")
  }

  tnames <- c(names(x$keywords_raw))[show_topic]

  if (!is.null(x$values_iter$pi_iter)) {
    pi_mat <- t(sapply(x$values_iter$pi_iter, unlist, use.names = FALSE))[, show_topic, drop = FALSE]
    pi_mat %>%
      tibble::as_tibble(.name_repair = ~ tnames) %>%
      dplyr::mutate(Iteration = x$values_iter$used_iter) %>%
      dplyr::filter(.data$Iteration >= start) %>%
      dplyr::select(-tidyselect::all_of("Iteration")) -> pi_mat

    if (nrow(pi_mat) == 0) {
      cli::cli_abort("Nothing left to plot. Please check arguments.")
    }

    pi_mat %>%
      tidyr::pivot_longer(cols = dplyr::everything(), names_to = "Topic") %>%
      dplyr::group_by(.data$Topic) %>%
      dplyr::summarise(x = list(tibble::enframe(calc_ci(.data$value, ci, method, point), "q", "value")), .groups = "drop_last") %>%
      tidyr::unnest(x) %>% tidyr::pivot_wider(names_from = tidyselect::all_of("q"), values_from = tidyselect::all_of("value")) -> temp

    p <- ggplot(temp, aes(y = .data$Point, x = .data$Topic)) +
         theme_bw() + geom_point() +
         geom_errorbar(aes(ymin = .data$Lower, ymax = .data$Upper), data = temp, width = 0.01, linewidth = 1) +
         xlab("Topic") + ylab("Probability") +
         ggtitle("Probability of words drawn from keyword topic-word distribution") +
         theme(plot.title = element_text(hjust = 0.5))
  } else {
    cli::cli_alert_info("Plotting pi from the final MCMC draw. Please set `store_pi` to `TRUE` if you want to plot pi over iterations.")
    x$pi %>%
      dplyr::mutate(Probability = .data$Proportion / 100) %>%
      dplyr::filter(.data$Topic %in% (!!show_topic)) %>%
      dplyr::mutate(Topic = tnames) -> temp

    p <- ggplot(temp, aes(x = .data$Topic, y = .data$Probability)) +
         geom_bar(stat = "identity") +
         theme_bw() +
         xlab("Topic") + ylab("Probability") +
         ggtitle("Probability of words drawn from keyword topic-word distribution") +
         theme(plot.title = element_text(hjust = 0.5))
  }
  p <- list(figure = p, values = temp)
  class(p) <- c("keyATM_fig", class(p))
  return(p)
}


#' Show the expected proportion of the corpus belonging to each topic
#'
#' @param x the output from a keyATM model (see [keyATM()]).
#' @param n The number of top words to show. Default is \code{3}.
#' @param show_topic an integer or a vector. Indicate topics to visualize. Default is \code{NULL}.
#' @param show_topwords logical. Show topwords. The default is \code{TRUE}.
#' @param order The order of topics.
#' @param label_topic a character vector. The name of the topics in the plot.
#' @param xmax a numeric. Indicate the max value on the x axis
#' @return keyATM_fig object
#' @import magrittr
#' @import ggplot2
#' @importFrom rlang .data
#' @seealso [save_fig()]
#' @export
plot_topicprop <- function(x, n = 3, show_topic = NULL, show_topwords = TRUE, label_topic = NULL, order = c("proportion", "topicid"), xmax = NULL)
{
  check_arg_type(x, "keyATM_output")
  order <- rlang::arg_match(order)

  total_k <- x$keyword_k + x$no_keyword_topics
  if (is.null(show_topic)) {
    show_topic <- 1:total_k
  } else {
    if (max(show_topic) > total_k | min(show_topic) < 1) {
      cli::cli_abort("Invalid topic ID in `show_topic`.")
    }
  }

  topwords <- top_words(x, n = n)[, show_topic]

  if (!is.null(label_topic)) {
    if (length(label_topic) != ncol(topwords)) {
      cli::cli_abort("The length of `label_topic` is incorrect.")
    }
    colnames(topwords) <- label_topic
  }

  topwords %>%
    dplyr::summarise(dplyr::across(dplyr::everything(), ~stringr::str_c(.x, collapse = ", "))) %>%
    tidyr::pivot_longer(dplyr::everything(),
                        values_to = "Topwords",
                        names_to = "Topic") -> topwords_commas

  theta_use <- x$theta[, show_topic]
  if (!is.null(label_topic)) {
    colnames(theta_use) <- label_topic
  }

  theta_use %>%
    tibble::as_tibble() %>%
    dplyr::summarise(dplyr::across(dplyr::everything(), ~mean(.x))) %>%
    tidyr::pivot_longer(dplyr::everything(),
                        values_to = "Topicprop",
                        names_to = "Topic") %>%
    dplyr::left_join(topwords_commas, by = "Topic") -> theta_use_tbl

  if (order == "proportion") {
    theta_use_tbl %>%
      dplyr::mutate(Topic = stringr::str_remove(.data$Topic, "^\\d+_")) -> theta_use_tbl
    theta_use_tbl %>%
      dplyr::arrange(dplyr::desc(.data$Topicprop)) %>%
      dplyr::pull(.data$Topic) -> use_order
  } else if (order == "topicid") {
    theta_use_tbl %>%
      dplyr::pull(.data$Topic) -> use_order
  }

  theta_use_tbl %>%
    dplyr::mutate(
      Topic = factor(.data$Topic, levels = rev(use_order)),
      xpos = max(.data$Topicprop) + 0.01
    ) %>%
    dplyr::arrange(dplyr::desc(.data$Topic)) -> plot_obj

  if (is.null(xmax)) {
    if (show_topwords) {
      xmax <- min(max(plot_obj$Topicprop) * 2.5, 1)
    } else {
      xmax <- max(plot_obj$Topicprop) + 0.02
    }
  }

  label_percent <- function(x) {
    paste0(round(x * 100, 2), "%")
  }

  p <- ggplot(plot_obj, aes(x = .data$Topicprop, y = .data$Topic)) +
        geom_col() +
        {if (show_topwords)
            geom_text(
              aes(x = .data$xpos, y = .data$Topic, label = .data$Topwords),
              hjust = 0, size = max(10 / n + 1, 2.5))
        } +
        scale_x_continuous(
          "Expected topic proportions", limits = c(0, xmax), labels = label_percent
        ) +
        theme_bw() +
        theme(panel.grid.major.x = element_blank(),
              panel.grid.minor.x = element_blank(),
              panel.grid.major.y = element_blank(),
              panel.grid.minor.y = element_blank())

  p <- list(figure = p, values = plot_obj)
  class(p) <- c("keyATM_fig", class(p))
  return(p)
}



#' Plot document-topic distribution by strata (for covariate models)
#'
#' @param x a strata_doctopic object (see [by_strata_DocTopic()]).
#' @param show_topic a vector or an integer. Indicate topics to visualize.
#' @param var_name the name of the variable in the plot.
#' @param by `topic` or `covariate`. Default is by `topic`.
#' @param ci value of the credible interval (between 0 and 1) to be estimated. Default is \code{0.9} (90%).
#' @param method method for computing the credible interval. The Highest Density Interval (\code{hdi}, default) or Equal-tailed Interval (\code{eti}).
#' @param point method for computing the point estimate. \code{mean} (default) or \code{median}.
#' @param width numeric. Width of the error bars.
#' @param show_point logical. Show point estimates. The default is \code{TRUE}.
#' @param ... additional arguments not used.
#' @return keyATM_fig object.
#' @import ggplot2
#' @import magrittr
#' @importFrom rlang .data
#' @seealso [save_fig()], [by_strata_DocTopic()]
#' @export
plot.strata_doctopic <- function(x, show_topic = NULL, var_name = NULL, by = c("topic", "covariate"),
                                 ci = 0.9, method = c("hdi", "eti"), point = c("mean", "median"),
                                 width = 0.1, show_point = TRUE, ...)
{
  by <- rlang::arg_match(by)
  method <- rlang::arg_match(method)
  point <- rlang::arg_match(point)

  tables <- summary.strata_doctopic(x, ci, method, point)
  by_var <- x$by_var
  by_values <- x$by_values
  if (!is.null(var_name)) {
    by_var <- var_name
  }

  if (is.null(show_topic)) {
    show_topic <- 1:nrow(tables[[1]])
  }

  tables <- dplyr::bind_rows(tables)
  tnames <- unique(tables$Topic)
  num_keytopic <- sum(!grepl("Other_[0-9]+", tnames))
  topic_parse <- function(s) {
    if (grepl("Other_[0-9]+", s)) {
      return(as.numeric(strsplit(s, "_")[[1]][2]) + num_keytopic)
    } else {
      return(as.numeric(strsplit(s, "_")[[1]][1]))
    }
  }
  tables$TopicID <- purrr::map_dbl(tables$Topic, topic_parse)
  tables$Topic <- factor(tables$Topic, levels = tnames[order(purrr::map_dbl(tnames, topic_parse))])
  tables <- tables %>% dplyr::filter(.data$TopicID %in% show_topic)

  variables <- unique(tables$label)

  if (point == "mean") {
    ylabel <- expression(paste("Mean of ", theta))
  } else {
    ylabel <- expression(paste("Median of ", theta))
  }

  p <- ggplot(tables) +
        coord_flip() +
        scale_x_discrete(limits = rev(variables)) +
        xlab(paste0(by_var)) +
        ylab(ylabel) +
        guides(color = guide_legend(title = "Topic")) +
        theme_bw()

  if (by == "topic") {
    p <- p + geom_errorbar(width = width, aes(x = .data$label, ymin = .data$Lower, ymax = .data$Upper,
                group = .data$Topic), position = position_dodge(width = -1/2)) + facet_wrap(~Topic)
    if (show_point)
      p <- p + geom_point(aes(x = .data$label, y = .data$Point))
  } else {
    p <- p + geom_errorbar(width = width, aes(x = .data$label, ymin = .data$Lower, ymax = .data$Upper,
                group = .data$Topic, colour = .data$Topic), position = position_dodge(width = -1/2))
    if (show_point)
      p <- p + geom_point(aes(x = .data$label, y = .data$Point, colour = .data$Topic), position = position_dodge(width = -1/2))
  }

  p <- list(figure = p, values = tables)
  class(p) <- c("keyATM_fig", class(p))
  return(p)
}


#' Plot time trend
#'
#' @param x the output from the dynamic keyATM model (see [keyATM()]).
#' @param show_topic an integer or a vector. Indicate topics to visualize. Default is \code{NULL}.
#' @param time_index_label a vector. The label for time index. The length should be equal to the number of documents (time index provided to [keyATM()]).
#' @param ci value of the credible interval (between 0 and 1) to be estimated. Default is \code{0.9} (90%). This is an option when calculating credible intervals (you need to set \code{store_theta = TRUE} in [keyATM()]).
#' @param method method for computing the credible interval. The Highest Density Interval (\code{hdi}, default) or Equal-tailed Interval (\code{eti}). This is an option when calculating credible intervals (you need to set \code{store_theta = TRUE} in [keyATM()]).
#' @param point method for computing the point estimate. \code{mean} (default) or \code{median}. This is an option when calculating credible intervals (you need to set \code{store_theta = TRUE} in [keyATM()]).
#' @param xlab a character.
#' @param scales character. Control the scale of y-axis (the parameter in [ggplot2::facet_wrap()][ggplot2::facet_wrap]): \code{free} adjusts y-axis for parameters. Default is \code{fixed}.
#' @param show_point logical. The default is \code{TRUE}. This is an option when calculating credible intervals.
#' @param ... additional arguments not used.
#' @return keyATM_fig object.
#' @import ggplot2
#' @import magrittr
#' @importFrom rlang .data
#' @seealso [save_fig()]
#' @export
plot_timetrend <- function(x, show_topic = NULL, time_index_label = NULL,
                           ci = 0.9, method = c("hdi", "eti"), point = c("mean", "median"),
                           xlab = "Time", scales = "fixed", show_point = TRUE, ...)
{
  method <- rlang::arg_match(method)
  point <- rlang::arg_match(point)
  check_arg_type(x, "keyATM_output")
  modelname <- extract_full_model_name(x)
  if (!modelname %in% c("hmm", "ldahmm")) {
    cli::cli_abort(paste0("This is not a model with time trends."))
  }

  if (!is.null(time_index_label)) {
    if (length(x$values_iter$time_index) != length(time_index_label)) {
      cli::cli_abort("The length of `time_index_label` does not match with the number of documents.")
    }
    time_index <- time_index_label
  } else {
    time_index <- x$values_iter$time_index
  }
  time_index_tbl <- tibble::tibble(time_index = time_index, time_index_raw = x$values_iter$time_index) %>% dplyr::distinct()

  if (is.null(show_topic)) {
    show_topic <- 1:x$keyword_k
  }

  format_theta <- function(theta, time_index, tnames) {
    theta[, show_topic, drop = FALSE] %>%
      tibble::as_tibble(.name_repair = ~tnames) %>%
      dplyr::mutate(time_index = time_index) %>%
      tidyr::pivot_longer(-tidyselect::all_of("time_index"), names_to = "Topic", values_to = "Proportion") %>%
      dplyr::group_by(.data$time_index, .data$Topic) %>%
      dplyr::summarize(Proportion = base::mean(.data$Proportion), .groups = "drop_last") -> res
    return(res)
  }

  tnames <- colnames(x$theta)

  if (is.null(x$values_iter$theta_iter)) {
    dat <- format_theta(x$theta, time_index, tnames[show_topic])
    p <- ggplot(dat, aes(x = .data$time_index, y = .data$Proportion, group = .data$Topic)) +
          geom_line(linewidth = 0.8, color = "blue") + geom_point(size = 0.9)
  } else {
    dat <- dplyr::bind_rows(lapply(x$values_iter$theta_iter, format_theta, time_index, tnames[show_topic])) %>%
      dplyr::group_by(.data$time_index, .data$Topic) %>%
      dplyr::summarise(x = list(tibble::enframe(calc_ci(.data$Proportion, ci, method, point), "q", "value"))) %>%
      tidyr::unnest(tidyselect::all_of("x")) %>% dplyr::ungroup() %>%
      tidyr::pivot_wider(names_from = tidyselect::all_of("q"), values_from = tidyselect::all_of("value")) %>%
      stats::setNames(c("time_index", "Topic", "Lower", "Point", "Upper"))
    p <- ggplot(dat, aes(x = .data$time_index, y = .data$Point, group = .data$Topic)) +
          geom_ribbon(aes(ymin = .data$Lower, ymax = .data$Upper), fill = "gray75") +
          geom_line(linewidth = 0.8, color = "blue")

    if (show_point)
      p <- p + geom_point(size = 0.9)
  }
  p <- p + xlab(xlab) + ylab(expression(paste("Mean of ", theta))) +
        facet_wrap(~.data$Topic, scales = scales) + theme_bw() + theme(panel.grid.minor = element_blank())
  dat <- dplyr::left_join(dat, time_index_tbl, by = "time_index") %>%
    dplyr::mutate(state_id = x$values_iter$R_iter_last[.data$time_index_raw])
  p <- list(figure = p, values = dat)
  class(p) <- c("keyATM_fig", class(p))
  return(p)
}