R/imputeCellReg.R

In VIM: Visualization and Imputation of Missing Values

#' Cellwise-robust regression imputation for mixed data
#'
#' IRMI-style imputation using cellwise-robust regression as the inner
#' engine.  Three engines are available: CRM (Filzmoser et al. 2020),
#' Shooting S (Öllerer et al. 2016), and a cellwise-weighted MM hybrid.
#'
#' @param data data.frame with missing values (mixed continuous + categorical)
#' @param engine regression engine: \code{"crm"} (default),
#'   \code{"cellwise-mm"}, or \code{"shooting-s"}
#' @param maxit maximum outer IRMI iterations (default: 50)
#' @param eps convergence tolerance (default: 5e-3)
#' @param uncert imputation uncertainty: \code{"pmm"} (default),
#'   \code{"normalerror"}, or \code{"none"}
#' @param trace logical; print progress
#'
#' @return A list with components:
#'   \item{data_imputed}{the imputed data.frame}
#'   \item{cellweights}{n x p matrix of cell weights (1 = clean)}
#'   \item{converged}{logical}
#'   \item{iterations}{number of outer iterations}
#'
#' @details
#' The function cycles through all variables with missing values
#' (IRMI framework), fitting a cellwise-robust regression of each
#' variable on all others.  The \code{engine} argument selects the
#' regression method:
#'
#' \describe{
#'   \item{\code{"crm"}}{CRM (Cellwise Robust M-regression) from the
#'     \code{crmReg} package.  Uses SPADIMO for cellwise outlier
#'     detection within each regression.  Requires \code{crmReg}.}
#'   \item{\code{"cellwise-mm"}}{Hybrid: compute cell weights via MCD
#'     conditional residuals, then fit MM-estimation (\code{lmrob})
#'     with row weights derived from cell weights.  The MM-estimator
#'     provides high breakdown point.}
#'   \item{\code{"shooting-s"}}{Shooting S-estimator (Öllerer et al.
#'     2016).  Iterates between cellwise detection and S-estimation.
#'     Implemented from the published algorithm.}
#' }
#'
#' Categorical variables are imputed via weighted multinomial logistic
#' regression, with row weights derived from the continuous cell weights.
#'
#' @references
#' P. Filzmoser, S. Höppner, I. Ortner, S. Serneels, S. Van Aelst (2020)
#' Cellwise robust M regression.  \emph{Computational Statistics and Data
#' Analysis}, 147, 106944.
#'
#' V. Öllerer, A. Alfons, C. Croux (2016) The shooting S-estimator for
#' robust regression.  \emph{Computational Statistics}, 31(3), 829--844.
#'
#' @author Matthias Templ
#' @importFrom stats as.formula predict median mad rnorm
#' @export
imputeCellReg <- function(data, engine = "crm",
                           maxit = 50, eps = 5e-3,
                           uncert = "pmm", trace = FALSE) {

  ## ---- input validation ----
  check_data(data)
  if (!is.data.frame(data)) {
    if (is.matrix(data)) data <- as.data.frame(data)
    else stop("data must be a data.frame or matrix")
  }
  engine <- match.arg(engine, c("crm", "cellwise-mm", "shooting-s"))
  uncert <- match.arg(uncert, c("pmm", "normalerror", "none"))

  if (engine == "crm" && !requireNamespace("crmReg", quietly = TRUE))
    stop("Engine 'crm' requires the crmReg package. Install with install.packages('crmReg').")

  if (ncol(data) < 2) stop("Need at least 2 variables.")

  # Convert character to factor
  for (j in seq_len(ncol(data))) {
    if (is.character(data[[j]])) data[[j]] <- as.factor(data[[j]])
  }

  if (!any(is.na(data))) {
    n <- nrow(data); p <- ncol(data)
    W <- matrix(1, n, p, dimnames = list(rownames(data), colnames(data)))
    return(list(data_imputed = data, cellweights = W,
                converged = TRUE, iterations = 0L))
  }
  if (any(apply(data, 1, function(x) all(is.na(x)))))
    stop("Unit non-responses detected. Remove them first.")

  ## ---- detect variable types ----
  rn <- rownames(data)
  cn <- colnames(data)
  n <- nrow(data); p <- ncol(data)
  is_continuous <- vapply(data, function(x) is.numeric(x) || is.integer(x), logical(1))
  is_categorical <- !is_continuous
  M <- is.na(data)
  vars_miss <- which(colMeans(M) > 0)

  ## work under safe positional names: pasted model formulas break with
  ## duplicated or non-syntactic column names (the response can resolve
  ## to the wrong column); original names are restored on exit
  cn_int <- .cw_safe_names(p)
  colnames(data) <- cn_int

  ## ---- initialise ----
  data <- initialise(data, mixed = NULL, method = "median")
  W <- matrix(1, n, p, dimnames = list(rn, cn))

  ## ---- outer loop ----
  converged <- FALSE
  iterations <- 0L
  d <- Inf

  while (d > eps && iterations < maxit) {
    iterations <- iterations + 1L
    data_previous <- data

    if (trace) message(paste("cellReg: iteration", iterations))

    ## ---- inner loop: cycle through variables ----
    for (j in vars_miss) {
      miss_j <- M[, j]
      n_miss <- sum(miss_j)
      if (n_miss == 0) next

      pred_cols <- setdiff(seq_len(p), j)

      if (is_continuous[j]) {
        ## ---- continuous response ----
        result <- switch(engine,
          crm          = .engine_crm(data, j, pred_cols, miss_j, cn_int),
          `cellwise-mm` = .engine_cellwise_mm(data, j, pred_cols, miss_j, W, cn_int),
          `shooting-s`  = .engine_shooting_s(data, j, pred_cols, miss_j, cn_int)
        )

        # Impute (deterministic during iteration)
        data[miss_j, j] <- result$fitted[miss_j]

        # Update cell weights for observed cells
        if (!is.null(result$cell_flags)) {
          # CRM/shooting-s return binary flags for predictor cells
          cont_preds <- pred_cols[is_continuous[pred_cols]]
          for (k in cont_preds) {
            k_local <- match(k, pred_cols)
            if (k_local <= ncol(result$cell_flags)) {
              obs_k <- !M[, k]
              W[obs_k, k] <- pmin(W[obs_k, k],
                                   1 - result$cell_flags[obs_k, k_local])
            }
          }
        }
        if (!is.null(result$response_weights)) {
          W[!miss_j, j] <- result$response_weights[!miss_j]
        }

      } else {
        ## ---- categorical response: weighted multinomial ----
        w_row <- apply(W[, pred_cols[is_continuous[pred_cols]], drop = FALSE],
                       1, function(ww) exp(mean(log(pmax(ww, 1e-10)))))

        form_j <- as.formula(paste0(cn_int[j], " ~ ",
                                     paste(cn_int[pred_cols], collapse = " + ")))
        multimod <- tryCatch(
          suppressMessages(nnet::multinom(form_j, data = data, weights = w_row,
                                          maxit = 50, trace = FALSE)),
          error = function(e) suppressMessages(
            nnet::multinom(form_j, data = data, maxit = 50, trace = FALSE))
        )
        prob_pred <- predict(multimod, newdata = data[miss_j, , drop = FALSE],
                              type = "probs")
        lvls <- levels(data[[j]])
        if (is.null(dim(prob_pred))) {
          if (length(lvls) == 2)
            prob_pred <- matrix(c(1 - prob_pred, prob_pred), nrow = n_miss)
          else
            prob_pred <- matrix(prob_pred, nrow = 1)
          colnames(prob_pred) <- lvls
        }
        imputed_cats <- apply(prob_pred, 1, function(pp) {
          pp <- pmax(pp, 0); colnames(prob_pred)[which.max(pp)]
        })
        data[miss_j, j] <- factor(imputed_cats, levels = lvls)
      }
    }  # end inner loop

    ## ---- check convergence ----
    d <- 0; n_imp <- 0
    for (jj in which(is_continuous)) {
      miss_jj <- M[, jj]
      if (any(miss_jj)) {
        prev <- data_previous[[jj]][miss_jj]
        curr <- data[[jj]][miss_jj]
        d <- d + sum((prev - curr)^2) / (sum(prev^2) + 1e-10)
        n_imp <- n_imp + 1L
      }
    }
    for (jj in which(is_categorical)) {
      miss_jj <- M[, jj]
      if (any(miss_jj)) {
        d <- d + sum(data_previous[[jj]][miss_jj] != data[[jj]][miss_jj]) / sum(miss_jj)
        n_imp <- n_imp + 1L
      }
    }
    if (n_imp > 0) d <- d / n_imp
    if (trace) message(paste("  convergence:", round(d, 6)))
    if (d <= eps) converged <- TRUE
  }

  ## ---- add uncertainty to final output ----
  if (uncert != "none") {
    for (j in vars_miss) {
      if (!is_continuous[j]) next
      miss_j <- M[, j]
      if (!any(miss_j)) next
      obs_j <- !miss_j
      pred_cols <- setdiff(seq_len(p), j)
      result <- switch(engine,
        crm           = .engine_crm(data, j, pred_cols, miss_j, cn_int),
        `cellwise-mm` = .engine_cellwise_mm(data, j, pred_cols, miss_j, W, cn_int),
        `shooting-s`  = .engine_shooting_s(data, j, pred_cols, miss_j, cn_int)
      )
      if (uncert == "pmm") {
        data[miss_j, j] <- .pmm_impute(
          pred_miss = result$fitted[miss_j],
          y_obs = data[[j]][obs_j],
          pred_obs = result$fitted[obs_j],
          n_donors = 5
        )
      } else {
        sigma <- mad(data[[j]][obs_j] - result$fitted[obs_j])
        if (sigma < .Machine$double.eps) sigma <- .robust_scale(data[[j]][obs_j] - result$fitted[obs_j])
        data[miss_j, j] <- result$fitted[miss_j] + rnorm(sum(miss_j), 0, sigma)
      }
    }
  }

  rownames(data) <- rn
  colnames(data) <- cn
  list(data_imputed = data, cellweights = W,
       converged = converged, iterations = iterations)
}


