R/plot_ALE.R

In gloverd2/admr: Contains standard functions for Admiral Advance Analytics

#' plot_ALM
#'
#' @description
#' Plot a 1D Accumulated Local Effects Plot (ALE)
#'
#' @param data dataframe - Data to use for the predictions. Should have same columns as training data (can be training data)
#' @param model model object - Model to produce ALM from
#' @param explain_col string - This is the factor to explain in the model. The columns must be in \code{data}
#' @param weight numeric - Vector of length \code{nrow(data)} contains weightings, if NULL even weighting is used
#' @param n_bins numeric - This is the number of points to calculate the PDP for
#' @param use_plotly Optional: boolean - If TRUE plotly object is returned else ggplot2 object
#'
#' @seealso plot_PDP
#'
#' @return plotly/ggplot object of ALE plot
#' @export
#'
#' @examples
#'
#' data <- data.frame(x1=runif(100, 0, 25), x2=runif(100, 0, 25)) %>%
#'   mutate(target=x1^2 * 0.01 + x2 + rnorm(n(),sd=5))
#'
#' #LM
#' model_lm <- glm(target ~ poly(x1, 2) + x2, data=data)
#'
#' plot_ALE(data, model_lm, explain_col="x1", n_bins=5)
#' plot_ALE(data, model_lm, explain_col="x2", n_bins=5)
#' #plot_ALE(data, model_lm, explain_col=c("x1","x2"), n_bins=5)
#'
#' #GLM
#' model_glm <- glm(target ~ poly(x1, 2) + x2, data=data)
#'
#' plot_ALE(data, model_glm, explain_col="x1", n_bins=5)
#' plot_ALE(data, model_glm, explain_col="x2", n_bins=5)
#' #plot_ALE(data, model_glm, explain_col=c("x1","x2"), n_bins=5)
#'
#' #GBM
#' model_gbm <- xgboost(data = as.matrix(data[,which(!(names(data)=="target"))]), label=data[["target"]], nrounds=20, verbose = 0)
#' plot_ALE(data[,which(!(names(data)=="target"))], model_gbm, explain_col="x1", n_bins=10)
#' plot_ALE(data[,which(!(names(data)=="target"))], model_gbm, explain_col="x2", n_bins=10)
#' #plot_ALE(data[,which(!(names(data)=="target"))], model_gbm, explain_col=c("x1","x2"), n_bins=10)
#'
plot_ALE <- function(data, model, explain_col, weight=NULL, n_bins=10, use_plotly=TRUE){

  # Use no weighting if none given
  if (is.null(weight)){
    weight <- rep(1, nrow(data))
  }else{
    checkmate::assert_numeric(weight, len=nrow(data), lower=0)
  }

  #Check inputs
  checkmate::assert_data_frame(data)
  checkmate::assert_character(explain_col, min.len = 1, max.len = 2)
  checkmate::assert_true(all(explain_col %in% names(data)))
  checkmate::assert_integerish(n_bins, len = 1)
  checkmate::assert_logical(use_plotly, len=1)



  plot.data <- plotting_numerical_buckets(var_to_band=data[[explain_col]], n_bins=n_bins, weight=weight, include_outliers=TRUE)


  # For each bin
  for (ii in 1:nrow(plot.data)){
    weight_ii <- weight[data[[explain_col]] >= plot.data[ii,"lower"] & data[[explain_col]] < plot.data[ii,"upper"]]
    if (sum(weight_ii) == 0){
      plot.data[ii,"gradient_pred"]=0
      plot.data[ii,"weight"]=0
    }else{


      data_ii_lower <- data %>% dplyr::filter(.data[[explain_col]] >= plot.data[ii,"lower"] & .data[[explain_col]] < plot.data[ii,"upper"])
      data_ii_upper <- data_ii_lower

      data_ii_lower[[explain_col]] <- plot.data[[ii,"lower"]] # Set the factor to the bucket min
      data_ii_upper[[explain_col]] <- plot.data[[ii,"upper"]] # Set the factor to the bucket max

      # Get average prediction
      if(any(class(model)=="xgb.Booster")){
        plot.data[ii,"lower_pred"] <- sum(predict(object=model, newdata=as.matrix(data_ii_lower)) %>% as.vector() * weight_ii) / sum(weight_ii)
        plot.data[ii,"upper_pred"] <- sum(predict(object=model, newdata=as.matrix(data_ii_upper)) %>% as.vector() * weight_ii) / sum(weight_ii)
      }else{
        plot.data[ii,"lower_pred"] <- sum(predict(object=model, newdata=data_ii_lower) %>% as.vector() * weight_ii) / sum(weight_ii)
        plot.data[ii,"upper_pred"] <- sum(predict(object=model, newdata=data_ii_upper) %>% as.vector() * weight_ii) / sum(weight_ii)
      }
      plot.data[ii,"gradient_pred"] <- (plot.data[ii,"upper_pred"] - plot.data[ii,"lower_pred"])

      plot.data[ii,"weight"] <- sum(weight_ii)
    }
  }

  plot.data <- plot.data %>%
    dplyr::mutate(ALM = cumsum(gradient_pred),
                  ALM0 = ALM - (sum(ALM * weight)/sum(weight)))

  if (use_plotly==TRUE){
      plotly::plot_ly(data=plot.data) %>%
      plotly::add_trace(y=~ALM0, x=~center, type="scatter", mode="lines+markers", name=explain_col) %>%
      plotly::add_trace(y=~weight, x=~center, type="bar", name="Exposure", yaxis="y2", opacity=0.2) %>%
      plotly::layout(
                    title = paste0("ALM for factor ", explain_col),
                    yaxis = list(title="Prediction (Accumulated Local Effect)"),
                    yaxis2 = list(overlaying = "y", side = "right", title="Exposure", showgrid = FALSE, rangemode="nonnegative"),
                    xaxis = list(title=explain_col)
      ) %>%
      return()
  }else{
    scale = mean(plot.data$ALM) / mean(plot.data$weight)

    out_plot <- ggplot2::ggplot(data=plot.data) +
      ggplot2::geom_line(ggplot2::aes(y=ALM0, x=center, color=explain_col)) +
      ggplot2::geom_point(ggplot2::aes(y=ALM0, x=center, color=explain_col)) +
      ggplot2::geom_rect(ggplot2::aes(ymax=weight  * scale, ymin=0, xmin=center-(width/2), xmax=center+(width/2), fill="Exposure"), alpha=0.2) +
      ggplot2::scale_y_continuous(name="Prediction (Accumulated Local Effect)", sec.axis=ggplot2::sec_axis(~.*scale ,name="Exposure")) +
      ggplot2::labs(title = paste0("ALM for factor ", explain_col), x=explain_col)

    return(out_plot)
  }
}