0.6-4 submission (#1935)

This commit is contained in:
Tong He
2017-01-04 23:31:05 -08:00
committed by GitHub
parent 8b827425b2
commit ce84af7923
5 changed files with 40 additions and 34 deletions

View File

@@ -78,7 +78,7 @@ xgb.plot.multi.trees <- function(model, feature_names = NULL, features_keep = 5,
nodes.dt <- tree.matrix[,.(Quality = sum(Quality)),by = .(abs.node.position, Feature)][,.(Text =paste0(Feature[1:min(length(Feature), features_keep)], " (", Quality[1:min(length(Quality), features_keep)], ")") %>% paste0(collapse = "\n")), by=abs.node.position]
edges.dt <- tree.matrix[Feature != "Leaf",.(abs.node.position, Yes)] %>% list(tree.matrix[Feature != "Leaf",.(abs.node.position, No)]) %>% rbindlist() %>% setnames(c("From", "To")) %>% .[,.N,.(From, To)] %>% .[,N:=NULL]
nodes <- DiagrammeR::create_nodes(nodes = nodes.dt[,abs.node.position],
nodes <- DiagrammeR::create_node_df(n = nrow(nodes.dt),
label = nodes.dt[,Text],
style = "filled",
color = "DimGray",
@@ -87,8 +87,8 @@ xgb.plot.multi.trees <- function(model, feature_names = NULL, features_keep = 5,
fontname = "Helvetica"
)
edges <- DiagrammeR::create_edges(from = edges.dt[,From],
to = edges.dt[,To],
edges <- DiagrammeR::create_edge_df(from = match(edges.dt[,From], nodes.dt[,abs.node.position]),
to = match(edges.dt[,To], nodes.dt[,abs.node.position]),
color = "DimGray",
arrowsize = "1.5",
arrowhead = "vee",
@@ -96,8 +96,7 @@ xgb.plot.multi.trees <- function(model, feature_names = NULL, features_keep = 5,
rel = "leading_to")
graph <- DiagrammeR::create_graph(nodes_df = nodes,
edges_df = edges,
graph_attrs = "rankdir = LR")
edges_df = edges)
DiagrammeR::render_graph(graph, width = plot_width, height = plot_height)
}

View File

@@ -49,7 +49,7 @@ xgb.plot.tree <- function(feature_names = NULL, model = NULL, n_first_tree = NUL
allTrees[, filledcolor:= "Beige"][Feature == "Leaf", filledcolor:= "Khaki"]
# rev is used to put the first tree on top.
nodes <- DiagrammeR::create_nodes(nodes = allTrees[,ID] %>% rev,
nodes <- DiagrammeR::create_node_df(n = length(allTrees[,ID] %>% rev),
label = allTrees[,label] %>% rev,
style = "filled",
color = "DimGray",
@@ -59,8 +59,8 @@ xgb.plot.tree <- function(feature_names = NULL, model = NULL, n_first_tree = NUL
fontname = "Helvetica"
)
edges <- DiagrammeR::create_edges(from = allTrees[Feature != "Leaf", c(ID)] %>% rep(2),
to = allTrees[Feature != "Leaf", c(Yes, No)],
edges <- DiagrammeR::create_edge_df(from = match(allTrees[Feature != "Leaf", c(ID)] %>% rep(2), allTrees[,ID] %>% rev),
to = match(allTrees[Feature != "Leaf", c(Yes, No)],allTrees[,ID] %>% rev),
label = allTrees[Feature != "Leaf", paste("<",Split)] %>% c(rep("",nrow(allTrees[Feature != "Leaf"]))),
color = "DimGray",
arrowsize = "1.5",
@@ -69,8 +69,7 @@ xgb.plot.tree <- function(feature_names = NULL, model = NULL, n_first_tree = NUL
rel = "leading_to")
graph <- DiagrammeR::create_graph(nodes_df = nodes,
edges_df = edges,
graph_attrs = "rankdir = LR")
edges_df = edges)
DiagrammeR::render_graph(graph, width = plot_width, height = plot_height)
}