# ============================================================================
# Engine 1: CRM (Filzmoser et al. 2020)
# ============================================================================

.engine_crm <- function(data, j, pred_cols, miss_j, cn) {
  form <- as.formula(paste0(cn[j], " ~ ",
                             paste(cn[pred_cols], collapse = " + ")))
  crm_res <- tryCatch(
    suppressMessages(crmReg::crm(form, data = data, verbose = FALSE)),
    error = function(e) NULL
  )

  if (is.null(crm_res)) {
    # Fallback to lmrob
    fit <- robustbase::lmrob(form, data = data)
    return(list(
      fitted = predict(fit, newdata = data),
      cell_flags = NULL,
      response_weights = fit$rweights
    ))
  }

  # Cell flags: 1 = outlying, 0 = clean. crm's $cellwiseoutliers holds the
  # (continuous) deviation of each flagged cell and 0 for clean cells, so
  # binarise it -- taking abs() directly and feeding it to 1 - flag downstream
  # produced negative cell weights for large deviations.
  cell_flags <- (crm_res$cellwiseoutliers != 0) * 1
  # Response weights from casewise diagnostics
  resp_w <- rep(1, nrow(data))
  if (!is.null(crm_res$weights)) resp_w <- crm_res$weights

  list(
    fitted = crm_res$fitted.values,
    cell_flags = cell_flags,
    response_weights = resp_w
  )
}


