Automatically predict NA for rows w/ NAs and learners that don't support missings

R/generateHyperParsEffect.R

@@ -362,7 +362,7 @@ plotHyperParsEffect = function(hyperpars.effect.data, x = NULL, y = NULL,
           regr.task = makeRegrTask(id = "interp", data = d.run[, c(x, y, z)],
             target = z)
           mod = train(lrn, regr.task)
-          prediction = predict(mod, newdata = grid)
+          prediction = predict(mod, newdata = grid[c(x, y)])
           grid[, z] = prediction$data[, prediction$predict.type]
           grid$learner_status = "Interpolated Point"
           grid$iteration = NA

R/predictLearner.R

@@ -58,9 +58,37 @@ predictLearner2 = function(.learner, .model, .newdata, ...) {
       .newdata[ns] = mapply(factor, x = .newdata[ns],
          levels = fls, SIMPLIFY = FALSE)
   }
-  p = predictLearner(.learner, .model, .newdata, ...)
+  if ("missings" %nin% getLearnerProperties(.learner))
+    no.na = removeNALines(.newdata)
+  else
+    no.na = list(newdata = .newdata, inserts = FALSE)
+  if (!nrow(no.na$newdata))
+    no.na = list(newdata = .newdata, inserts = FALSE)  # no choice if all lines contain NA
+  p = predictLearner(.learner, .model, no.na$newdata, ...)
   p = checkPredictLearnerOutput(.learner, .model, p)
-  return(p)
+  return(insertLines(p, no.na$inserts))
+}
+
+removeNALines = function(newdata) {
+  namat = is.na(newdata)
+  narows = apply(namat, 1, any)
+  return(list(newdata = newdata[!narows, , drop = FALSE], inserts = narows))
+}
+
+insertLines = function(prediction, inserts) {
+#  if (!any(inserts))
+#    return(prediction)
+  if (is.matrix(prediction)) {
+    ret = matrix(nrow = nrow(prediction) + sum(inserts), ncol = ncol(prediction))
+    ret[!inserts, ] = prediction
+    colnames(ret) = colnames(prediction)
+  } else {
+    ret = rep(NA, length(prediction) + sum(inserts))
+    ret[!inserts] = prediction
+    attributes(ret) = attributes(prediction)
+    names(ret) = NULL
+  }
+  return(ret)
 }
 
 checkPredictLearnerOutput = function(learner, model, p) {

tests/testthat/test_base_predict.R

@@ -144,3 +144,17 @@ test_that("predict works with data.table as newdata", {
   expect_warning(predict(mod, newdata = data.table(iris)), regexp = "Provided data for prediction is not a pure data.frame but from class data.table, hence it will be converted.")
 })
 
+test_that("predict with NA rows for learners that don't support missings automatically returns NA", {
+  modknn = train("classif.knn", pid.task)
+  modrf = train(makeLearner("classif.randomForest", mtry = 1), pid.task)
+  newdata = getTaskData(pid.task, target.extra = TRUE)$data
+  newdata.na = newdata
+  newdata.na[[1]][1] = NA
+  for (mod in list(modknn, modrf)) {
+    prediction = predict(mod, newdata = newdata)
+    prediction.na = predict(mod, newdata = newdata.na)
+    expect_equal(which(is.na(prediction.na$data$response[1])), 1)
+    expect_equal(prediction.na$data[-1, ], prediction$data[-1, ])
+  }
+})
+