FreesearchR/R/regression_plot.R

#' Regression coef plot from gtsummary. Slightly modified to pass on arguments
#'
#' @param x (`tbl_regression`, `tbl_uvregression`)\cr
#'   A 'tbl_regression' or 'tbl_uvregression' object
##  #' @param remove_header_rows (scalar `logical`)\cr
##  #'   logical indicating whether to remove header rows
##  #'   for categorical variables. Default is `TRUE`
##  #' @param remove_reference_rows (scalar `logical`)\cr
##  #'   logical indicating whether to remove reference rows
##  #'   for categorical variables. Default is `FALSE`.
#' @param ... arguments passed to `ggstats::ggcoef_plot(...)`
#'
#' @returns ggplot object
#' @export
#'
#' @examples
#' \dontrun{
#' mod <- lm(mpg ~ ., default_parsing(mtcars))
#' p <- mod |>
#'   gtsummary::tbl_regression() |>
#'   plot(colour = "variable")
#' }
#'
plot.tbl_regression <- function(x,
                                plot_ref = TRUE,
                                remove_header_rows = TRUE,
                                remove_reference_rows = FALSE,
                                ...) {
  # check_dots_empty()
  gtsummary:::check_pkg_installed("ggstats")
  gtsummary:::check_not_missing(x)
  # gtsummary:::check_scalar_logical(remove_header_rows)
  # gtsummary:::check_scalar_logical(remove_reference_rows)

  df_coefs <- x$table_body

  if (isTRUE(remove_header_rows)) {
    df_coefs <- df_coefs |> dplyr::filter(!header_row %in% TRUE)
  }
  if (isTRUE(remove_reference_rows)) {
    df_coefs <- df_coefs |> dplyr::filter(!reference_row %in% TRUE)
  }

  # Removes redundant label
  df_coefs$label[df_coefs$row_type == "label"] <- ""

  # Add estimate value to reference level
  if (plot_ref == TRUE){
  df_coefs[df_coefs$var_type == "categorical" & is.na(df_coefs$reference_row),"estimate"] <- if (x$inputs$exponentiate) 1 else 0}

  p <- df_coefs |>
    ggstats::ggcoef_plot(exponentiate = x$inputs$exponentiate, ...)

  if (x$inputs$exponentiate){
    p <- symmetrical_scale_x_log10(p)
  }
  p
}


#' Wrapper to pivot gtsummary table data to long for plotting
#'
#' @param list a custom regression models list
#' @param model.names names of models to include
#'
#' @returns list
#' @export
#'
merge_long <- function(list, model.names) {
  l_subset <- list$tables[model.names]

  l_merged <- l_subset |> tbl_merge()

  df_body <- l_merged$table_body

  sel_list <- lapply(seq_along(l_subset), \(.i){
    endsWith(names(df_body), paste0("_", .i))
  }) |>
    setNames(names(l_subset))

  common <- !Reduce(`|`, sel_list)

  df_body_long <- sel_list |>
    purrr::imap(\(.l, .i){
      d <- dplyr::bind_cols(
        df_body[common],
        df_body[.l],
        model = .i
      )
      setNames(d, gsub("_[0-9]{,}$", "", names(d)))
    }) |>
    dplyr::bind_rows() |> dplyr::mutate(model=as_factor(model))

  l_merged$table_body <- df_body_long

  l_merged$inputs$exponentiate <- !identical(class(list$models$Multivariable$model), "lm")

  l_merged
}


#' Easily round log scale limits for nice plots
#'
#' @param data data
#' @param fun rounding function (floor/ceiling)
#' @param ... ignored
#'
#' @returns numeric vector
#' @export
#'
#' @examples
#' limit_log(-.1,floor)
#' limit_log(.1,ceiling)
#' limit_log(-2.1,ceiling)
#' limit_log(2.1,ceiling)
limit_log <- function(data,fun,...){
  fun(10^-floor(data)*10^data)/10^-floor(data)
}

#' Ensure symmetrical plot around 1 on a logarithmic x scale for ratio plots
#'
#' @param plot ggplot2 plot
#' @param breaks breaks used and mirrored
#' @param ... ignored
#'
#' @returns ggplot2 object
#' @export
#'
symmetrical_scale_x_log10 <- function(plot,breaks=c(1,2,3,5,10),...){
  rx <- ggplot2::layer_scales(plot)$x$get_limits()

  x_min <- floor(10*rx[1])/10
  x_max <- ceiling(10*rx[2])/10

  rx_min <- limit_log(rx[1],floor)
  rx_max <- limit_log(rx[2],ceiling)

  max_abs_x <- max(abs(c(x_min,x_max)))

  ticks <- log10(breaks)+(ceiling(max_abs_x)-1)
  browser()
  plot + ggplot2::scale_x_log10(limits=c(rx_min,rx_max),breaks=create_log_tics(10^ticks[ticks<=max_abs_x]))
}