# ============================================================================
# Engine 2: Cellwise-Weighted MM (hybrid)
# ============================================================================

.engine_cellwise_mm <- function(data, j, pred_cols, miss_j, W, cn) {
  form <- as.formula(paste0(cn[j], " ~ ",
                             paste(cn[pred_cols], collapse = " + ")))

  # Row weights = geometric mean of cell weights across continuous predictors
  cont_pred_idx <- pred_cols[vapply(data[pred_cols], function(x)
    is.numeric(x) || is.integer(x), logical(1))]

  if (length(cont_pred_idx) > 0) {
    w_row <- apply(W[, cont_pred_idx, drop = FALSE], 1, function(ww)
      exp(mean(log(pmax(ww, 1e-10)))))
    # Combine with response weight
    w_row <- w_row * W[, j]
  } else {
    w_row <- W[, j]
  }

  # Fit MM-estimation with row weights
  fit <- tryCatch(
    robustbase::lmrob(form, data = data, weights = w_row),
    error = function(e) NULL
  )

  if (is.null(fit)) {
    # Fallback: unweighted lmrob
    fit <- tryCatch(
      robustbase::lmrob(form, data = data),
      error = function(e) {
        # Last resort: lm
        stats::lm(form, data = data)
      }
    )
  }

  fitted_vals <- predict(fit, newdata = data)

  # Response weights from residuals
  r <- data[[j]] - fitted_vals
  sigma <- mad(r[!miss_j])
  if (sigma < .Machine$double.eps) sigma <- .robust_scale(r[!miss_j])
  resp_w <- tukey_weight(r / sigma, k = 4.685)

  list(
    fitted = fitted_vals,
    cell_flags = NULL,
    response_weights = resp_w
  )
}


