Conditional Average Treatment Effect estimation

Conditional Average Treatment Effect estimation with cross-fitting.

Usage

cate(
  response.model,
  treatment.model,
  cate.model = ~1,
  calibration.model = NULL,
  missing.model = NULL,
  auxiliary.model = NULL,
  data,
  contrast,
  nfolds = 1,
  rep = 1,
  id = NULL,
  silent = FALSE,
  stratify = FALSE,
  mc.cores = NULL,
  var.type = "IC",
  pseudo.outcome = c("regression", "dr"),
  second.order = TRUE,
  ...
)

Arguments

response.model: formula or learner object (formula => learner_glm)
treatment.model: formula or learner object (formula => learner_glm)
cate.model: formula specifying regression design for conditional average treatment effects
calibration.model: linear calibration model. Specify covariates in addition to predicted potential outcomes to include in the calibration.
missing.model: formula or learner object (formula => learner_glm(family = binomial); default NULL. Model for the missingness mechanism $P(R=1 \mid W, A)$. Required when the outcome in response.model contains NAs. If the formula LHS is omitted, the observation indicator is used automatically. When stratify = TRUE the missing model is fit separately per treatment arm. When supplied, the AIPW score is inverse-probability-of-observation weighted and (if second.order = TRUE) an additional second-order term is added to the influence function.
auxiliary.model: when the missing data model depends on post-randomization variables $Z$, this argument defines the nuisance model for $E(Y|W,A,Z)$
data: data.frame
contrast: treatment contrast (default 1 vs 0)
nfolds: number of folds (positive integer), or a pre-specified list of fold indices where each element is an integer vector of observation indices forming a partition of 1:nrow(data).
rep: number of replications of cross-fitting procedure by averaging estimates and influence functions from each replication
id: (integer or character) optional subject id vector of length nrow(data). The id can also be specified as part of the cate.model argument with a formula syntax: ~ 1 + cluster(id).
silent: suppress all messages and progressbars
stratify: if TRUE the response.model, missing.model and auxiliary.model (if provided) will be stratified by treatment
mc.cores: (optional) number of cores. parallel::mcmapply used instead of future
var.type: when equal to "IC" the asymptotic variance is derived from the influence function. Otherwise, based on expressions in Bannick et al. (2025) valid under different covariate-adaptive randomization schemes (only available for ATE and when calibration.model is also specified)
pseudo.outcome: character specifying the pseudo-outcome used to fit response.model when auxiliary.model is supplied. With "regression" (default), use predictions from auxiliary.model, $h(W,A,Z)$. With "dr", use $U = h(W,A,Z) + R\{Y-h(W,A,Z)\}/\rho(W,A,Z)$. The latter has conditional expectation $E[Y\mid A,W]$ when either the auxiliary outcome model or the observation model is correctly specified, under MAR. Note, for binary outcomes, "regression" requires a learner accepting fractional responses (e.g. quasibinomial regression). The "dr" option requires a learner accepting responses outside the outcome's original bounds. Ignored when auxiliary.model is not used.
second.order: add seconder order term to IF to handle misspecification of outcome models
...: additional arguments to future.apply::future_mapply

Value

cate.targeted object

Details

We have observed data $(Y,A,W)$ where $Y$ is the response variable, $A$ the binary treatment, and $W$ covariates. We further let $V$ be a subset of the covariates. Define the conditional potential mean outcome $$\psi_{a}(P)(V) = E_{P}[E_{P}(Y\mid A=a, W)|V]$$ and let $m(V; \beta)$ denote a parametric working model, then the target parameter is the mean-squared error $$\beta(P) = \operatorname{argmin}_{\beta} E_{P}[\{\Psi_{1}(P)(V)-\Psi_{0}(P)(V)\} - m(V; \beta)]^{2}$$

Missing data is handled under a Missing At Random assumption (MAR). Let $R$ denote the indicator for data not being missing, $R\perp Y|W,A,Z$ where the conditioning set must be observed for all subjects, and $Z$ denotes potentiel post-treatment variable.

The nuisance models are $Q(w,a) = E(Y|W=w, A=a)$, $g(w,a) = P(A=a|W=w)$, and $\rho(w, a, z) = P(R=1|W=w, A=a, Z=z)$, $h(w,a,z) = E(Y|W=w, A=a, Z=z)$.

For the expected potential outcome $E[Y(a)]$, the EIF then takes the form

$$\frac{I(A=a)}{g(W,a)}\left[ \frac{R}{\rho(W, a, Z)}\{Y-Q(W, a)\} + \big\{1 - \frac{R}{\rho(W, a, Z)}\big\}\{h(W,a,Z)-Q(W,a)\} \right] + Q(W,a) - \psi_a $$

Note that when the MAR assumption depends on $Z$, we can no longer obtain a consistent estimate of $Q(W,A)$ from $E[Y|A,W,R=1]$. Instead we an estimate of $Q(W,A)$ is obtained as $E[h(W,A,Z) | A, W]$. A more general doubly-robust estimate can be obtained from the regression $E[U|A,W]$ using the DR pseudo-outcome $U = h(W,A,Z) + R/\rho(W,A,Z)(Y-h(W,A,Z))$.

References

Mark J. van der Laan (2006) Statistical Inference for Variable Importance, The International Journal of Biostatistics.

Bannick, Shao & Liu et al. (2025) A General Form of Covariate Adjustment in Clinical Trials under Covariate-Adaptive Randomization, Biometrika.

Author

Klaus Kähler Holst, Andreas Nordland

Examples

sim1 <- function(n=1000, ...) {
  w1 <- rnorm(n)
  w2 <- rnorm(n)
  a <- rbinom(n, 1, plogis(-1 + w1))
  y <- cos(w1) + w2*a + 0.2*w2^2 + a + rnorm(n)
  data.frame(y, a, w1, w2)
}

d <- sim1(5000)
## ATE
cate(cate.model=~1,
     response.model=y~a*(w1+w2),
     treatment.model=a~w1+w2,
     data=d)
#>             Estimate Std.Err   2.5%  97.5%    P-value
#> E[y(1)]       1.8047 0.04831 1.7100 1.8994 2.099e-305
#> E[y(0)]       0.8308 0.01984 0.7919 0.8697  0.000e+00
#> ───────────                                          
#> (Intercept)   0.9740 0.05337 0.8694 1.0786  2.054e-74
## CATE
cate(cate.model=~1+w2,
     response.model=y~a*(w1+w2),
     treatment.model=a~w1+w2,
     data=d)
#>             Estimate Std.Err   2.5%  97.5%    P-value
#> E[y(1)]       1.8047 0.04831 1.7100 1.8994 2.099e-305
#> E[y(0)]       0.8308 0.01984 0.7919 0.8697  0.000e+00
#> ───────────                                          
#> (Intercept)   0.9502 0.05280 0.8467 1.0536  2.093e-72
#> w2            1.0377 0.04756 0.9445 1.1309 1.586e-105

if (FALSE) { # \dontrun{
# superlearner
mod1 <- list(
   glm = learner_glm(y~w1+w2),
   gam = learner_gam(y~s(w1) + s(w2))
)
s1 <- learner_sl(mod1, nfolds=5)
cate(cate.model=~1,
     response.model=s1,
     treatment.model=learner_glm(a~w1+w2, family=binomial),
     data=d,
     stratify=TRUE)
} # }

## Missing data
sim_missing_cate <- function(n = 5000, seed = 1) {
  set.seed(seed)
  w1 <- rnorm(n)
  w2 <- rnorm(n)
  a  <- rbinom(n, 1, 0.5) # randomized trial
  y_full <- 1 + a + w1 + 0.5 * w2 + rnorm(n)
  pR <- plogis(0.5 - 1 * w2 * a + 0.5 * a)
  R  <- rbinom(n, 1, pR)
  y  <- ifelse(R == 1, y_full, NA_real_)
  data.frame(y0 = y_full, y = y, a = a, w1 = w1, w2 = w2)
}
d <- sim_missing_cate()

# ignoring missing data (complete-case analysis)
cate(cate.model = ~1,
     response.model = y ~ a * w2, # wrong outcome model
     treatment.model = a ~ 1,
     data = na.omit(d), nfolds = 1L)
#>             Estimate Std.Err   2.5%  97.5%    P-value
#> E[y(1)]       1.9271 0.03483 1.8588 1.9953  0.000e+00
#> E[y(0)]       0.9076 0.03607 0.8370 0.9783 9.391e-140
#> ───────────                                          
#> (Intercept)   1.0194 0.04879 0.9238 1.1150  6.127e-97
# MAR analysis
fit <- cate(cate.model = ~1,
            response.model = y ~ a * w2,
            treatment.model = a ~ 1,
            missing.model  = ~ a * (w1 + w2),
            data = d, nfolds = 1L)
fit
#>             Estimate Std.Err   2.5% 97.5%    P-value
#> E[y(1)]       1.9807 0.03372 1.9146 2.047  0.000e+00
#> E[y(0)]       0.9725 0.03237 0.9091 1.036 2.485e-198
#> ───────────                                         
#> (Intercept)   1.0082 0.04573 0.9186 1.098 1.033e-107

# missingness depending on a post-randomization variable Z
# (doubly-robust estimation via `auxiliary.model`)

sim_aux <- function(n = 5000, seed = 1) {
  set.seed(seed)
  w <- rnorm(n)
  a <- rbinom(n, 1, 0.5)
  z <- rnorm(n, mean = 0.5 * a, sd=0.5)
  r <- rbinom(n, 1, plogis(-0.5 - 0.5 * a * z + 0.1 * w))
  y0 <- 1 + 0.5 * a + 2 * z^2 + w + rnorm(n, sd=0.5)
  y <- ifelse(r == 1, y0, NA_real_)
  data.frame(w, a, z, y0, y)
}
d <- sim_aux()
# Truth: E[Y(1)] = 1 + 0.5 + 2 * (1 + 0.25) = 4, E[Y(0)] = 3, ATE = 1
cate(y ~ a + w, a ~ 1, data=na.omit(d))
#>             Estimate Std.Err   2.5%  97.5%   P-value
#> E[y(1)]       2.3849 0.04900 2.2889 2.4810 0.000e+00
#> E[y(0)]       1.5743 0.03687 1.5021 1.6466 0.000e+00
#> ───────────                                         
#> (Intercept)   0.8106 0.05051 0.7116 0.9096 5.817e-58
fit <- cate(response.model = y ~ a + w,
            treatment.model = a ~ 1,
            missing.model = ~ a * z + w,
            auxiliary.model = y ~ a + I(z^2) + w,
            data = d)
fit
#>             Estimate Std.Err   2.5% 97.5%    P-value
#> E[y(1)]       2.4744 0.03405 2.4077 2.541  0.000e+00
#> E[y(0)]       1.5108 0.02668 1.4585 1.563  0.000e+00
#> ───────────                                         
#> (Intercept)   0.9636 0.03814 0.8889 1.038 7.847e-141