R/analysis_covariates.R

In keyATM: Keyword Assisted Topic Models

#' Estimate subsetted topic-word distribution
#'
#' @param x the output from a keyATM model (see [keyATM()]).
#' @param keyATM_docs an object generated by [keyATM_read()].
#' @param by a vector whose length is the number of documents.
#'
#' @return strata_topicword object (a list).
#' @import magrittr
#' @export
by_strata_TopicWord <- function(x, keyATM_docs, by)
{
  # Check inputs
  if (!is.vector(by)) {
    cli::cli_abort("`by` should be a vector.")
  }
  if (!"Z" %in% names(x$kept_values)) {
    cli::cli_abort("`Z` and `S` should be in the output. Please check `keep` option in `keyATM()`.")
  }
  if (!"S" %in% names(x$kept_values)) {
    cli::cli_abort("`Z` and `S` should be in the output. Please check `keep` option in `keyATM()`.")
  }
  if (length(keyATM_docs$W_raw) != length(by)) {
    cli::cli_abort("The length of `by` should be the same as the length of documents.")
  }


  # Get unique values of `by`
  unique_val <- unique(by)
  tnames <- rownames(x$phi)

  # Get phi for each
  obj <- lapply(unique_val,
                function(val) {
                  doc_index <- which(by == val)
                  all_words <- unlist(keyATM_docs$W_raw[doc_index], use.names = FALSE)
                  all_topics <- as.integer(unlist(x$kept_values$Z[doc_index]), use.names = FALSE)
                  all_s <- as.integer(unlist(x$kept_values$S[doc_index]), use.names = FALSE)
                  pi_estimated <- keyATM_output_pi(x$kept_values$Z[doc_index],
                                                   x$kept_values$S[doc_index],
                                                   x$priors$gamma)
                  vocab <- sort(unique(all_words))
                  phi_obj <- keyATM_output_phi_calc_key(all_words, all_topics, all_s, pi_estimated,
                                                        x$keywords_raw,
                                                        vocab, x$priors, tnames, model = x)
                }
               )
  names(obj) <- unique_val

  res <- list(phi = obj, theta = x$theta, keywords_raw = x$keywords_raw)

  class(res) <- c("strata_topicword", class(res))
  return(res)
}


#' Show covariates information
#'
#' @param x the output from the covariate keyATM model (see [keyATM()]).
#' @export
covariates_info <- function(x) {
  if (x$model != "covariates" | !("keyATM_output" %in% class(x))) {
    cli::cli_abort("This is not an output of the covariate model")
  }
  cat(paste0("Colnames: ", paste(colnames(x$kept_values$model_settings$covariates_data_use), collapse = ", "),
             "\nStandardization: ", as.character(x$kept_values$model_settings$standardize),
             "\nFormula: ", paste(as.character(x$kept_values$model_settings$covariates_formula), collapse = " "), "\n\nPreview:\n"))
  print(utils::head(x$kept_values$model_settings$covariates_data_use))
}


#' Return covariates used in the iteration
#'
#' @param x the output from the covariate keyATM model (see [keyATM()])
#' @export
covariates_get <- function(x) {
  if (x$model != "covariates" | !("keyATM_output" %in% class(x))) {
    cli::cli_abort("This is not an output of the covariate model")
  }
  return(x$kept_values$model_settings$covariates_data_use)
}


