Fixed a bug with display of confusion matrix

panzelon · Aug 30, 2013 · ac05e20 · ac05e20
1 parent 7d50308
commit ac05e20
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 14 deletions.
diff --git a/R/h2o-package/R/Algorithms.R b/R/h2o-package/R/Algorithms.R
@@ -38,11 +38,13 @@ setMethod("h2o.glm", signature(x="character", y="character", data="H2OParsedData
               result$threshold = res$validations[[1]]$threshold
 
               # Build confusion matrix
-              temp = res$validations[[1]]$cm
-              temp[[1]] = NULL; temp = lapply(temp, function(x) { x[[1]] = NULL; x })
-              result$confusion = t(matrix(unlist(temp), nrow = length(temp)))
-              myNames = c("False", "True", "Error")
-              dimnames(result$confusion) = list(Actual = myNames, Predicted = myNames)
+              if(family == "binomial") {
+                temp = res$validations[[1]]$cm
+                temp[[1]] = NULL; temp = lapply(temp, function(x) { x[[1]] = NULL; x })
+                result$confusion = t(matrix(unlist(temp), nrow = length(temp)))
+                myNames = c("False", "True", "Error")
+                dimnames(result$confusion) = list(Actual = myNames, Predicted = myNames)
+              }
               result
             }
 

diff --git a/R/h2o-package/R/Classes.R b/R/h2o-package/R/Classes.R
@@ -35,17 +35,25 @@ setMethod("show", "H2OGLMModel", function(object) {
   cat("\nNull Deviance:    ", round(model$null.deviance,1))
   cat("\nResidual Deviance:", round(model$deviance,1), " AIC:", round(model$aic,1))
 
-  cat("\n\nAvg Training Error:", model$training.err)
-  cat("\nBest Threshold:", model$threshold, " AUC:", model$auc)
-  cat("\n\nConfusion Matrix:\n"); print(model$confusion)
+  if(model$family == "binomial") {
+    cat("\n\nAvg Training Error:", model$training.err)
+    cat("\nBest Threshold:", model$threshold, " AUC:", model$auc)
+    cat("\n\nConfusion Matrix:\n"); print(model$confusion)
+  }
 
   if(length(object@xval) > 0) {
-    cat("\nCross-Validation:\n")
-    xval = sapply(object@xval, function(x) {
-      c(x@model$threshold, x@model$auc, x@model$confusion[1,3], x@model$confusion[2,3])
-    })
-    xval = t(xval)
-    colnames(xval) = c("Best Threshold", "AUC", "Err(0)", "Err(1)")
+    if(model$family == "binomial") {
+      cat("\nCross-Validation:\n")
+      xval = sapply(object@xval, function(x) {
+        c(x@model$threshold, x@model$auc, x@model$confusion[1,3], x@model$confusion[2,3])
+      })
+      xval = t(xval)
+      colnames(xval) = c("Best Threshold", "AUC", "Err(0)", "Err(1)")
+    } else {
+      cat("\n\nCross-Validation:\n")
+      xval = data.frame(sapply(object@xval, function(x) { x@model$training.err }))
+      colnames(xval) = "Error"
+    }
     rownames(xval) = paste("Model", 1:nrow(xval))
     print(xval)
   }