diff --git a/R-package/R/xgb.plot.multi.tree.R b/R-package/R/xgb.plot.multi.tree.R
index 314e2157a..f61540dae 100644
--- a/R-package/R/xgb.plot.multi.tree.R
+++ b/R-package/R/xgb.plot.multi.tree.R
@@ -1,6 +1,6 @@
library(stringr)
library(data.table)
-
+library(xgboost)
data(agaricus.train, package='xgboost')
@@ -10,8 +10,8 @@ data(agaricus.train, package='xgboost')
#Each column of the sparse Matrix is a feature in one hot encoding format.
train <- agaricus.train
-bst <- xgboost(data = train$data, label = train$label, max.depth = 5,
- eta = 1, nthread = 2, nround = 2,objective = "binary:logistic")
+bst <- xgboost(data = train$data, label = train$label, max.depth = 3,
+ eta = 1, nthread = 2, nround = 4,objective = "binary:logistic")
#agaricus.test$data@Dimnames[[2]] represents the column names of the sparse matrix.
tree.matrix <- xgb.model.dt.tree(agaricus.train$data@Dimnames[[2]], model = bst)
@@ -21,22 +21,39 @@ tree.matrix <- xgb.model.dt.tree(agaricus.train$data@Dimnames[[2]], model = bst)
# root init
root.nodes <- tree.matrix[str_detect(ID, "\\d+-0"), ID]
-tree.matrix[ID == root.nodes, Abs.Position:=root.nodes]
+tree.matrix[ID == root.nodes, Abs.Position:=root.nodes %>% str_replace("-", "_")]
-precedent.nodes <- root.nodes
+precedent.nodes <- root.nodes %>% str_replace("-", "_")
while(tree.matrix[,sum(is.na(Abs.Position))] > 0) {
yes.row.nodes <- tree.matrix[Abs.Position %in% precedent.nodes & !is.na(Yes)]
no.row.nodes <- tree.matrix[Abs.Position %in% precedent.nodes & !is.na(No)]
- yes.nodes.abs.pos <- yes.row.nodes[, Abs.Position] %>% paste0("-0")
- no.nodes.abs.pos <- no.row.nodes[, Abs.Position] %>% paste0("-1")
+ yes.nodes.abs.pos <- yes.row.nodes[, Abs.Position] %>% paste0("_0")
+ no.nodes.abs.pos <- no.row.nodes[, Abs.Position] %>% paste0("_1")
tree.matrix[ID == yes.row.nodes[, Yes], Abs.Position := yes.nodes.abs.pos]
tree.matrix[ID == no.row.nodes[, No], Abs.Position := no.nodes.abs.pos]
precedent.nodes <- c(yes.nodes.abs.pos, no.nodes.abs.pos)
}
-tree.matrix
+tree.matrix[!is.na(Yes),Yes:= paste0(Abs.Position, "_0")]
+tree.matrix[!is.na(No),No:= paste0(Abs.Position, "_1")]
+tree.matrix[,ID:= Abs.Position]
+tree.matrix[Feature!="Leaf" ,yesPath:= paste(ID,"(", Feature, "
Cover: ", Cover, "
Gain: ", Quality, ")-->|< ", Split, "|", Yes, ">", Yes.Feature, "]", sep = "")]
+tree.matrix[Feature!="Leaf" ,noPath:= paste(ID,"(", Feature, ")-->|>= ", Split, "|", No, ">", No.Feature, "]", sep = "")]
+
+CSSstyle <- "classDef greenNode fill:#A2EB86, stroke:#04C4AB, stroke-width:2px\nclassDef redNode fill:#FFA070, stroke:#FF5E5E, stroke-width:2px"
+
+
+yes <- tree.matrix[Feature!="Leaf", c(Yes)] %>% paste(collapse = ",") %>% paste("class ", ., " greenNode", sep = "")
+
+no <- tree.matrix[Feature!="Leaf", c(No)] %>% paste(collapse = ",") %>% paste("class ", ., " redNode", sep = "")
+
+path <- tree.matrix[Feature!="Leaf", c(yesPath, noPath)] %>% .[order(.)] %>% paste(sep = "", collapse = "\n") %>% paste("graph LR", .,collapse = "", sep = "\n") %>% paste(CSSstyle, yes, no, sep = "\n")
+DiagrammeR::mermaid(path)
+
+# path <- "graph LR;0-0-0(spore-print-color=green)-->|>= 2.00001|0-0-0-1>Leaf"
+setnames(tree.matrix, old = c("ID", "Yes", "No"), c("nodes", "edge_from", "edge_to"))