#' Estimate document-topic distribution by strata (for covariate models)
#'
#' @param x the output from the covariate keyATM model (see [keyATM()]).
#' @param by_var character. The name of the variable to use.
#' @param labels character. The labels for the values specified in `by_var` (ascending order).
#' @param by_values numeric. Specific values for `by_var`, ordered from small to large. If it is not specified, all values in `by_var` will be used.
#' @param ... other arguments passed on to the [predict.keyATM_output()] function.
#' @return strata_topicword object (a list).
#' @import magrittr
#' @importFrom stats predict
#' @export
by_strata_DocTopic <- function(x, by_var, labels, by_values = NULL, ...)
{
  # Check inputs
  variables <- colnames(x$kept_values$model_settings$covariates_data_use)
  if (length(by_var) != 1)
    cli::cli_abort("`by_var` should be a single variable.")
  if (!by_var %in% variables)
    cli::cli_abort(paste0(by_var, " is not in the set of covariates in keyATM model. Check with `covariates_info()`.",
                "Covariates provided are: ",
                paste(colnames(x$kept_values$model_settings$covariates_data_use), collapse = " , ")))

  # Info
  if (is.null(by_values)) {
    by_values <- sort(unique(x$kept_values$model_settings$covariates_data_use[, by_var]))
  }
  if (length(by_values) != length(labels)) {
    cli::cli_abort("Length mismatches. Please check `labels`.")
  }

  apply_predict <- function(i, ...) {
    value <- by_values[i]
    new_data <- x$kept_values$model_settings$covariates_data_use
    new_data[, by_var] <- value
    obj <- predict(object = x, newdata = new_data, raw_values = TRUE, ...)
    return(obj)
  }

  set.seed(x$options$seed)
  res <- lapply(1:length(by_values), apply_predict, ...)
  names(res) <- by_values
  res <- tibble::as_tibble(res)

  # Making CI
  tables <- lapply(1:length(by_values),
                  function(index) {
                     theta <- res[[index]]
                     return(strata_doctopic_CI(theta[, 1:(ncol(theta)-1)],
                                               label = labels[index], ...))
                  })
  names(tables) <- labels


  # Making a return object
  obj <- list(theta = res, tables = tables, by_values = by_values, by_var = by_var, labels = labels)
  class(obj) <- c("strata_doctopic", class(obj))
  return(obj)
}


#' @noRd
#' @export
print.strata_doctopic <- function(x, ...)
{
  cat(paste0("strata_doctopic object for the ", x$by_var, "\n"))
}


#' @noRd
#' @export
summary.strata_doctopic <- function(object, ...)
{
  return(object$tables)
}


#' @noRd
#' @import magrittr
#' @keywords internal
strata_doctopic_CI <- function(theta, ci = 0.9, method = c("hdi", "eti"), point = c("mean", "median"), label = NULL, ...)
{
  method <- rlang::arg_match(method)
  point <- rlang::arg_match(point)

  q <- as.data.frame(apply(theta, 2, calc_ci, ci, method, point))
  q$CI <- c("Lower", "Point", "Upper")
  q %>%
    tidyr::gather(key = "Topic", value = "Value", -"CI") %>%
    tidyr::spread(key = "CI", value = "Value") %>%
    dplyr::mutate(TopicID = 1:(dplyr::n())) %>% tibble::as_tibble() -> res

  if (!is.null(label)) {
    res %>% dplyr::mutate(label = label) -> res
  }
  return(res)
}


#' @noRd
#' @keywords internal
calc_ci <- function(vec, ci, method, point)
{
  # Check bayestestR package and Kruschke's book
  if (point == "mean") {
    point <- mean(vec)
  }
  if (point == "median") {
    point <- stats::median(vec)
  }

  if (method == "hdi") {
    sorted_points <- sort.int(vec, method = "quick")
    window_size <- ceiling(ci * length(sorted_points))
    nCIs <-  length(sorted_points) - window_size

    if (window_size < 2 | nCIs < 1) {
      cli::cli_warn("`ci` is too small or interations are not enough, using `eti` option instead.")
      return(calc_ci(vec, ci, method = "eti", point))
    }

    ci_width <- sapply(1:nCIs, function(x) {sorted_points[x + window_size] - sorted_points[x]})
    slice_min <- which.min(ci_width)

    res <- c(sorted_points[slice_min], point, sorted_points[slice_min + window_size])
  }

  if (method == "eti") {
    qua <- stats::quantile(vec, probs = c((1 - ci) / 2, (1 + ci) / 2))
    res <- c(qua[1], point, qua[2])
  }

  names(res) <- c("Lower", "Point", "Upper")
  return(res)
}