PAaSO/fun/functions.R

create_log_tics <- function(data){
  sort(round(unique(c(1/data,data)),2))
}


#' Test if input is within specified range or range of vector
#'
#' @param data numeric vector
#' @param vec numeric vector
#'
#' @returns logical
#' @export
#'
#' @examples
#' in_range(2,c(3,5))
#' in_range(c(2,4),c(3,5))
in_range <- function(data, vec) {
  ## Extra safety to ensure correct order of min/max
  rng <- range(vec)
  ## Test
  sapply(data, \(.x){
    .x >= rng[1] & .x <= rng[2]
  })
}

simple_shap_plot <- function(data, title = "Coefficients") {
  df <- data |> dplyr::mutate(var = forcats::as_factor(as.character(var)))
  
  d.tics <- c(.5, .8, 1, 1.5, 3)
  # browser()
  
  all.tics <- create_log_tics(d.tics)
  
  if (all(range(df$val) < 1)) {
    if (!in_range(min(df$val), all.tics)) {
      all.tics <- create_log_tics(c(d.tics, floor(min(df$val))))
    }
    # log(df$val)
    x.tics <- all.tics[all.tics <= 1]
  } else if (all(range(df$val) > 1)) {
    if (!in_range(max(df$val), all.tics)) {
      all.tics <- create_log_tics(c(d.tics, ceiling(max(df$val))))
    }
    
    x.tics <- all.tics[all.tics >= 1]
  } else {
    x.tics <- all.tics
  }
  
  ggplot2::ggplot(df, ggplot2::aes(x = val, y = var)) +
    ggplot2::geom_bar(fill = "#fca50a", width = 2 / 3, stat = "identity") +
    # ggplot2::geom_text(ggplot2::aes(label=val), hjust=-1)+
    ggplot2::scale_y_discrete(limits = rev) +
    # ggplot2::scale_x_continuous(labels=function(x)x+1)+
    ggplot2::scale_x_continuous(
      transform = "log",
      breaks = x.tics,
      labels = x.tics,
      limits = range(x.tics)
    ) +
    ggplot2::ggtitle(title) +
    ggplot2::labs(x = "median(OR)", y = ggplot2::element_blank()) +
    ggplot2::theme_bw() +
    ggplot2::theme(
      panel.grid.minor = ggplot2::element_blank(),
      # legend.title = ggplot2::element_text(""),
      legend.position = "bottom"
    )
}


forest_plot <- function(data,
                        only_sign=FALSE) {
  
  df <- data |> 
    dplyr::mutate(var = forcats::as_factor(as.character(var)))
    
  # browser()
  df$is_sign <- apply(df[c("lo","hi")],1,\(.x){
    !in_range(0,.x) 
    })
  
  if (!isTRUE(only_sign)){
    cols <- c("gray50","#82AE36")
  } else {
    cols <- "#82AE36"
    df <- df[df$is_sign,]
  }
  
  df |>
    ggplot2::ggplot(ggplot2::aes(x = beta, y = var,color=is_sign))+
    ggplot2::geom_vline(ggplot2::aes(xintercept = 0), linewidth = .5, linetype = "dashed") +
    ggplot2::geom_point(#ggplot2::aes(shape = model),
                        # position = ggplot2::position_dodge(width = dodge.width),
                        size = 6
    ) +
    ggplot2::geom_errorbarh(ggplot2::aes(xmax = hi, xmin = lo),
                            # position = ggplot2::position_dodge(width = dodge.width),
                            size = .8,
                            height = .5,
                            color = "gray50"
    ) +
    # ggplot2::position_dodge(width = 2, preserve = "total")+
    ggplot2::scale_x_continuous(
      # breaks = seq(min(df$lo,na.rm=TRUE),max(df$hi,na.rm=TRUE))
    ) +
    ggplot2::scale_y_discrete(limits = rev) +
    ggplot2::scale_color_manual(values = cols)+
    ggplot2::theme_bw() +
    ggplot2::theme(
      panel.grid.minor = ggplot2::element_blank(),
      # legend.title = ggplot2::element_text(""),
      legend.position = "none"
    ) +
    ggplot2::ylab("") +
    ggplot2::xlab("Mean PASE score difference")
}


