make predict an S3 method

catboost/R-package/R/catboost.R

@@ -1708,16 +1708,19 @@ catboost.save_model <- function(model, model_path,
 
 
 #' @name catboost.predict
-#' @title Apply the model
+#' @title Get predictions from a CatBoost model
 #'
-#' @description Apply the model to the given dataset.
+#' @description Get predictions from a CatBoost model on new data.
+#' @details The function `catboost.predict` is a synonym for `predict.catboost.Model`,
+#' which is an S3 method (i.e. called like `predict(model, newdata)`).
 #'
-#'              Peculiarities: In case of multiclassification the prediction is returned in the form of a matrix.
-#'              Each line of this matrix contains the predictions for one object of the input dataset.
-#' @param model The model obtained as the result of training.
+#' In case of multiclassification the prediction is returned in the form of a matrix.
+#' Each row of this matrix contains the predictions for one row of the input dataset.
+#' @param object The model obtained as the result of training.
 #'
 #' Default value: Required argument
-#' @param pool The input dataset.
+#' @param newdata The input data on which to make predictions. Should be a `catboost.Pool`
+#' object.
 #'
 #' Default value: Required argument
 #' @param verbose Verbose output to stdout.
@@ -1754,26 +1757,41 @@ catboost.save_model <- function(model, model_path,
 #' @return Vector of predictions (matrix for multi-class classification).
 #' @export
 #' @seealso \url{https://catboost.ai/docs/concepts/r-reference_catboost-predict.html}
-catboost.predict <- function(model, pool,
-                             verbose = FALSE, prediction_type = "RawFormulaVal",
-                             ntree_start = 0, ntree_end = 0, thread_count = -1) {
-    if (!inherits(model, "catboost.Model"))
-        stop("Expected catboost.Model, got: ", class(model))
-    if (!inherits(pool, "catboost.Pool"))
-        stop("Expected catboost.Pool, got: ", class(pool))
-    if (is.null.handle(pool))
+predict.catboost.Model <- function(object, newdata,
+                                   verbose = FALSE, prediction_type = "RawFormulaVal",
+                                   ntree_start = 0, ntree_end = 0, thread_count = -1) {
+    if (!inherits(object, "catboost.Model"))
+        stop("Expected catboost.Model, got: ", class(object))
+    if (!inherits(newdata, "catboost.Pool"))
+        stop("Expected catboost.Pool, got: ", class(newdata))
+    if (is.null.handle(newdata))
         stop("Pool object is invalid.")
 
-    if (is.null.handle(model$handle))
-        model$handle <- .Call("CatBoostDeserializeModel_R", model$raw)
-    prediction <- .Call("CatBoostPredictMulti_R", model$handle, pool,
+    if (is.null.handle(object$handle))
+        object$handle <- .Call("CatBoostDeserializeModel_R", object$raw)
+    prediction <- .Call("CatBoostPredictMulti_R", object$handle, newdata,
                         verbose, prediction_type, ntree_start, ntree_end, thread_count)
-    if (length(prediction) != nrow(pool)) {
-        prediction <- matrix(prediction, nrow = nrow(pool), byrow = TRUE)
+    if (length(prediction) != nrow(newdata)) {
+        prediction <- matrix(prediction, nrow = nrow(newdata), byrow = TRUE)
     }
     return(prediction)
 }
 
+#' @rdname catboost.predict
+#' @export
+#' @usage catboost.predict(
+#'   object,
+#'   newdata,
+#'   verbose = FALSE,
+#'   prediction_type = "RawFormulaVal",
+#'   ntree_start = 0,
+#'   ntree_end = 0,
+#'   thread_count = -1
+#' )
+catboost.predict <- function(...) {
+    return(predict.catboost.Model(...))
+}
+
 
 #' @name catboost.staged_predict
 #' @title Apply the model for each tree