xgboostImpute fix for binary variables

statistikat · Nov 9, 2023 · 225f064 · 225f064
1 parent 6f05df0
commit 225f064
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/R/xgboostImpute.R b/R/xgboostImpute.R
@@ -96,7 +96,11 @@ xgboostImpute <- function(formula, data, imp_var = TRUE,
       predictions <- 
         predict(mod, model.matrix(formPred,subset(data, !rhs_na & lhs_na)))
       if(currentClass=="factor"){
-        data[!rhs_na & lhs_na, lhsV] <- levels(dattmp[,lhsV])[predictions+1]
+        if(is.null(num_class)){
+          data[!rhs_na & lhs_na, lhsV] <- levels(dattmp[,lhsV])[as.numeric(predictions>.5)+1]  
+        }else{
+          data[!rhs_na & lhs_na, lhsV] <- levels(dattmp[,lhsV])[predictions+1]  
+        }
       }else{
         data[!rhs_na & lhs_na, lhsV] <- predictions  
       }