pase_table_clean <- function(data){
  suppressWarnings(
    list(Univariate=data[1:3],
         Multivariate=data[c(1,4,5)],
         "Imputed multivariate"=data[c(1,6,7)]) |> 
      purrr::imap(\(.x,.i){
        out <- setNames(.x,c("var","val","ci")) |> 
          tidyr::separate(
            col = "ci", into = c("lo", "hi"), sep = ", ", convert = TRUE
          )|>
          dplyr::mutate(dplyr::across(tidyselect::all_of(c("val", "lo", "hi")), as.numeric))
        # browser()
        out$var <- sapply(out$var,\(.s){
          l <- nchar(.s)
          substr(.s,5,l)
        })
        # out[1,2:4] <- 1
        out$model <- .i
        out$var <- sapply(out$var,\(.x){
          paste(strwrap(.x, 25), collapse = "\n")})
        
        out
      }) |> 
      dplyr::bind_rows()
  )
}

forest_plot_grp <- function(data,
                            # group.colors = viridisLite::viridis(4,option = "D"),
                            group.colors=NULL,
                            dodge.width=-.8,
                            rev_y=TRUE) {
  
  df <- data |> 
    dplyr::mutate(var = forcats::as_factor(as.character(var)),
                  model = forcats::as_factor(as.character(model)))
  
  d.tics <- c(.5, 1, 1.5, 3)
  # browser()
  
  all.tics <- create_log_tics(d.tics)
  
  if (!in_range(min(df$lo,na.rm=TRUE), all.tics)) {
    all.tics <- create_log_tics(c(d.tics, floor(min(df$lo,na.rm=TRUE)*10)/10))
  }
  
  if (!in_range(max(df$hi,na.rm=TRUE), all.tics)) {
    all.tics <- create_log_tics(c(d.tics, ceiling(max(df$hi,na.rm=TRUE))))
  }
  
  if (all(range(df$val) < 1)) {
    x.tics <- all.tics[all.tics <= 1]
  } else if (all(range(df$val) > 1)) {
    x.tics <- all.tics[all.tics >= 1]
  } else {
    x.tics <- all.tics
  }
  
  df$is_sign <- apply(df[c("lo","hi")],1,\(.x){
    out <- !in_range(1,.x)
    out[is.na(out)] <- FALSE
    out
  })
  # 
  # if (!isTRUE(only_sign)){
  #   cols <- c("gray50","#82AE36")
  # } else {
  #   cols <- "#82AE36"
  #   df <- df[df$is_sign,]
  # }
  # browser()
  
  out <- df |>
    ggplot2::ggplot(ggplot2::aes(x = val, y = var, xmax = hi, xmin = lo))+
    ggplot2::geom_vline(ggplot2::aes(xintercept = 1), linewidth = .5, linetype = "dashed") +
    ## Adding the grouping solved the errorbar not dodging correctly
    ## found here: https://www.reddit.com/r/rstats/comments/1ew20sh/ggplot2s_position_dodge_doesnt_work_with_geom/
    ggplot2::geom_errorbarh(ggplot2::aes(group = interaction(var,model)),
                            position = ggplot2::position_dodge(width = dodge.width,preserve = "total"),
                            size = .8,
                            height = .5,
                            color = "gray50"
    ) +
    ggplot2::geom_point(ggplot2::aes(color = var, fill=var, shape=model),
      position = ggplot2::position_dodge(width = dodge.width,preserve = "total"),
      size = 6
    ) +
    # ggplot2::position_dodge(width = 2, preserve = "total")+
    ggplot2::scale_x_continuous(
      transform = "log",
      breaks = x.tics,
      labels = x.tics,
      limits = range(x.tics)
    )  +
    # ggplot2::scale_color_manual(values = cols)+
    ggplot2::theme_bw() +
    ggplot2::theme(
      panel.grid.minor = ggplot2::element_blank(),
      # legend.title = ggplot2::element_text(""),
      legend.position = "bottom"
    ) +
    ggplot2::ylab("") +
    ggplot2::xlab("OR")
  
  if (isTRUE(rev_y)){
    out <- out +
      ggplot2::scale_y_discrete(limits = rev)
  } 
  
  if (!is.null(group.colors)){
    out <- out+
      ggplot2::scale_color_manual(values = group.colors) +
      ggplot2::scale_fill_manual(values = group.colors)
  }
  
  out
}