Compute optimal treatment assignment for treatment effect precision

neymanAllocationStrata.R

#' Compute Neyman allocation propensity scores for inference-optimal treatment assignment in data table [very fast]
#' @param df data.table
#' @param y outcome name
#' @param w treatment name
#' @param x covariate names (must all be discrete)
#' @return data.table with strata level conditional means, variances, propensity scores,
#'  and neyman allocation propensities.
#' @export
neymanAllocation = function(df, y, w, x){
    df1 = copy(df); N = nrow(df1)
    setnames(df1, c(y, w), c("y", "w"))
    nTreat = length(unique(df1$w))
    # conditional means and standard deviations within treat X covariate strata
    wxStrata = df1[, .(muHat = mean(y), sigHat = sd(y), num = .N), by = c(x, 'w')]
    # reshape to strata level
    wxStrataWide = dcast(wxStrata,
            as.formula(glue::glue("{glue::glue_collapse(x, '+' )} ~ w")),
            value.var = c("muHat", "sigHat", "num"))
    countCols = grep("num_", colnames(wxStrataWide), value = T)
    wxStrataWide[, nX := Reduce(`+`, .SD), .SDcols = countCols]
    wxStrataWide[, `:=`(fX = nX/N)]
    # empirical pscore
    nCols = grep("num_", colnames(wxStrataWide), value = T)
    wxStrataWide[, paste0("piHat_", nCols) := lapply(.SD, \(x) x/nX),
        .SDcols = nCols]
    # neyman allocation
    sigCols = grep("sigHat_", colnames(wxStrataWide), value = T)
    wxStrataWide[, sigDenom := Reduce(`+`, .SD), .SDcols = sigCols]
    wxStrataWide[, paste0("piNeyman_", nCols) := lapply(.SD, \(x) x/sigDenom ),
        .SDcols = sigCols]
    if(nTreat == 2){ # also returns treatment effects with 2 treatments; else compute desired contrasts yourself 
        wxStrataWide[, tauX := muHat_1 - muHat_0]
        tau = wxStrataWide[, sum(fX * tauX)]
        se  = wxStrataWide[, sqrt(mean(sigHat_1^2/piHat_1 +
                                       sigHat_0^2/(1-piHat_0) +
                                       (tauX - tau)^2)) / sqrt(sum(nX))
                          ]
        return(list(strataLev = wxStrataWide, tau = tau, se = se))
    } else{
        return(wxStrataWide)
    }
}

#' Compute unconstrained optimal treatment probability
#' @param pNey neyman allocation propensity
#' @param p1 propensity score in first phase
#' @param kappa share of obs in first phase share
#' @return vector of variance-optimal treatment propensities
#' @export
ssProb = function(pNey, p1, kappa=.5){
    (1/(1 - kappa)) * (pNey - kappa * p1)
}