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

Fix label errors in graph visualization (#6369)

Browse Source
This commit is contained in:
Jiaming Yuan 2020-11-12 09:44:59 +08:00 committed by GitHub
parent debeae2509
commit d711d648cb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 13 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
src/tree/tree_model.cc
View File

@ -552,20 +552,23 @@ class GraphvizGenerator : public TreeGenerator {
{"{params}", param_.condition_node_params}}); {"{params}", param_.condition_node_params}});
static std::string const kEdgeTemplate = static std::string const kEdgeTemplate =
" {nid} -> {child} [label=\"{is_missing}\" color=\"{color}\"]\n"; " {nid} -> {child} [label=\"{branch}\" color=\"{color}\"]\n";
auto MatchFn = SuperT::Match; // mingw failed to capture protected fn. auto MatchFn = SuperT::Match; // mingw failed to capture protected fn.
auto BuildEdge = auto BuildEdge =
[&tree, nid, MatchFn, this](int32_t child) { [&tree, nid, MatchFn, this](int32_t child, bool left) {
// Is this the default child for missing value?
bool is_missing = tree[nid].DefaultChild() == child; bool is_missing = tree[nid].DefaultChild() == child;
std::string branch = std::string {left ? "yes" : "no"} +
std::string {is_missing ? ", missing" : ""};
std::string buffer = MatchFn(kEdgeTemplate, { std::string buffer = MatchFn(kEdgeTemplate, {
{"{nid}", std::to_string(nid)}, {"{nid}", std::to_string(nid)},
{"{child}", std::to_string(child)}, {"{child}", std::to_string(child)},
{"{color}", is_missing ? param_.yes_color : param_.no_color}, {"{color}", is_missing ? param_.yes_color : param_.no_color},
{"{is_missing}", is_missing ? "yes, missing": "no"}}); {"{branch}", branch}});
return buffer; return buffer;
}; };
result += BuildEdge(tree[nid].LeftChild()); result += BuildEdge(tree[nid].LeftChild(), true);
result += BuildEdge(tree[nid].RightChild()); result += BuildEdge(tree[nid].RightChild(), false);
return result; return result;
}; };

5
tests/cpp/tree/test_tree_model.cc
View File

@ -343,6 +343,11 @@ TEST(Tree, DumpDot) {
str = tree.DumpModel(fmap, true, R"(dot:{"graph_attrs": {"bgcolor": "#FFFF00"}})"); str = tree.DumpModel(fmap, true, R"(dot:{"graph_attrs": {"bgcolor": "#FFFF00"}})");
ASSERT_NE(str.find(R"(graph [ bgcolor="#FFFF00" ])"), std::string::npos); ASSERT_NE(str.find(R"(graph [ bgcolor="#FFFF00" ])"), std::string::npos);
// Default left for root.
ASSERT_NE(str.find(R"(0 -> 1 [label="yes, missing")"), std::string::npos);
// Default right for node 1
ASSERT_NE(str.find(R"(1 -> 4 [label="no, missing")"), std::string::npos);
} }
TEST(Tree, JsonIO) { TEST(Tree, JsonIO) {