add NA's to confusion matrix if present

hiredd · Jul 17, 2017 · 9997ed1 · 9997ed1
1 parent 621f23f
commit 9997ed1
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/R/ranger.R b/R/ranger.R
@@ -637,7 +637,8 @@ ranger <- function(formula = NULL, data = NULL, num.trees = 500, mtry = NULL,
                                             levels(response))
     true.values <- integer.to.factor(unlist(data.final[, dependent.variable.name]),
                                      levels(response))
-    result$confusion.matrix <- table(true.values, result$predictions, dnn = c("true", "predicted"))
+    result$confusion.matrix <- table(true.values, result$predictions, 
+                                     dnn = c("true", "predicted"), useNA = "ifany")
   } else if (treetype == 5) {
     if (is.list(result$predictions)) {
       result$predictions <- do.call(rbind, result$predictions)

diff --git a/tests/testthat/test_classification.R b/tests/testthat/test_classification.R
@@ -126,6 +126,12 @@ test_that("confusion matrix rows are the true classes", {
                as.numeric(table(iris$Species)))
 })
 
+test_that("confusion matrix rows are the true classes if using case weights", {
+  rf <- ranger(Species ~ ., data = iris, num.trees = 5, 
+               case.weights = c(rep(100, 5), rep(5, 145)))
+  expect_equal(as.numeric(rowSums(rf$confusion.matrix)), 
+               as.numeric(table(iris$Species)))
+})
 ## Splitrule
 test_that("default splitrule is Gini for classification", {
   set.seed(42)