R/honest_splitting.R

In CRE: Interpretable Discovery and Inference of Heterogeneous Treatment Effects

#' @title
#' Honest splitting
#'
#' @description
#' Splits data into discovery and inference sub-samples.
#'
#' @param y An observed response vector.
#' @param z A treatment vector.
#' @param X A features matrix.
#' @param ratio_dis A double number indicating the ratio of data delegated to
#' the discovery sub-sample.
#' @param ite A vector of estimated ITE.
#'
#' @return
#' A list containing the discovery and inference sub-samples.
#'
#' @keywords internal
#'
honest_splitting <- function(y, z, X, ratio_dis, ite = NULL) {

  logger::log_debug("(Honest) Splitting the dataset...")

  n <- length(y)

  index <- sample(1:n, round(n * ratio_dis), replace = FALSE)

  X <- as.matrix(X)
  y <- as.matrix(y)
  z <- as.matrix(z)
  if (!is.null(ite)) {
    ite <- as.matrix(ite)
  }

  discovery <- list(y = y[index],
                    z = z[index],
                    X = X[index, ],
                    ite = ite[index, ])
  inference <- list(y = y[-index],
                    z = z[-index],
                    X = X[-index, ],
                    ite = ite[-index, ])
  logger::log_debug("Done with (Honest) Splitting the dataset.")
  return(list(discovery = discovery, inference = inference))
}