R/helper_FRTCI.R

In hettx: Fisherian and Neymanian Methods for Detecting and Measuring Treatment Effect Variation

## Test sweeping over range of plausibe taus for conservative p-value
FRTCI <- function(Y, Z, X = NULL, test.stat = SKS.stat, B=500, 
                  gamma=0.0001, grid.gamma=100*gamma, 
                  grid.size=151,
                  te.vec=NULL, return.matrix=FALSE,
                  n.cores=1,
                  verbose=TRUE, ... ) {
    
    
    if ( is.null(te.vec) ) {
        if ( grid.size %% 2 == 0 ) {
            grid.size <- grid.size+1
        }

        if( is.null(X) ){
            te.vec <- get.tau.vector( Y, Z, gamma=gamma, grid.size=grid.size, grid.gamma=grid.gamma )
        }else{
            te.vec <- get.tau.vector( Y, Z, X, gamma=gamma, grid.size=grid.size, grid.gamma=grid.gamma )
        }
    } else {
        grid.size = length( te.vec )
    }
    te.hat <- attr( te.vec, "te.hat" )
    te.se <- attr( te.vec, "te.se" )
    te.MOE <- attr( te.vec, "te.MOE" )
    
    ## IMPUTE MISSING POTENTIAL OUTCOMES
    Y1.mat <- sapply(te.vec, function(te) ifelse(Z, Y, Y + te) )
    Y0.mat <- sapply(te.vec, function(te) ifelse(Z, Y - te, Y) )

    if( is.null(X) ){
        res <- generate.permutations( Y, Z=Z, test.stat=test.stat, Y0.mat=Y0.mat, Y1.mat=Y1.mat, B=B, n.cores, verbose=verbose, ... )
    }else{
        res <- generate.permutations( Y, Z=Z, test.stat=test.stat, Y0.mat=Y0.mat, Y1.mat=Y1.mat, B=B, n.cores, verbose=verbose, X=X, ... )
    }
    
    ci.p = res$ci.p + gamma
    
    t = res$ks.obs
    p.value = max( ci.p )
    n=length(Y)
    method = "FRT CI Test for Treatment Effect Heterogeneity"
    DAT = paste( n, " observations", sep="")
    
    if ( !return.matrix ) {
        ks.mat=NULL
    } else {
        ks.mat = res$ks.mat
    }
    
    structure(list(statistic = t, p.value = p.value,
                   p.value.plug = ci.p[(grid.size+1)/2],
                   method=method,
                   data.name = DAT,
                   Y=Y, Z=Z, n=n, ci.p=ci.p, te.vec=te.vec,
                   te.hat=te.hat,
                   te.SE=te.se,
                   te.MOE = te.MOE,
                   B=B, gamma=gamma, ks.mat=ks.mat ),
              
              class = "FRTCI.test")
} 

## Test using plug-in sample average treatment effect
FRTplug <- function( Y, Z, test.stat=SKS.stat, tau.hat=mean(Y[Z == 1]) - mean(Y[Z == 0]), ... ){
    mth = FRTCI( Y, Z, test.stat=test.stat, te.vec=c(tau.hat), n.cores = 1, ...)
    mth$method = "FRT Plug-in Test for Treatment Effect Heterogeneity"
    mth
}

## Test with treatment-covariate interactions to adjust for known heterogeneity
FRTCI.interact <- function( Y, Z, W, X=NULL, test.stat = SKS.stat.int.cov, B=500, 
                            gamma=0.0001, grid.gamma=100*gamma, 
                            grid.size=151, return.matrix=FALSE, 
                            n.cores=1, verbose=TRUE, ... ) {
    
    grid.info = get.testing.grid( Y, Z, W=W, X=X, gamma, grid.size )
    
    te.MOE = NA
    
    te.grid = grid.info$te.grid
    Y1.mat <- grid.info$Y1.mat
    Y0.mat <- grid.info$Y0.mat

    if( is.null(X) ){
        res <- generate.permutations( Y, Z, test.stat, Y0.mat, Y1.mat, B=B, n.cores=n.cores, verbose=verbose, W=W, ... )
    }else{
        res <- generate.permutations( Y, Z, test.stat, Y0.mat, Y1.mat, B=B, n.cores=n.cores, verbose=verbose, X=X, W=W, ... )
    }
    
    t = res$ks.obs
    ci.p = res$ci.p + gamma
    p.value = max( ci.p )
    n=length(Y)
    method = "FRT CI Test for Tx Effect Heterogeneity Beyond a Systematic Model"
    DAT = paste( n, " observations", sep="")
    
    if ( !return.matrix ) {
        ks.mat=NULL
    } else {
        ks.mat = res$ks.mat
    }
    
    structure(list(statistic = t, p.value = p.value,
                   p.value.plug = ci.p[1],
                   method=method,
                   data.name = DAT,
                   Y=Y, Z=Z, n=n, ci.p=ci.p, te.grid=te.grid,
                   B=B, gamma=gamma, ks.mat=ks.mat,
                   W=W, X=X ),
              
              class = "FRTCI.test")
}