Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix some bug + improve display + code clean

Browse Source
This commit is contained in:
unknown 2015-11-07 22:24:37 +01:00
parent 996645dc17
commit 7cb34e3ad6
1 changed files with 8 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
R-package/R/xgb.plot.tree.R
View File

@ -4,10 +4,7 @@
#' Plotting only works for boosted tree model (not linear model). #' Plotting only works for boosted tree model (not linear model).
#' #'
#' @importFrom data.table data.table #' @importFrom data.table data.table
#' @importFrom data.table set
#' @importFrom data.table rbindlist
#' @importFrom data.table := #' @importFrom data.table :=
#' @importFrom data.table copy
#' @importFrom magrittr %>% #' @importFrom magrittr %>%
#' @param feature_names names of each feature as a character vector. Can be extracted from a sparse matrix (see example). If model dump already contains feature names, this argument should be \code{NULL}. #' @param feature_names names of each feature as a character vector. Can be extracted from a sparse matrix (see example). If model dump already contains feature names, this argument should be \code{NULL}.
#' @param filename_dump the path to the text file storing the model. Model dump must include the gain per feature and per tree (parameter \code{with.stats = T} in function \code{xgb.dump}). Possible to provide a model directly (see \code{model} argument). #' @param filename_dump the path to the text file storing the model. Model dump must include the gain per feature and per tree (parameter \code{with.stats = T} in function \code{xgb.dump}). Possible to provide a model directly (see \code{model} argument).
@ -62,22 +59,18 @@ xgb.plot.tree <- function(feature_names = NULL, filename_dump = NULL, model = NU
allTrees <- xgb.model.dt.tree(feature_names = feature_names, model = model, n_first_tree = n_first_tree) allTrees <- xgb.model.dt.tree(feature_names = feature_names, model = model, n_first_tree = n_first_tree)
} }
allTrees[Feature != "Leaf" ,yesPath := paste(ID,"(", Feature, "<br/>Cover: ", Cover, "<br/>Gain: ", Quality, ")-->|< ", Split, "|", Yes, ">", Yes.Feature, "]", sep = "")]
allTrees[Feature != "Leaf" ,noPath := paste(ID,"(", Feature, ")-->|>= ", Split, "|", No, ">", No.Feature, "]", sep = "")]
allTrees[, label:= paste0(Feature, "\nCover: ", Cover, "\nGain: ", Quality)] allTrees[, label:= paste0(Feature, "\nCover: ", Cover, "\nGain: ", Quality)]
allTrees[, shape:= "rectangle"][Feature == "Leaf", shape:= "oval"] allTrees[, shape:= "rectangle"][Feature == "Leaf", shape:= "oval"]
allTrees[, filledcolor:= "Beige"][Feature == "Leaf", filledcolor:= "Khaki"] allTrees[, filledcolor:= "Beige"][Feature == "Leaf", filledcolor:= "Khaki"]
nodes <- DiagrammeR::create_nodes(nodes = allTrees[,ID], # rev is used to put the first tree on top.
label = allTrees[,label], nodes <- DiagrammeR::create_nodes(nodes = allTrees[,ID] %>% rev,
#type = c("lower", "lower", "upper", "upper"), label = allTrees[,label] %>% rev,
style = "filled", style = "filled",
color = "DimGray", color = "DimGray",
fillcolor= allTrees[,filledcolor], fillcolor= allTrees[,filledcolor] %>% rev,
shape = allTrees[,shape], shape = allTrees[,shape] %>% rev,
data = allTrees[,Feature], data = allTrees[,Feature] %>% rev,
fontname = "Helvetica" fontname = "Helvetica"
) )
@ -100,4 +93,4 @@ xgb.plot.tree <- function(feature_names = NULL, filename_dump = NULL, model = NU
# Avoid error messages during CRAN check. # Avoid error messages during CRAN check.
# The reason is that these variables are never declared # The reason is that these variables are never declared
# They are mainly column names inferred by Data.table... # They are mainly column names inferred by Data.table...
globalVariables(c("Feature", "ID", "Cover", "Quality", "Split", "Yes", "No", ".", "shape", "filledcolor")) globalVariables(c("Feature", "ID", "Cover", "Quality", "Split", "Yes", "No", ".", "shape", "filledcolor", "label"))