R/ctree_stacked.R

In ctreeMI: Conditional Inference Trees with Stacked Multiple Imputation

#' Conditional Inference Tree on Stacked Multiply Imputed Data
#'
#' @description
#' Fits a conditional inference tree (ctree) on stacked multiply imputed
#' datasets using the Stack / M rescaling procedure described in
#' Sherlock et al. (2026). Multiply imputed datasets are concatenated
#' vertically ("stacked"), and the significance threshold used for
#' node-level pruning is divided by the number of imputations (M) to
#' counteract the artificially inflated sample size. This yields a single,
#' coherent, interpretable tree that incorporates imputation variability
#' without requiring the pooling of structurally different trees.
#'
#' @param formula A model formula, passed to [partykit::ctree()].
#' @param data A `mids` object from [mice::mice()], a list of imputed
#'   data frames, or a single complete data frame. If a single complete
#'   data frame is supplied (no missing data handling needed), the function
#'   falls back to a standard [partykit::ctree()] call with a warning.
#' @param m Integer. Number of imputations to use. If `data` is a `mids`
#'   object or a list, defaults to the number of datasets available.
#'   Ignored when `data` is a plain data frame.
#' @param alpha Numeric. The nominal significance threshold for node-level
#'   splitting (default 0.05). It is applied to p-values recomputed from
#'   node statistics that have been divided by `m`; `alpha` itself is not
#'   rescaled. Must be strictly between 0 and 1.
#' @param verbose Logical. If `TRUE` (default), prints a message summarising
#'   the stacking and correction applied.
#' @param ... Additional arguments passed to [partykit::ctree_control()].
#'   Note: `alpha` in `...` is ignored in favour of the `alpha` argument
#'   above.
#'
#' @return An object of class `c("ctreeMI", "constparty", "party")`. This
#'   inherits all methods from [partykit::ctree()] output and additionally
#'   carries a `ctreeMI_info` attribute with the following elements:
#'   \describe{
#'     \item{`m`}{Number of imputations used.}
#'     \item{`n_original`}{Number of rows in a single imputed dataset.}
#'     \item{`n_stacked`}{Total rows in the stacked dataset.}
#'     \item{`alpha`}{The significance threshold applied to the rescaled
#'       p-values.}
#'     \item{`correction`}{Character, `"statistic/M"`.}
#'     \item{`node_stats`}{Data frame of per-node raw statistics, degrees
#'       of freedom, rescaled statistics, p-values, and retention.}
#'     \item{`n_splits_before`, `n_splits_after`}{Splits before and after
#'       the correction was applied.}
#'     \item{`formula`}{The model formula.}
#'     \item{`call`}{The matched call.}
#'   }
#'
#' @details
#' ## Methodological background
#'
#' When data contain missing values, a common and principled approach is
#' multiple imputation (Rubin, 1987), wherein M completed datasets are
#' generated by drawing from the posterior predictive distribution of the
#' missing values. For most statistical models, results from M datasets
#' are pooled using Rubin's rules.
#'
#' Conditional inference trees (ctree; Hothorn, Hornik & Zeileis, 2006)
#' are not straightforward to pool across imputations because the tree
#' structure itself -- which variables are split, at what values, and in
#' what order -- can differ across imputations. Pooling structurally
#' different trees produces inconsistent subgroup definitions and
#' uninterpretable results.
#'
#' Rodgers et al. (2021) proposed stacking the M imputed datasets
#' vertically and fitting a single tree to the combined data. This
#' circumvents the structural-mismatch problem and produces one
#' interpretable tree. However, stacking inflates the nominal sample size
#' by a factor of M, causing node-level test statistics to be similarly
#' inflated and the tree to split more aggressively than warranted.
#'
#' Sherlock et al. (2026) proposed and validated the **Stack / M**
#' correction: each node-level test statistic computed on the stacked data
#' is divided by M, the p-value is recomputed from the rescaled statistic,
#' the Bonferroni adjustment for the number of candidate splitting
#' variables is reapplied, and nodes that no longer meet `alpha` are
#' pruned. Monte Carlo simulations under MCAR confirmed sub-nominal
#' (conservative) type-I error and reduced but acceptable power.
#'
#' ## Correction applied to the statistic, not to alpha
#'
#' Rescaling the statistic is **not** equivalent to dividing the
#' significance threshold by M. Writing `q(p, df)` for the chi-square
#' quantile function, the two rules are:
#'
#' \itemize{
#'   \item statistic rescaling (correct): reject when
#'     `X > M * q(1 - alpha, df)`
#'   \item threshold rescaling (incorrect): reject when
#'     `X > q(1 - alpha / M, df)`
#' }
#'
#' These coincide only at `M = 1`. At `df = 1`, `alpha = 0.05`, `M = 30`
#' the first requires `X > 115.2` and the second only `X > 9.9`, so
#' threshold rescaling under-corrects by an order of magnitude and grows
#' substantially larger trees than the published method.
#'
#' Versions 0.1.0 and 0.2.0 of this package implemented threshold
#' rescaling. Trees fitted with those versions are under-corrected and
#' should be refitted.
#'
#' ## Usage with `mice`
#'
#' ```r
#' library(mice)
#' imp <- mice(my_data, m = 30, printFlag = FALSE)
#' fit <- ctree_stacked(outcome ~ ., data = imp, alpha = 0.05)
#' print(fit)
#' plot(fit)
#' ```
#'
#' ## Usage with a list of data frames
#'
#' ```r
#' imp_list <- lapply(1:30, function(i) complete(imp, i))
#' fit <- ctree_stacked(outcome ~ ., data = imp_list, alpha = 0.05)
#' ```
#'
#' @references
#' Sherlock, P., Mansolf, M., Hofheimer, J., Hockett, C. W., O'Connor, T. G.,
#'   Roubinov, D., Graff, J. C., Lai, J.-S., Bush, N. R., Wright, R. J., &
#'   Chiu, Y.-H. M. (2026). Beyond linear risk: A machine learning approach
#'   to understanding perinatal depression in context.
#'   *Multivariate Behavioral Research*, 1-16.
#'   \doi{10.1080/00273171.2026.2661244}
#'
#' Hothorn, T., Hornik, K., & Zeileis, A. (2006). Unbiased recursive
#'   partitioning: A conditional inference framework. *Journal of
#'   Computational and Graphical Statistics*, 15(3), 651-674.
#'   \doi{10.1198/106186006X133933}
#'
#' Hothorn, T., & Zeileis, A. (2015). partykit: A modular toolkit for
#'   recursive partitioning in R. *Journal of Machine Learning Research*,
#'   16, 3905-3909.
#'
#' Rodgers, J., Khoo, S.-T., & L?dtke, O. (2021). Handling missing data in
#'   structural equation models using multiple imputation and stacking.
#'   *Structural Equation Modeling*, 28(6), 915-930.
#'   \doi{10.1080/10705511.2021.1916925}
#'
#' Rubin, D. B. (1987). *Multiple imputation for nonresponse in surveys.*
#'   Wiley.
#'
#' @seealso
#' [partykit::ctree()], [partykit::ctree_control()], [mice::mice()],
#' [stack_imputations()], [rescale_statistic()], [prune_stackM()]
#'
#' @examples
#' \dontrun{
#' library(mice)
#'
#' # Introduce missingness into the airquality dataset
#' set.seed(42)
#' aq <- airquality
#' aq$Ozone[sample(nrow(aq), 20)] <- NA
#' aq$Solar.R[sample(nrow(aq), 15)] <- NA
#'
#' # Impute
#' imp <- mice(aq, m = 10, printFlag = FALSE)
#'
#' # Fit ctree with Stack/M correction
#' fit <- ctree_stacked(Ozone ~ Solar.R + Wind + Temp + Month,
#'                      data = imp,
#'                      alpha = 0.05)
#' print(fit)
#' plot(fit)
#' }
#'
#' @export
ctree_stacked <- function(formula,
                          data,
                          m       = NULL,
                          alpha   = 0.05,
                          verbose = TRUE,
                          ...) {

  cl <- match.call()

  # ## Input validation########################################################
  if (!inherits(formula, "formula")) {
    stop("`formula` must be a formula object, e.g. y ~ x1 + x2.")
  }
  if (alpha <= 0 || alpha >= 1) {
    stop("`alpha` must be strictly between 0 and 1.")
  }

  # ## Resolve data source####################################################-
  data_list <- resolve_data(data, m)
  m_actual  <- length(data_list)
  n_orig    <- nrow(data_list[[1]])

  if (m_actual == 1) {
    warning(
      "Only one dataset provided. Running standard ctree() without ",
      "Stack/M correction. Supply m > 1 imputed datasets for the ",
      "stacked-imputation workflow."
    )
    fit <- partykit::ctree(formula, data = data_list[[1]],
                           control = partykit::ctree_control(alpha = alpha, ...))
    return(fit)
  }

  # ## Stack datasets##########################################################-
  stacked  <- stack_imputations(data_list)
  n_stack  <- nrow(stacked)

  if (verbose) {
    message(sprintf(
      "[ctreeMI] Stacked %d imputed datasets (n = %d each; %d total rows).",
      m_actual, n_orig, n_stack))
  }

  # ## Grow, then correct######################################################
  #
  # The tree is grown at the NOMINAL alpha and the Stack/M correction is
  # applied afterwards by rescaling node statistics. Two facts make this
  # valid and efficient:
  #
  #  1. partykit selects the splitting variable and cut point independently
  #     of alpha; alpha only decides whether to stop. So pruning a tree by a
  #     criterion is equivalent to having grown it under that criterion.
  #  2. The correction is strictly STRICTER than the nominal threshold for
  #     any m >= 1, because M * q(1 - alpha, df) >= q(1 - alpha, df). Every
  #     split the correction could retain is therefore already present in a
  #     tree grown at the nominal alpha, so there is no need to grow a
  #     maximal tree first.
  #
  # `testtype = "Univariate"` is forced so that the stored p-values are raw
  # per-variable values. The multiplicity adjustment is reapplied in
  # prune_stackM() AFTER the statistic has been rescaled, which is the order
  # the method requires; leaving partykit to Bonferroni-adjust beforehand
  # would make the stored p-values impossible to invert reliably.
  dots <- list(...)
  dots$alpha    <- alpha
  dots$teststat <- "quadratic"
  dots$testtype <- "Univariate"
  ctrl <- do.call(partykit::ctree_control, dots)

  fit_full <- partykit::ctree(formula, data = stacked, control = ctrl)

  n_before <- length(setdiff(partykit::nodeids(fit_full),
                             partykit::nodeids(fit_full, terminal = TRUE)))

  pr  <- prune_stackM(fit_full, m = m_actual, alpha = alpha, verbose = verbose)
  fit <- pr$tree

  n_after <- length(setdiff(partykit::nodeids(fit),
                            partykit::nodeids(fit, terminal = TRUE)))

  # ## Attach ctreeMI metadata##################################################
  attr(fit, "ctreeMI_info") <- list(
    m               = m_actual,
    n_original      = n_orig,
    n_stacked       = n_stack,
    alpha           = alpha,
    correction      = "statistic/M",
    node_stats      = pr$node_stats,
    n_splits_before = n_before,
    n_splits_after  = n_after,
    formula         = formula,
    call            = cl
  )

  class(fit) <- c("ctreeMI", class(fit))
  fit
}


# ## Internal helper: resolve_data############################################

#' @keywords internal
resolve_data <- function(data, m) {

  # mice mids object
  if (methods::is(data, "mids")) {
    total_m <- data$m
    use_m   <- if (is.null(m)) total_m else min(m, total_m)
    if (!is.null(m) && m > total_m) {
      warning(sprintf(
        "`m` (%d) exceeds the number of imputations in the mids object (%d). ",
        "Using all %d imputations.", m, total_m, total_m
      ))
    }
    return(lapply(seq_len(use_m), function(i) mice::complete(data, i)))
  }

  # list of data frames
  if (is.list(data) && all(sapply(data, is.data.frame))) {
    use_m <- if (is.null(m)) length(data) else min(m, length(data))
    return(data[seq_len(use_m)])
  }

  # plain data frame -- no imputation
  if (is.data.frame(data)) {
    return(list(data))
  }

  stop(
    "`data` must be a `mids` object (from mice), a list of data frames, ",
    "or a single data frame."
  )
}