gbm_contour_eg.R

library("gbm")

plot_gbm_contour <- function(gbm_obj, i.var, ...){
  ##############################################################################
  # Plot a 3D contour plot for a gbm's partial dependencies.                   #
  # Args:                                                                      #
  #   gbm_obj: The gbm object.                                                 #
  #   i.var:   Character vector of variables to show.                          #
  #   ...:     Additional args passed to wireframe.                            #
  # Returns:                                                                   #
  #   Plot object.                                                             #
  # Author: [email protected]                                                 #
  ##############################################################################                                                                      
  library("lattice")
  inter_grid <- gbm:::plot.gbm(gbm_obj, i.var = i.var, return.grid = TRUE)
  inter_formula <- as.formula(paste("y ~", paste(i.var, collapse = "*")))
  wireframe(inter_formula, data = inter_grid, ...)
}

N <- 4000
df <- data.frame(noise_1 = runif(N),
                 inter_a = runif(N),
                 inter_b = runif(N))
 
df$response <- 3 * (df$inter_a * df$inter_b) ^ 2

g <- gbm(response ~ ., 
         data              = df,
         n.trees           = 4000,
         interaction.depth = 2,
         shrinkage         = 0.01,
         verbose           = TRUE)

plot_gbm_contour(g, c("inter_a", "inter_b"), 
                 drape    = TRUE,
                 colorkey = TRUE,
                 screen   = list(z = 25, x = -70))