lgb.Predictor.R

Predictor <- R6Class(
  classname = "lgb.Predictor",
  cloneable = FALSE,
  public = list(
    
    # Finalize will free up the handles
    finalize = function() {
      
      # Check the need for freeing handle
      if (private$need_free_handle && !lgb.is.null.handle(private$handle)) {
        
        # Freeing up handle
        lgb.call("LGBM_BoosterFree_R", ret = NULL, private$handle)
        private$handle <- NULL
        
      }
      
    },
    
    # Initialize will create a starter model
    initialize = function(modelfile) {
      
      # Create new lgb handle
      handle <- lgb.new.handle()
      
      # Check if handle is a character
      if (is.character(modelfile)) {
        
        # Create handle on it
        handle <- lgb.call("LGBM_BoosterCreateFromModelfile_R", ret = handle, lgb.c_str(modelfile))
        private$need_free_handle <- TRUE
        
      } else if (is(modelfile, "lgb.Booster.handle")) {
        
        # Check if model file is a booster handle already
        handle <- modelfile
        private$need_free_handle <- FALSE
        
      } else {
        
        # Model file is unknown
        stop("lgb.Predictor: modelfile must be either a character filename or an lgb.Booster.handle")
        
      }
      
      # Override class and store it
      class(handle) <- "lgb.Booster.handle"
      private$handle <- handle
      
    },
    
    # Get current iteration
    current_iter = function() {
      
      cur_iter <- 0L
      lgb.call("LGBM_BoosterGetCurrentIteration_R",  ret = cur_iter, private$handle)
      
    },
    
    # Predict from data
    predict = function(data,
                       num_iteration = NULL,
                       rawscore = FALSE,
                       predleaf = FALSE,
                       header = FALSE,
                       reshape = FALSE) {
      
      # Check if number of iterations is existing - if not, then set it to -1 (use all)
      if (is.null(num_iteration)) {
        num_iteration <- -1
      }
      
      # Set temporary variable
      num_row <- 0L
      
      # Check if data is a file name
      if (is.character(data)) {
        
        # Data is a filename, create a temporary file with a "lightgbm_" pattern in it
        tmp_filename <- tempfile(pattern = "lightgbm_")
        on.exit(unlink(tmp_filename), add = TRUE)
        
        # Predict from temporary file
        lgb.call("LGBM_BoosterPredictForFile_R", ret = NULL, private$handle, data,
          as.integer(header),
          as.integer(rawscore),
          as.integer(predleaf),
          as.integer(num_iteration),
          lgb.c_str(tmp_filename))
        
        # Get predictions from file
        preds <- read.delim(tmp_filename, header = FALSE, seq = "\t")
        num_row <- nrow(preds)
        preds <- as.vector(t(preds))
        
      } else {
        
        # Not a file, we need to predict from R object
        num_row <- nrow(data)
        npred <- 0L
        
        # Check number of predictions to do
        npred <- lgb.call("LGBM_BoosterCalcNumPredict_R",
                          ret = npred,
                          private$handle,
                          as.integer(num_row),
                          as.integer(rawscore),
                          as.integer(predleaf),
                          as.integer(num_iteration))
        
        # Pre-allocate empty vector
        preds <- numeric(npred)
        
        # Check if data is a matrix
        if (is.matrix(data)) {
          preds <- lgb.call("LGBM_BoosterPredictForMat_R",
                            ret = preds,
                            private$handle,
                            data,
                            as.integer(nrow(data)),
                            as.integer(ncol(data)),
                            as.integer(rawscore),
                            as.integer(predleaf),
                            as.integer(num_iteration))
          
        } else if (is(data, "dgCMatrix")) {
          
          # Check if data is a dgCMatrix (sparse matrix, column compressed format)
          preds <- lgb.call("LGBM_BoosterPredictForCSC_R",
                            ret = preds,
                            private$handle,
                            data@p,
                            data@i,
                            data@x,
                            length(data@p),
                            length(data@x),
                            nrow(data),
                            as.integer(rawscore),
                            as.integer(predleaf),
                            as.integer(num_iteration))
          
        } else {
          
          # Cannot predict on unknown class
          # to-do: predict from lgb.Dataset
          stop("predict: cannot predict on data of class ", sQuote(class(data)))
          
        }
      }
      
      # Check if number of rows is strange (not a multiple of the dataset rows)
      if (length(preds) %% num_row != 0) {
        stop("predict: prediction length ", sQuote(length(preds))," is not a multiple of nrows(data): ", sQuote(num_row))
      }
      
      # Get number of cases per row
      npred_per_case <- length(preds) / num_row
      
      
      # Data reshaping
      
      if (predleaf) {
        
        # Predict leaves only, reshaping is mandatory
        preds <- matrix(preds, ncol = npred_per_case, byrow = TRUE)
        
      } else if (reshape && npred_per_case > 1) {
        
        # Predict with data reshaping
        preds <- matrix(preds, ncol = npred_per_case, byrow = TRUE)
        
      }
      
      # Return predictions
      return(preds)
      
    }
    
  ),
  private = list(handle = NULL,
                 need_free_handle = FALSE)
)