From 3dd202a19eb318a16bfa0b98fd0f70fb78a927b4 Mon Sep 17 00:00:00 2001 From: El Potaeto Date: Tue, 6 Jan 2015 18:18:55 +0100 Subject: [PATCH] Add stat indicators in plot --- R-package/NAMESPACE | 1 + R-package/R/xgb.plot.tree.R | 11 ++++++----- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/R-package/NAMESPACE b/R-package/NAMESPACE index f2baee12b..f68eafbc5 100644 --- a/R-package/NAMESPACE +++ b/R-package/NAMESPACE @@ -20,6 +20,7 @@ importClassesFrom(Matrix,dgeMatrix) importFrom(DiagrammeR,DiagrammeR) importFrom(data.table,":=") importFrom(data.table,as.data.table) +importFrom(data.table,copy) importFrom(data.table,data.table) importFrom(data.table,rbindlist) importFrom(data.table,set) diff --git a/R-package/R/xgb.plot.tree.R b/R-package/R/xgb.plot.tree.R index 4863fd7ca..a263fe989 100644 --- a/R-package/R/xgb.plot.tree.R +++ b/R-package/R/xgb.plot.tree.R @@ -6,10 +6,11 @@ #' @importFrom data.table data.table #' @importFrom data.table set #' @importFrom data.table rbindlist +#' @importFrom data.table := +#' @importFrom data.table copy #' @importFrom magrittr %>% #' @importFrom magrittr not #' @importFrom magrittr add -#' @importFrom data.table := #' @importFrom stringr str_extract #' @importFrom stringr str_split #' @importFrom stringr str_extract @@ -93,10 +94,11 @@ xgb.plot.tree <- function(feature_names = NULL, filename_dump = NULL, n_first_tr coverLeaf <- extract(leaf, "cover=\\d*\\.*\\d*") dt <- data.table(ID = c(idBranch, idLeaf), Feature = c(featureBranch, featureLeaf), Split = c(splitBranch, splitLeaf), Yes = c(yesBranch, yesLeaf), No = c(noBranch, noLeaf), Missing = c(missingBranch, missingLeaf), Quality = c(qualityBranch, qualityLeaf), Cover = c(coverBranch, coverLeaf))[order(ID)][,Tree:=i] - set(dt, i = which(dt[,Feature]!= "Leaf"), j = "YesFeature", value = dt[ID == dt[,Yes], Feature]) - set(dt, i = which(dt[,Feature]!= "Leaf"), j = "NoFeature", value = dt[ID == dt[,No], Feature]) + set(dt, i = which(dt[,Feature]!= "Leaf"), j = "YesFeature", value = merge(copy(dt)[,ID:=Yes][, .(ID)], dt[,.(ID, Feature, Quality, Cover)], by = "ID")[,paste(Feature, "
Cover: ", Cover, sep = "")]) - dt[Feature!="Leaf" ,yesPath:= paste(ID,"[", Feature, "]-->|< ", Split, "|", Yes, "[", YesFeature, "]", sep = "")] + set(dt, i = which(dt[,Feature]!= "Leaf"), j = "NoFeature", value = merge(copy(dt)[,ID:=No][, .(ID)], dt[,.(ID, Feature, Quality, Cover)], by = "ID")[,paste(Feature, "
Cover: ", Cover, sep = "")]) + + dt[Feature!="Leaf" ,yesPath:= paste(ID,"[", Feature, "
Cover: ", Cover, "
Gain: ", Quality, "]-->|< ", Split, "|", Yes, "[", YesFeature, "]", sep = "")] dt[Feature!="Leaf" ,noPath:= paste(ID,"[", Feature, "]-->|>= ", Split, "|", No, "[", NoFeature, "]", sep = "")] @@ -112,6 +114,5 @@ xgb.plot.tree <- function(feature_names = NULL, filename_dump = NULL, n_first_tr no <- allTrees[Feature!="Leaf", c(No)] %>% paste(collapse = ",") %>% paste("class ", ., " redNode", sep = "") path <- allTrees[Feature!="Leaf", c(yesPath, noPath)] %>% .[order(.)] %>% paste(sep = "", collapse = ";") %>% paste("graph LR", .,collapse = "", sep = ";") %>% paste(styles, yes, no, sep = ";") - DiagrammeR(path) }