xgb.plot.tree节点颜色填充和文本修饰

时间:2019-01-08 02:55:02

标签: xgboost diagrammer

我正在尝试使通过xgboost创建的树图更具可读性。我发现其他一些建议使用“ diagrammer”包来更改字体和颜色特征的帖子,但我发现仅在创建新图形时设置这些参数。

此外,我找到了这篇文章,但是当我自己尝试编辑xgb.plot.tree函数的源代码时,它给我一个错误,即找不到源中的“检查弃用”功能代码。

xgb.plot.tree layout in r

这是源代码

xgb.plot.tree <- function(feature_names = NULL, model = NULL, trees = NULL, 
plot_width = NULL, plot_height = NULL,
                      render = TRUE, show_node_id = FALSE, ...){
check.deprecation(...)
if (!inherits(model, "xgb.Booster")) {
stop("model: Has to be an object of class xgb.Booster")
}

if (!requireNamespace("DiagrammeR", quietly = TRUE)) {
stop("DiagrammeR package is required for xgb.plot.tree", call. = FALSE)
}

dt <- xgb.model.dt.tree(feature_names = feature_names, model = model, trees 
= trees)

dt[, label:= paste0(Feature, "\nCover: ", Cover, ifelse(Feature == "Leaf", 
"\nValue: ", "\nGain: "), Quality)]
if (show_node_id)
dt[, label := paste0(ID, ": ", label)]
dt[Node == 0, label := paste0("Tree ", Tree, "\n", label)]
dt[, shape:= "rectangle"][Feature == "Leaf", shape:= "oval"]
dt[, filledcolor:= "Beige"][Feature == "Leaf", filledcolor:= "Khaki"]
# in order to draw the first tree on top:
dt <- dt[order(-Tree)]

nodes <- DiagrammeR::create_node_df(
n         = nrow(dt),
ID        = dt$ID,
label     = dt$label,
fillcolor = dt$filledcolor,
shape     = dt$shape,
data      = dt$Feature,
fontcolor = "black")

edges <- DiagrammeR::create_edge_df(
from  = match(dt[Feature != "Leaf", c(ID)] %>% rep(2), dt$ID),
to    = match(dt[Feature != "Leaf", c(Yes, No)], dt$ID),
label = dt[Feature != "Leaf", paste("<", Split)] %>%
        c(rep("", nrow(dt[Feature != "Leaf"]))),
style = dt[Feature != "Leaf", ifelse(Missing == Yes, "bold", "solid")] %>%
        c(dt[Feature != "Leaf", ifelse(Missing == No, "bold", "solid")]),
rel   = "leading_to")

graph <- DiagrammeR::create_graph(
  nodes_df = nodes,
  edges_df = edges,
  attr_theme = NULL
  ) %>%
DiagrammeR::add_global_graph_attrs(
  attr_type = "graph",
  attr  = c("layout", "rankdir"),
  value = c("dot", "LR")
  ) %>%
DiagrammeR::add_global_graph_attrs(
  attr_type = "node",
  attr  = c("color", "style", "fontname"),
  value = c("DimGray", "filled", "Helvetica")
  ) %>%
DiagrammeR::add_global_graph_attrs(
  attr_type = "edge",
  attr  = c("color", "arrowsize", "arrowhead", "fontname"),
  value = c("DimGray", "1.5", "vee", "Helvetica"))

if (!render) return(invisible(graph))

DiagrammeR::render_graph(graph, width = plot_width, height = plot_height)
}

,我只是添加了黑色和灰色,而不是米色和卡其色,看是否有帮助。

xgb.plot.tree <- function(feature_names = NULL, model = NULL, trees = NULL, 
plot_width = NULL, plot_height = NULL,
                      render = TRUE, show_node_id = FALSE, ...){
check.deprecation(...)
if (!inherits(model, "xgb.Booster")) {
stop("model: Has to be an object of class xgb.Booster")
}

if (!requireNamespace("DiagrammeR", quietly = TRUE)) {
stop("DiagrammeR package is required for xgb.plot.tree", call. = FALSE)
}

dt <- xgb.model.dt.tree(feature_names = feature_names, model = model, trees 
= trees)

dt[, label:= paste0(Feature, "\nCover: ", Cover, ifelse(Feature == "Leaf", 
"\nValue: ", "\nGain: "), Quality)]
if (show_node_id)
dt[, label := paste0(ID, ": ", label)]
dt[Node == 0, label := paste0("Tree ", Tree, "\n", label)]
dt[, shape:= "rectangle"][Feature == "Leaf", shape:= "oval"]
dt[, filledcolor:= "Black"][Feature == "Leaf", filledcolor:= "Grey"]
# in order to draw the first tree on top:
dt <- dt[order(-Tree)]

nodes <- DiagrammeR::create_node_df(
n         = nrow(dt),
ID        = dt$ID,
label     = dt$label,
fillcolor = dt$filledcolor,
shape     = dt$shape,
data      = dt$Feature,
fontcolor = "black")

edges <- DiagrammeR::create_edge_df(
from  = match(dt[Feature != "Leaf", c(ID)] %>% rep(2), dt$ID),
to    = match(dt[Feature != "Leaf", c(Yes, No)], dt$ID),
label = dt[Feature != "Leaf", paste("<", Split)] %>%
        c(rep("", nrow(dt[Feature != "Leaf"]))),
style = dt[Feature != "Leaf", ifelse(Missing == Yes, "bold", "solid")] %>%
        c(dt[Feature != "Leaf", ifelse(Missing == No, "bold", "solid")]),
rel   = "leading_to")

graph <- DiagrammeR::create_graph(
  nodes_df = nodes,
  edges_df = edges,
  attr_theme = NULL
  ) %>%
DiagrammeR::add_global_graph_attrs(
  attr_type = "graph",
  attr  = c("layout", "rankdir"),
  value = c("dot", "LR")
  ) %>%
DiagrammeR::add_global_graph_attrs(
  attr_type = "node",
  attr  = c("color", "style", "fontname"),
  value = c("DimGray", "filled", "Helvetica")
  ) %>%
DiagrammeR::add_global_graph_attrs(
  attr_type = "edge",
  attr  = c("color", "arrowsize", "arrowhead", "fontname"),
  value = c("DimGray", "1.5", "vee", "Helvetica"))

if (!render) return(invisible(graph))

DiagrammeR::render_graph(graph, width = plot_width, height = plot_height)
}

任何帮助都将非常感谢!

0 个答案:

没有答案