R/predict.srlars.R

In srlars: Fast and Scalable Cellwise-Robust Ensemble

#' @title Predictions for srlars Object
#'
#' @description \code{predict.srlars} returns the predictions for a srlars object.
#'
#' @method predict srlars
#'
#' @param object An object of class srlars.
#' @param newx New data matrix for predictions.
#' @param model_index Indices of the sub-models to include in the ensemble. Default is NULL (all models).
#' @param dynamic Logical. If TRUE, and the model was trained robustly, the new data \code{newx} is cleaned using
#'        \code{\link[cellWise]{DDCpredict}} before prediction. This ensures consistency with the robust training phase.
#'        Default is TRUE.
#' @param ... Additional arguments for compatibility.
#'
#' @return A numeric vector of predictions.
#'
#' @export
#'
#' @author Anthony-Alexander Christidis, \email{anthony.christidis@stat.ubc.ca}
#'
#' @seealso \code{\link{srlars}}
#'
#' @importFrom cellWise DDCpredict
#'
#' @examples
#' # Required libraries
#' library(mvnfast)
#' library(cellWise)
#' library(robustbase)
#'
#' # Simulation parameters
#' n <- 50
#' p <- 100
#' rho.within <- 0.8
#' rho.between <- 0.2
#' p.active <- 20
#' group.size <- 5
#' snr <- 3
#' contamination.prop <- 0.1
#'
#' # Setting the seed
#' set.seed(0)
#'
#' # Block correlation structure
#' sigma.mat <- matrix(0, p, p)
#' sigma.mat[1:p.active, 1:p.active] <- rho.between
#' for(group in 0:(p.active/group.size - 1))
#'   sigma.mat[(group*group.size+1):(group*group.size+group.size),
#'   (group*group.size+1):(group*group.size+group.size)] <- rho.within
#' diag(sigma.mat) <- 1
#'
#' # Simulation of beta vector
#' true.beta <- c(runif(p.active, 0, 5)*(-1)^rbinom(p.active, 1, 0.7), rep(0, p - p.active))
#'
#' # Setting the SD of the variance
#' sigma <- as.numeric(sqrt(t(true.beta) %*% sigma.mat %*% true.beta)/sqrt(snr))
#'
#' # Simulation of uncontaminated data
#' x <- mvnfast::rmvn(n, mu = rep(0, p), sigma = sigma.mat)
#' colnames(x) <- paste0("V", 1:p)
#' y <- x %*% true.beta + rnorm(n, 0, sigma)
#'
#' # Cellwise contamination
#' contamination_indices <- sample(1:(n * p), round(n * p * contamination.prop))
#' x_train <- x
#' x_train[contamination_indices] <- runif(length(contamination_indices), -10, 10)
#'
#' # FSCRE Ensemble model
#' ensemble_fit <- srlars(x_train, y,
#'                        n_models = 5,
#'                        tolerance = 1e-4,
#'                        x_preprocess = "ddc",
#'                        y_preprocess = "wrap",
#'                        cor_estimator = "wrap",
#'                        cv_preprocess = "global",
#'                        cv_fit = "ls",
#'                        cv_loss = "huber",
#'                        compute_coef = TRUE)
#'
#' # Generate Test Data
#' x_test <- mvnfast::rmvn(50, mu = rep(0, p), sigma = sigma.mat)
#' colnames(x_test) <- paste0("V", 1:p)
#' y_test <- x_test %*% true.beta + rnorm(50, 0, sigma)
#'
#' # Predict on Test Data
#' preds <- predict(ensemble_fit, x_test)
#' 
#' # Calculate MSPE
#' mspe <- mean((y_test - preds)^2)
#' print(paste("MSPE:", mspe))
#'
predict.srlars <- function(object,
                           newx,
                           model_index = NULL,
                           dynamic = TRUE,
                           ...) {

    # 1. Validate Inputs
    newx <- as.matrix(newx)

    if (is.null(colnames(newx))) {
        colnames(newx) <- paste0("V", 1:ncol(newx))
    }

    if(is.null(model_index)){
        model_index <- 1:object$n_models
    } else{
        if(any(!(model_index %in% 1:object$n_models)))
            stop("The model_index contains invalid indices.")
    }

    # 2. Dynamic Imputation (Robust Prediction)
    x_for_pred <- newx

    # Check if we should (and can) perform robust cleaning
    if (dynamic && isTRUE(object$robust) && !is.null(object$ddc.object)) {

        # object$ddc.object was fit on the predictors alone (computeRobustFoundation runs DDC
        # on X only, separately from y, to avoid target leakage), so newx is passed as-is.
        ddc_pred <- tryCatch({
            cellWise::DDCpredict(Xnew = newx, InitialDDC = object$ddc.object)
        }, error = function(e) {
            warning(paste("DDCpredict failed:", e$message, "Falling back to raw newx."))
            return(NULL)
        })

        if (!is.null(ddc_pred)) {
            x_for_pred <- ddc_pred$Ximp
        }
    }

    # 3. Compute Predictions
    n_test <- nrow(x_for_pred)
    final_preds <- numeric(n_test)
    n_groups <- length(model_index)

    for (k in model_index) {
        # Extract params
        beta_k <- object$coefficients[[k]]
        intercept_k <- object$intercepts[k]

        # Linear predictor: alpha + X * beta
        preds_k <- intercept_k + (x_for_pred %*% beta_k)

        # Accumulate
        final_preds <- final_preds + preds_k
    }

    # Average
    if (n_groups > 0) {
        final_preds <- final_preds / n_groups
    }

    return(as.numeric(final_preds))
}