R/PT.Khmaladze.fit.R

In RATest: Randomization Tests

#' @title  Permutation Test for Heterogeneous Treatment Effects with a Nuisance Parameter
#'
#' @description A permutation test of the two-sample goodness-of-fit hypothesis in the presence of an estimated nuisance parameter. The permutation test considered here is based on the Khmaladze transformation of the empirical process (Khmaladze (1981)), and adapted by Chung and Olivares (2021).
#'
#' @param y1 Numeric. A vector containing the response variable of the treatment group.
#' @param y0 Numeric. A vector containing the response variable of the control group.
#' @param alpha Numeric. Nominal level for the test. The default is 0.05.
#' @param n.perm Numeric. Number of permutations needed for the stochastic approximation of the p-values. The default is n.perm=999.
#' @param akj.control Optional named list of tuning parameters passed to
#'   \code{quantreg::akj}, used to estimate the (extended) score function via
#'   adaptive kernel density estimation. Recognized entries are \code{alpha},
#'   \code{kappa}, \code{h}, and \code{iker1}; see \code{?quantreg::akj}. Any
#'   entry not supplied falls back to the \code{akj} defaults, which are used by
#'   default and are recommended: in the simulations reported in Chung and
#'   Olivares (2021) they maximize power, and alternative (e.g. larger-bandwidth)
#'   settings that reduce the score's estimation error were found to *lower*
#'   power. This argument is provided only for users who have a specific reason
#'   to prefer different settings for their application.
#'
#' @return An object of class "PT.Khmaladze.fit" containing at least the following components:
#'
#'  \item{n_populations}{Number of groups.}
#'  \item{N}{Sample Size.}
#'  \item{T.obs}{Observed test statistic.}
#'  \item{shift}{The estimated nuisance parameter (average treatment effect).}
#'  \item{cv}{Critical Value. This value is used in the general construction of a randomization test.}
#'  \item{pvalue}{P-value.}
#'  \item{T.perm}{Vector. Test statistic recalculated for all permutations used in the stochastic approximation.}
#'  \item{n_perm}{Number of permutations.}
#'  \item{sample_sizes}{Groups size.}
#'
#' @author Mauricio Olivares
#' @references
#' Khmaladze, E. (1981). Martingale Approach in the Theory of Goodness-of-fit Tests. Theory of Probability and its Application, 26: 240-257.
#' Chung, E. and Olivares, M. (2021). Permutation Test for Heterogeneous Treatment Effects with a Nuisance Parameter. Journal of Econometrics.
#'
#' @keywords permutation test goodness-of-fit Khmaladze Transformation
#' @include group.action.R
#' @include randomization.test.R
#' @import quantreg
#' @importFrom stats quantile cor var runif ecdf
#' @importFrom utils modifyList
#' @examples
#'\dontrun{
#' Y0 <- rnorm(100, 1, 1)
#' # Treatment Group with constant shift equals to 1
#' Y1 <- Y0 + 1
#' Tx = sample(100) <= 0.5*(100)
#' # Observed Outcome
#' Y = ifelse( Tx, Y1, Y0 )
#' dta <- data.frame(Y = Y, Z = as.numeric(Tx))
#' pt.GoF<-PT.Khmaladze.fit(dta$Y[dta$Z==1],dta$Y[dta$Z==0],n.perm = 49)
#' summary(pt.GoF)
#' }
#' @export



PT.Khmaladze.fit <- function(y1, y0, alpha=0.05, n.perm=999, akj.control=list()){

  if (anyNA(c(y1,y0))) stop("NAs in first or second argument")
  if (!is.numeric(c(y1,y0)))  stop("Arguments must be numeric")

  # ---------------------------- #
  #   Calculates the statistic   #
  # ---------------------------- #

  # Sample sizes. m = treatment group, N - m = control group.
  N <- length(c(y1,y0))
  m <- length(y1)
  n <- N - m                      # control-group size = length(y0)
  lengths <- list(m, n)
  # Calculate the shift (average treatment effect)
  shift <- mean(y1) - mean(y0)
  # Recentering
  Y1.star <- y1 - shift
  Y0.star <- y0
  Z <- c(Y1.star, Y0.star)        # Stacked data

  # Evaluation grid for the score estimation and numeric integration.
  #
  # The two-sample empirical process is evaluated at the CONTROL quantiles
  # Q0hat(t); its resolution is therefore capped by the number of distinct
  # control order statistics, i.e. by n = length(y0). A grid finer than n
  # places evaluation points on flat stretches of the process, which makes the
  # score column of the recursive-least-squares design collinear with the
  # constant column and silently zeroes the score-direction projection of the
  # compensator. We therefore set p = n (control-group size). This is robust to
  # group imbalance because the control sample is the binding resolution
  # constraint regardless of the treatment-group size.
  p <- n
  eps <- .Machine$double.eps
  t <- seq(eps, 1, length.out = p)

  # Generate random permutations. The first column of Sn is the OBSERVED
  # (recentered) data; the remaining columns are permutations of it.
  S_perm <- group.action(Z, n.perm, "permutations")
  Sn <- cbind(Z, S_perm)

  # Calculation of the test statistic for all permutations.
  stat <- apply(Sn, 2, function(x) Khm.trans(x, m, n, t, p, akj.control))

  # Observed test statistic
  T.obs <- stat[1]
  # Test statistic for the permuted samples
  T.perm <- stat[-1]

  # Critical value
  cv <- randomization.test(T.obs, T.perm, alpha)
  # Indicator rule (permutation p-value)
  ind.rule <- mean(ifelse(T.perm >= T.obs, 1, 0))


  object_perm <- list() # Collects the info required for summary
  object_perm$N <- N
  object_perm$T.obs <- T.obs
  object_perm$shift <- shift
  object_perm$cv <- cv[2]
  object_perm$pvalue <- ind.rule
  object_perm$T.perm <- T.perm
  object_perm$n_perm <- n.perm
  object_perm$sample_sizes <- c(lengths)


  class(object_perm) <- "PT.Khmaladze.fit"

  return(object_perm)

}


