R/cov_plots.R

In networkscaleup: Network Scale-Up Models for Aggregated Relational Data

#' Covariance plots
#'
#' Plots of the estimated covariance structure from a given fitted model
#'
#' @param ard ard matrix
#' @param model_fit a fitted object from [fit_mle()] or [fit_map()]
#' @param x_cov covariate matrix
#' @param resid_type the type of residuals to use
#' @param method the method to use
#' @param se whether to compute standard errors of estimates
#'
#' @return a list of ggplots, corresponding to covariance structure
#' @export
#'
#' @importFrom rlang .data
cov_plots <- function(ard,
                      model_fit,
                      x_cov,
                      resid_type = c("rqr", "pearson_residuals"),
                      method = "lm",
                      se = F) {
  ## Grab family
  family <- model_fit$family
  ## Grab residual type
  resid_type <- match.arg(resid_type, c("rqr", "pearson_residuals"))
  ## Obtain residuals
  resid_mat <- model_fit[[resid_type]]
  alpha_est <- model_fit$alphas
  ## Convert ard to data.frame, if not already
  if (!inherits(ard, "data.frame")) {
    ard <- data.frame(ard)
  }
  ## Convert x_cov to data.frame, if not already
  if (!inherits(x_cov, "data.frame")) {
    x_cov <- data.frame(x_cov) ## Auto-names to X1:Xp if no names
  }
  ## Standardize x_cov to have same min-max
  x_cov <- as.data.frame(lapply(x_cov, function(x) {
    (x - min(x)) / (max(x) - min(x))
  }))

  ard_x <- data.frame(resid = resid_mat, cov = x_cov, alpha = alpha_est)
  ard_long <- ard_x |>
    tidyr::pivot_longer(
      cols = tidyselect::starts_with("resid."),
      names_to = "Group",
      values_to = "resid"
    )
  ard_longer <- ard_long |>
    tidyr::pivot_longer(
      cols = tidyselect::starts_with("cov."),
      names_to = "cov_names",
      values_to = "CovValue"
    ) |>
    dplyr::mutate(
      cov_label = stringr::str_remove(.data$cov_names, "^cov\\."),
      cov_label = factor(.data$cov_label, levels = colnames(x_cov))
    )
  ## Produce plot 1, group-specific plots
  gg1 <- ggplot2::ggplot(ard_longer, ggplot2::aes(
    x = .data$CovValue,
    y = .data$resid,
    col = .data$Group,
    group = .data$Group
  )) +
    ggplot2::geom_smooth(method = method, se = se) +
    ggplot2::facet_wrap(~cov_label, scales = "free") +
    ggplot2::theme_minimal() +
    ggplot2::labs(x = "Covariate Value", y = "Residual", title = "Residuals vs Covariates") +
    ggplot2::scale_color_discrete(guide = "none")


  ## Find endpoints to add covariate label
  label_df <- ard_longer |>
    split(ard_longer$cov_label) |>
    purrr::imap_dfr(~ {
      sm <- suppressMessages(
        ggplot2::ggplot_build(
          ggplot2::ggplot(.x, ggplot2::aes(x = CovValue, y = alpha)) +
            ggplot2::geom_smooth(method = method, se = se)
        )$data[[1]]
      )

      sm |>
        dplyr::filter(x == max(x, na.rm = TRUE)) |>
        dplyr::mutate(cov_label = .y)
    })

  # Produce plot 2, averaged over groups
  gg2 <- ggplot2::ggplot(
    ard_longer,
    ggplot2::aes(
      x = .data$CovValue,
      y = .data$alpha,
      color = .data$cov_label
    )
  ) +
    ggplot2::geom_smooth(method = method, se = se) +
    ggplot2::geom_text(
      data = label_df,
      ggplot2::aes(
        x = .data$x, y = .data$y,
        label = .data$cov_label, color = .data$cov_label
      ),
      hjust = -0.1,
      show.legend = FALSE
    ) +
    ggplot2::expand_limits(x = max(label_df$x) * 1.1) +
    ggplot2::theme_minimal() +
    ggplot2::labs(
      x = "Covariate Value",
      y = "Estimated Respondent Effect",
      title = "Respondent Effect vs Covariate",
      color = "Covariate"
    )

  return(list(Group_plot = gg1, Respondent_plot = gg2))
}