Some fixes + refactoring

[agghoo.git] / R / R6_Model.R

diff --git a/R/R6_Model.R b/R/R6_Model.R
index 8912cdb..05cb7d8 100644 (file)
--- a/R/R6_Model.R
+++ b/R/R6_Model.R
@@ -40,7 +40,7 @@ Model <- R6::R6Class("Model",
       if (is.null(params))
         # Here, gmodel is a string (= its family),
         # because a custom model must be given with its parameters.
-        params <- as.list(private$getParams(gmodel, data, target))
+        params <- as.list(private$getParams(gmodel, data, target, task))
       private$params <- params
       if (is.character(gmodel))
         gmodel <- private$getGmodel(gmodel, task)
@@ -49,12 +49,17 @@ Model <- R6::R6Class("Model",
     },
     #' @description
     #' Returns the model at index "index", trained on dataHO/targetHO.
-    #' index is between 1 and self$nmodels.
     #' @param dataHO Matrix or data.frame
     #' @param targetHO Vector of targets (generally numeric or factor)
     #' @param index Index of the model in 1...nmodels
     get = function(dataHO, targetHO, index) {
       private$gmodel(dataHO, targetHO, private$params[[index]])
+    },
+    #' @description
+    #' Returns the parameter at index "index".
+    #' @param index Index of the model in 1...nmodels
+    getParam = function(index) {
+      private$params[[index]]
     }
   ),
   private = list(
@@ -68,9 +73,23 @@ Model <- R6::R6Class("Model",
         function(dataHO, targetHO, param) {
           require(rpart)
           method <- ifelse(task == "classification", "class", "anova")
+          if (is.null(colnames(dataHO)))
+            colnames(dataHO) <- paste0("V", 1:ncol(dataHO))
           df <- data.frame(cbind(dataHO, target=targetHO))
           model <- rpart::rpart(target ~ ., df, method=method, control=list(cp=param))
-          function(X) predict(model, X)
+          if (task == "regression")
+            type <- "vector"
+          else {
+            if (is.null(dim(targetHO)))
+              type <- "class"
+            else
+              type <- "prob"
+          }
+          function(X) {
+            if (is.null(colnames(X)))
+              colnames(X) <- paste0("V", 1:ncol(X))
+            predict(model, as.data.frame(X), type=type)
+          }
         }
       }
       else if (family == "rf") {
@@ -104,22 +123,21 @@ Model <- R6::R6Class("Model",
       }
     },
     # Return a default list of parameters, given a gmodel family
-    getParams = function(family, data, target) {
+    getParams = function(family, data, target, task) {
       if (family == "tree") {
         # Run rpart once to obtain a CV grid for parameter cp
         require(rpart)
         df <- data.frame(cbind(data, target=target))
         ctrl <- list(
+          cp = 0,
           minsplit = 2,
           minbucket = 1,
-          maxcompete = 0,
-          maxsurrogate = 0,
-          usesurrogate = 0,
-          xval = 0,
-          surrogatestyle = 0,
-          maxdepth = 30)
-        r <- rpart(target ~ ., df, method="class", control=ctrl)
+          xval = 0)
+        method <- ifelse(task == "classification", "class", "anova")
+        r <- rpart(target ~ ., df, method=method, control=ctrl)
         cps <- r$cptable[-1,1]
+        if (length(cps) <= 1)
+          stop("No cross-validation possible: select another model")
         if (length(cps) <= 11)
           return (cps)
         step <- (length(cps) - 1) / 10
@@ -129,7 +147,7 @@ Model <- R6::R6Class("Model",
         p <- ncol(data)
         # Use caret package to obtain the CV grid of mtry values
         require(caret)
-        caret::var_seq(p, classification = (task == "classificaton"),
+        caret::var_seq(p, classification = (task == "classification"),
                        len = min(10, p-1))
       }
       else if (family == "ppr")