R/apply_function_to_imputed_data.R

In basecamb: Utilities for Streamlined Data Import, Imputation and Modelling

#' Apply function to dataframes in a mice object
#'
#' Wrapper function to apply a function on each dataframe in an imputed dataset
#'   created with `mice::mice()`.
#'
#' @param mice_data a mids object generated by `mice::mice()`.
#' @param fun the function to apply to each dataframe. May only take one
#'   positional argument of type data.frame.
#' @param ... other arguments passed to fun()
#'
#' @return a mids object with transformed data.
#'
#' @importFrom assertthat assert_that
#' @importFrom mice as.mids is.mids
#' @importFrom mice complete
#'
#' @export
#'
#' @author J. Peter Marquardt
apply_function_to_imputed_data <- function(mice_data, fun, ...) {
  assertthat::assert_that(mice::is.mids(mice_data))
  assertthat::assert_that(is.function(fun))

  long <- mice::complete(mice_data, action = "long", include = TRUE) # convert imputations into a long dataframe
  return_dat <- data.frame()  # initialize empty long dataframe for return

  #  apply method to each imputed dataset
  for (n_imp in unique(long$.imp)) {
    data <- fun(long[long$.imp == n_imp, ], ...)  # perform fun on one iteration only

    return_dat <- rbind.data.frame(return_dat, data)  # append results to long-type return df
  }

  return(mice::as.mids(return_dat)) # reconvert long dataframe to mids object
}