# ============================================================================
# Engine 3: Shooting S-estimator (Öllerer, Alfons, Croux 2016)
# ============================================================================

.engine_shooting_s <- function(data, j, pred_cols, miss_j, cn) {
  # Implementation of the Shooting S algorithm
  # Iterates between cell detection and S-estimation

  y <- data[[j]]
  # Build design matrix (numeric only for shooting S)
  X_parts <- list()
  for (k in pred_cols) {
    if (is.numeric(data[[k]]) || is.integer(data[[k]])) {
      X_parts[[length(X_parts) + 1]] <- data[[k]]
    } else if (is.factor(data[[k]])) {
      mm <- model.matrix(~ data[[k]] - 1)[, -1, drop = FALSE]
      X_parts[[length(X_parts) + 1]] <- mm
    }
  }
  if (length(X_parts) == 0) {
    return(list(fitted = rep(median(y, na.rm = TRUE), length(y)),
                cell_flags = NULL, response_weights = rep(1, length(y))))
  }
  X <- do.call(cbind, X_parts)
  n <- nrow(X); q <- ncol(X)

  # Step 1: Initial S-estimate on full data
  X_int <- cbind(1, X)
  s_fit <- tryCatch(
    robustbase::lmrob.S(X_int, y,
      control = robustbase::lmrob.control(method = "S", k.max = 200)),
    error = function(e) NULL
  )
  if (is.null(s_fit)) {
    beta <- .weighted_qr_solve(X_int, y,
      matrix(1, n, q + 1), rep(1, n), rep(1, n))
  } else {
    beta <- s_fit$coefficients
  }

  # Step 2: Shooting iterations
  cell_flags <- matrix(0, nrow = n, ncol = q)
  max_shooting <- 5

  for (shoot_iter in seq_len(max_shooting)) {
    # Compute residuals
    r <- y - X_int %*% beta
    sigma <- mad(r)
    if (sigma < .Machine$double.eps) sigma <- .robust_scale(r)

    # For each predictor cell: compute its contribution to the residual
    # A cell is flagged if removing it substantially reduces the residual
    flags_new <- matrix(0, n, q)
    for (k in seq_len(q)) {
      # Partial residual: contribution of column k
      contrib_k <- X[, k] * beta[k + 1]  # +1 for intercept
      # Standardized contribution
      u_k <- abs(contrib_k) / sigma
      # Flag cells with large contribution relative to the residual
      flags_new[, k] <- as.integer(u_k > 3 & abs(r / sigma) > 2)
    }

    if (identical(flags_new, cell_flags)) break
    cell_flags <- flags_new

    # Re-estimate with flagged cells set to zero in design
    X_clean <- X
    X_clean[cell_flags == 1] <- 0
    X_clean_int <- cbind(1, X_clean)

    s_fit2 <- tryCatch(
      robustbase::lmrob.S(X_clean_int, y,
        control = robustbase::lmrob.control(method = "S", k.max = 200)),
      error = function(e) NULL
    )
    if (!is.null(s_fit2)) {
      beta <- s_fit2$coefficients
    }
  }

  fitted_vals <- as.numeric(X_int %*% beta)

  # Response weights
  r <- y - fitted_vals
  sigma <- mad(r)
  if (sigma < .Machine$double.eps) sigma <- .robust_scale(r)
  resp_w <- tukey_weight(r / sigma, k = 4.685)

  # Map cell_flags back to predictor columns
  # (cell_flags columns correspond to expanded X, not original pred_cols)
  # For simplicity, return the full flag matrix
  list(
    fitted = fitted_vals,
    cell_flags = cell_flags,
    response_weights = resp_w
  )
}