Parse splits on inf in xgb.model.dt.tree (#3900)

dmlc · Mar 2, 2021 · a2b61f9 · a2b61f9
1 parent a9b4a95
commit a2b61f9
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/R-package/R/xgb.model.dt.tree.R b/R-package/R/xgb.model.dt.tree.R
@@ -95,7 +95,7 @@ xgb.model.dt.tree <- function(feature_names = NULL, model = NULL, text = NULL,
 
   add.tree.id <- function(node, tree) if (use_int_id) node else paste(tree, node, sep = "-")
 
-  anynumber_regex <- "[-+]?[0-9]*\\.?[0-9]+([eE][-+]?[0-9]+)?"
+  anynumber_regex <- "[-+]?[0-9]*\\.?[0-9]+([eE][-+]?[0-9]+)?|[-+]?[Ii]nf"
 
   td <- data.table(t = text)
   td[position, Tree := 1L]

diff --git a/R-package/tests/testthat/test_helpers.R b/R-package/tests/testthat/test_helpers.R
@@ -281,6 +281,18 @@ test_that("xgb.model.dt.tree throws error for gblinear", {
   expect_error(xgb.model.dt.tree(model = bst.GLM))
 })
 
+test_that("xgb.model.dt.tree works with splits on inf", {
+  tree.dump <- c(
+    "booster[0]",
+    "0:[f1<inf] yes=1,no=3,missing=3,gain=0.1,cover=3",
+    "1:[f2<-inf] yes=2,no=3,missing=3,gain=0.1,cover=1",
+    "2:[f1<2] yes=4,no=3,missing=3,gain=0.3,cover=4",
+    "3:leaf=0.2,cover=1",
+    "4:leaf=0.5,cover=1"
+  )
+  expect_equal(xgb.model.dt.tree(text = tree.dump)$Split, c(Inf, -Inf, 2, NA, NA))
+})
+
 test_that("xgb.importance works with and without feature names", {
   importance.Tree <- xgb.importance(feature_names = feature.names, model = bst.Tree)
   if (!flag_32bit)