# ------------------------------------------------------------------------- #
#  Internal: uniform-scale two-sample empirical process, eq. (7) of         #
#  Chung and Olivares (2021).                                               #
#                                                                           #
#  v(t) = Fhat_1(Q0hat(t)) - Fhat_0(Q0hat(t)),                              #
#                                                                           #
#  with Q0hat(t) the CONTROL quantile function, generalized-inverse         #
#  convention (type = 1). Returned WITHOUT the sqrt(mn/N) scaling; that      #
#  factor is applied once in Khm.trans. Computed here ONCE and reused by     #
#  both the statistic and the compensator, so the two cannot drift apart.   #
# ------------------------------------------------------------------------- #
".emp.process.khm" <- function(Y1, Y0, t){
  q0 <- as.numeric(stats::quantile(Y0, probs = t, type = 1))
  stats::ecdf(Y1)(q0) - stats::ecdf(Y0)(q0)
}

# ------------------------------------------------------------------------- #
#  Internal: extended score gdot(s) = (1, f0'(Q0(s)) / f0(Q0(s)))',         #
#  estimated by adaptive kernel density estimation (quantreg::akj).          #
#  akj returns psi = -f'/f, so the score component is -psi.                  #
#  Evaluation points are the CONTROL quantiles (outcome scale), NOT the      #
#  probability grid t.                                                       #
# ------------------------------------------------------------------------- #
".gdot.khm" <- function(Y0, t, p, akj.control = list()){
  q0 <- as.numeric(stats::quantile(Y0, probs = t, type = 1))
  args <- utils::modifyList(
    list(x = sort(Y0), z = q0),   # akj requires sorted x
    akj.control
  )
  sc <- do.call(quantreg::akj, args)
  gdot2 <- -sc$psi
  cbind(rep(1, p - 1), gdot2[1:(p - 1)])
}

# ------------------------------------------------------------------------- #
#  Martingale-transformed two-sample KS statistic.                          #
#  Z is the (recentered) stacked sample; m = treatment size, n = control.   #
# ------------------------------------------------------------------------- #
"Khm.trans" <- function(Z, m, n, t, p, akj.control = list()){
  Y1 <- Z[1:m]
  Y0 <- Z[(m + 1):(m + n)]

  # Uniform-scale empirical process, computed once.
  v <- .emp.process.khm(Y1, Y0, t)
  # Extended score on the control quantiles.
  gdot0 <- .gdot.khm(Y0, t, p, akj.control)
  # Compensator (recursive-least-squares projection onto the score).
  comp <- K.hat(v, gdot0, p)

  # The Khmaladze transformation of the uniform empirical process:
  # like the residuals of a regression of the process on the (extended)
  # score function.
  sqrt(n * m / (m + n)) * max(abs(v - comp))
}

# ------------------------------------------------------------------------- #
#  Compensator via recursive least squares.                                 #
#  Takes the ALREADY-COMPUTED process v and the score-basis matrix gdot,     #
#  so it can no longer rebuild (and desynchronize from) the process.        #
#                                                                           #
#  X and y are left unscaled: omitting sqrt(dt) from X and 1/sqrt(dt) from   #
#  y gives bhat = bhat_paper / q, and omitting the leading 1/q from the      #
#  cumulative sum restores the compensator exactly. The two omissions        #
#  cancel.                                                                   #
# ------------------------------------------------------------------------- #
"K.hat" <- function(v, gdot, p){
  X <- gdot[(p - 1):1, , drop = FALSE]
  y <- rev(diff(v))
  bhat <- quantreg::lm.fit.recursive(X, y, int = FALSE)
  bhat <- bhat[, (p - 1):1, drop = FALSE]
  dk <- colSums(t(gdot) * bhat)   # = diag(gdot %*% bhat), O(p) not O(p^2)
  c(0, cumsum(dk))
}