Merge pull request #569 from ilia-kats/xgb_varimp

return full list of variables in xgb*$varImp
topepo · Jan 9, 2017 · 854d128 · 854d128
2 parents bf5d8c8 + 91d7a83
commit 854d128
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/models/files/xgbLinear.R b/models/files/xgbLinear.R
@@ -94,6 +94,12 @@ modelInfo <- list(label = "eXtreme Gradient Boosting",
                     rownames(imp) <- as.character(imp[,1])
                     imp <- imp[,2,drop = FALSE]
                     colnames(imp) <- "Overall"
+
+                    missing <- object$xNames[!(object$xNames %in% rownames(imp))]
+                    missing_imp <- data.frame(Overall=rep(0, times=length(missing)))
+                    rownames(missing_imp) <- missing
+                    imp <- rbind(imp, missing_imp)
+
                     imp   
                   },
                   levels = function(x) x$obsLevels,

diff --git a/models/files/xgbTree.R b/models/files/xgbTree.R
@@ -184,6 +184,12 @@ modelInfo <- list(label = "eXtreme Gradient Boosting",
                     rownames(imp) <- as.character(imp[,1])
                     imp <- imp[,2,drop = FALSE]
                     colnames(imp) <- "Overall"
+
+                    missing <- object$xNames[!(object$xNames %in% rownames(imp))]
+                    missing_imp <- data.frame(Overall=rep(0, times=length(missing)))
+                    rownames(missing_imp) <- missing
+                    imp <- rbind(imp, missing_imp)
+
                     imp
                   },
                   levels = function(x) x$obsLevels,