\documentclass[notitlepage]{article}
\usepackage{Math}
%\usepackage{../../Math598}
\usepackage{listings}
\usepackage{numprint}
\usepackage{enumerate}
\usepackage{bm}
\usepackage{bbm}
\usepackage{verbatim}
\usepackage{mathtools}
\usepackage{scalerel}
\usepackage{dsfont}

\usepackage{amsfonts,amsmath,amssymb}

\usepackage{tikz}
\usetikzlibrary{shapes,decorations,arrows,calc,arrows.meta,fit,positioning}


\usepackage{pgfplots}
\pgfplotsset{compat=1.18}

\newcommand*\circled[1]{\tikz[baseline=(char.base)]{
            \node[shape=circle,draw,inner sep=2pt] (char) {#1};}}

\usepackage{xcolor}

\newtheorem{alg}{Algorithm}
\newtheorem{ex}{Example}
\newtheorem{note}{Note}
\newenvironment{proof}{\par\noindent{\normalfont{\bf{Proof
}}}}
{\hfill \small$\blacksquare$\\[2mm]}


\setlength{\parindent}{0in}

\def\E{\mathbbmss{E}}
\def\Var{\text{Var}}
\def\Cov{\text{Cov}}
\def\Corr{\text{Corr}}

\def\bW{\mathbf{W}}
\def\bH{\mathbf{H}}

\newcommand{\bs}[1]{\bm{#1}}
\newcommand{\Rho}{\mathrm{P}}
\newcommand{\cif}{\text{ if }}
\newcommand{\Ra}{\Longrightarrow}

\def\bphi{\boldsymbol \phi}


\lstloadlanguages{R}

\definecolor{keywordcolor}{rgb}{0,0.6,0.6}
\definecolor{delimcolor}{rgb}{0.461,0.039,0.102}
\definecolor{Rcommentcolor}{rgb}{0.101,0.043,0.432}

\lstdefinestyle{Rsettings}{
  basicstyle=\ttfamily,
  breaklines=true,
  showstringspaces=false,
  keywords={if, else, function, theFunction, tmp}, % Write as many keywords
  otherkeywords={},
  commentstyle=\itshape\color{Rcommentcolor},
  keywordstyle=\color{keywordcolor},
  moredelim=[s][\color{delimcolor}]{"}{"},
}

\lstset{basicstyle=\ttfamily, numbers=none, literate={~} {$\sim$}{2}}

\def\Xb{\mathbf{X}_{\beta}}
\def\Xp{\mathbf{X}_{\psi}}
\def\beps{\boldsymbol{\epsilon}}
\def\betahat{\widehat \beta}
\def\psihat{\widehat \psi}

\begin{document}

<<setup, cache=FALSE, echo=FALSE>>=
library(knitr)
# global chunk options
opts_chunk$set(cache=FALSE, autodep=TRUE)
options(scipen=6)
options(repos=c(CRAN="https://cloud.r-project.org/"))
knit_hooks$set(document = function(x) {sub('\\usepackage[]{color}', '\\usepackage{xcolor}', x, fixed = TRUE)})
@
\begin{center}
{\textsclarge{AIPW with a Continuous Exposure}}
\end{center}

\tikzset{
    -Latex,auto,node distance =1 cm and 1 cm,semithick,
    state/.style ={circle, draw, minimum width = 0.7 cm},
    box/.style ={rectangle, draw, minimum width = 0.7 cm, fill=lightgray},
    point/.style = {circle, draw, inner sep=0.08cm,fill,node contents={}},
    bidirected/.style={Latex-Latex,dashed},
    el/.style = {inner sep=3pt, align=left, sloped}

}

The inverse probability weighting (IPW) estimation for a continuous treatment uses IPW estimators with weights that target the outcome at (potential) value $z$ of the form
\[
w(x,t) = \frac{\Ind_{(z-d,z+d)}(t)}{f_{Z|X}(t|x)}
\]
or via the stabilized version, that is
\[
w(x,t) = \frac{\Ind_{(z-d,z+d)}(t) f_Z(t)}{f_{Z|X}(t|x)}.
\]
to yield the $\widehat \mu_{\text{IPW}} (z)$ estimator
\[
\widehat \mu_{\text{IPW}} (z) = \sum\limits_{i=1}^n W_{iz} Y_i
\]
with
\[
W_{iz} = \dfrac{w(x_i,z_i)}{\sum\limits_{j=1}^n w(x_j,z_j) }.
\]
To extend this to the augmented (AIPW) estimator, we propose and fit a mean model $m(x,z)$ to form the estimator
\[
\widehat \mu_{\text{AIPW}} (z) = \sum\limits_{i=1}^n W_{iz} (Y_i - m(X_i,Z_i)) + \frac{1}{n} \sum_{i=1}^n m(X_i,z)
\]


\bigskip

\textbf{EXAMPLE:} In this example, we have a data generating process with $k=10$ predictors, joint Normally distributed $X \sim Normal(\bm{\mu},\Sigma)$, with treatment model given by
\[
Z | X = x \sim Normal(\bx_\alpha \alpha, \sigma_Z^2)
\]
and outcome model
\[
Y | X = x, Z=z \sim Normal(\mu(x,z), \sigma_Y^2)
\]
with
\[
\mu(x,z) = \bx_0 \beta + \psi_0 z = \beta_0 + \sum_{l=1}^k x_l \beta_l + \psi_0 z.
\]
In the model, all predictors are predictors of $Z$, but only first three components are confounders.  In the analysis
\[
\alpha = (4,1,1,1,1,1,1,-2,-2,2)^\top \qquad \beta = (2,-2,2.2,3.6,0,0,0,0,0,0)^\top
\]
and $\psi_0 = 1$.  Here, $\sigma_Z = 5$ and $\sigma_Y = 1$.  The quantity $\psi_0$ measures the unconfounded effect of treatment, as, marginally,
\[
\E[Y(z)] = \psi_0 z + \mu_0 \beta.
\]
where $\mu_0 = (1,\bm{\mu}^\top)$.  In the studies below, we consider
\begin{itemize}
  \item IPW with ordinary weights
  \item IPW with stabilized weights
  \item AIPW with ordinary weights
  \item AIPW with stabilized weights
\end{itemize}
where for the AIPW methods, the (mis-specified) mean model
\[
m(x,z) = \beta_0 + \beta_1 x_1 + \psi_0 z
\]
is used.  In each case, the weights are standardized to sum to 1 as in the formula for $W_{iz}$.


\bigskip

We carry out a simulation study of 200 replicates with $n=1000$.



<<A0,comment='+', fig.width=7.25, fig.height=5.0, fig.align='center', warning=F,cache=TRUE,eval=TRUE,echo=FALSE>>=
#Set-up
set.seed(865177)
library(mvnfast)
n<-10000
k<-10
Mu<-sample(-5:5,size=k,rep=T)/2
Sigma<-0.1*0.8^abs(outer(1:k,1:k,'-'))
al<-c(4,rep(1,6),rep(-2,4))
be<-rep(0,k+2)
be[1:5]<-c(2,1,-2.0,2.2,3.6)
sigZ<-5;sigY<-1

#Exploratory sample
X<-rmvn(n,mu=Mu,Sigma)
Xal<-cbind(1,X)
ps.true<-Xal %*% al
Z<-rnorm(n,ps.true,sigZ)

#Find a plotting grid
Z0<-round(mean(Z),0)
zgrid<-seq(Z0-5,Z0+5,by=0.1)
nZ<-length(zgrid)
#Intercept
intercept.val<-be[1]+sum(Mu*be[-c(1:2)])
Xm<-cbind(1,Z,X)
muval<-as.vector(Xm %*% be)
Y<-rnorm(n,muval,sigY)
@

<<A1,comment='+', fig.width=7.25, fig.height=5.0, fig.align='center', warning=F,cache=TRUE,eval=TRUE,echo=FALSE>>=
#Monte Carlo study
n<-1000
nreps<-200
d<-0.05
zhat.samp<-matrix(0,nrow=nreps,ncol=nZ)
zstab.samp<-matrix(0,nrow=nreps,ncol=nZ)
zhat.dr.samp<-matrix(0,nrow=nreps,ncol=nZ)
zstab.dr.samp<-matrix(0,nrow=nreps,ncol=nZ)
for(irep in 1:nreps){
    X<-rmvn(n,mu=Mu,Sigma)
    X1<-X[,1]
    Xal<-cbind(1,X)
    ps.true<-Xal %*% al
    Z<-rnorm(n,ps.true,sigZ)
    Xm<-cbind(1,Z,X)
    muval<-as.vector(Xm %*% be)
    Y<-rnorm(n,muval,sigY)

    fnum<-lm(Z~1)
    fden<-lm(Z~X)
    den<-dnorm(Z,predict(fden),sd(fden$residuals))
    ratio<-dnorm(Z,predict(fnum),sd(fnum$residuals))/den
    
    fitm<-lm(Y~Z+X1)
    mhat<-fitted(fitm)

    zhat<-zstab<-rep(0,nZ)
    zhat.dr<-zstab.dr<-rep(0,nZ)
    for(ix in 1:nZ){
        zl<-zgrid[ix]-d
        zu<-zgrid[ix]+d
        ind<-(Z > zl & Z < zu)
        w<-ind/den
        mval<-mean(predict(fitm,newdata=data.frame(Z=zgrid[ix],X1=X1)))
        zhat[ix]<-sum(w*Y,na.rm=T)/sum(w,na.rm=T)
        zhat.dr[ix]<-sum(w*(Y-mhat),na.rm=T)/sum(w,na.rm=T)+mval
        w<-ind*ratio
        zstab[ix]<-sum(w*Y,na.rm=T)/sum(w,na.rm=T)
        zstab.dr[ix]<-sum(w*(Y-mhat),na.rm=T)/sum(w,na.rm=T)+mval
    }
    zhat.samp[irep,]<-zhat
    zstab.samp[irep,]<-zstab
    zhat.dr.samp[irep,]<-zhat.dr
    zstab.dr.samp[irep,]<-zstab.dr
}
@



<<A3,comment='+', fig.width=7.25, fig.height=4.5, fig.align='center',eval=TRUE,echo=FALSE>>=
par(mar=c(4,5,3,0))
zhat.summary<-apply(zhat.samp,2,quantile,prob=c(0.025,0.50,0.975),na.rm=TRUE)
zhat.dr.summary<-apply(zhat.dr.samp,2,quantile,prob=c(0.025,0.50,0.975),na.rm=TRUE)
plot(zgrid,zhat.summary[2,],type='n',xlab='z',
        ylab=expression(paste(hat(E),"[Y(z)]")),col='red')
lines(zgrid,zhat.summary[1,],lty=2,col='red')
lines(zgrid,zhat.summary[3,],lty=2,col='red')
lines(zgrid,zhat.dr.summary[1,],lty=2,col='cyan')
lines(zgrid,zhat.dr.summary[3,],lty=2,col='cyan')
legend(zgrid[1],28,c('True','IPW','AIPW'),
        col=c('black','red','cyan'),lty=c(1,2,2))
title(expression(paste(hat(E),"[Y(z)]")))
lines(zgrid,intercept.val+be[2]*zgrid)
fit.true<-lm(zhat.summary[2,]~zgrid)
coef(summary(fit.true))
fit.true<-lm(zhat.dr.summary[2,]~zgrid)
coef(summary(fit.true))


zstab.summary<-apply(zstab.samp,2,quantile,prob=c(0.025,0.50,0.975),na.rm=TRUE)
zstab.dr.summary<-apply(zstab.dr.samp,2,quantile,prob=c(0.025,0.50,0.975),na.rm=TRUE)
plot(zgrid,zstab.summary[2,],type='n',xlab='z',
        ylab=expression(paste(hat(E),"[Y(z)]")),col='red')
lines(zgrid,zstab.summary[1,],lty=2,col='red')
lines(zgrid,zstab.summary[3,],lty=2,col='red')
lines(zgrid,zstab.dr.summary[1,],lty=2,col='cyan')
lines(zgrid,zstab.dr.summary[3,],lty=2,col='cyan')
legend(zgrid[1],28,c('True','IPW','AIPW'),
        col=c('black','red','cyan'),lty=c(1,2,2))
title(expression(paste(hat(E),"[Y(z)] - Stabilized")))
lines(zgrid,intercept.val+be[2]*zgrid)


fit.true<-lm(zstab.summary[2,]~zgrid)
coef(summary(fit.true))
fit.true<-lm(zstab.dr.summary[2,]~zgrid)
coef(summary(fit.true))
@

\pagebreak

\textbf{EXAMPLE:} The methods also work for quadratic data generating models
\[
\mu(x,z) = \bx_0 \beta + \psi_0 z + \psi_1 z^2
\]

<<A6,comment='+', fig.width=7.25, fig.height=5.0, fig.align='center', warning=F,cache=TRUE,eval=TRUE,echo=FALSE>>=
#Monte Carlo study
zhat.samp<-zstab.samp<-matrix(0,nrow=nreps,ncol=nZ)
zhat.dr.samp<-matrix(0,nrow=nreps,ncol=nZ)
zstab.dr.samp<-matrix(0,nrow=nreps,ncol=nZ)
for(irep in 1:nreps){
    X<-rmvn(n,mu=Mu,Sigma)
    Xal<-cbind(1,X)
    ps.true<-Xal %*% al
    Z<-rnorm(n,ps.true,sigZ)
    Xm<-cbind(1,Z,X)
    muval<-as.vector(Xm %*% be)
    Y<-rnorm(n,muval+0.1*(Z-Z0)^2,sigY)

    fnum<-lm(Z~1)
    fden<-lm(Z~X)
    den<-dnorm(Z,predict(fden),sd(fden$residuals))
    ratio<-dnorm(Z,predict(fnum),sd(fnum$residuals))/den

    fitm<-lm(Y~Z+X1)
    mhat<-fitted(fitm)

    zhat<-zstab<-rep(0,nZ)
    for(ix in 1:nZ){
        zl<-zgrid[ix]-d
        zu<-zgrid[ix]+d
        ind<-(Z > zl & Z < zu)
        w0<-ind/den
        w1<-ind*ratio
        mval<-mean(predict(fitm,newdata=data.frame(Z=zgrid[ix],X1=X1)))
        zhat[ix]<-sum(w0*Y,na.rm=T)/sum(w0,na.rm=T)
        zstab[ix]<-sum(w1*Y,na.rm=T)/sum(w1,na.rm=T)
        zhat.dr[ix]<-sum(w0*(Y-mhat),na.rm=T)/sum(w0,na.rm=T)+mval
        zstab.dr[ix]<-sum(w1*(Y-mhat),na.rm=T)/sum(w1,na.rm=T)+mval
    }
    zhat.samp[irep,]<-zhat
    zstab.samp[irep,]<-zstab
    zhat.dr.samp[irep,]<-zhat.dr
    zstab.dr.samp[irep,]<-zstab.dr
}
@

<<A7,comment='+', fig.width=7.25, fig.height=4.5, fig.align='center',eval=TRUE,echo=FALSE>>=
par(mar=c(4,5,3,0))
zhat.summary<-apply(zhat.samp,2,quantile,prob=c(0.025,0.50,0.975),na.rm=TRUE)
zhat.dr.summary<-apply(zhat.dr.samp,2,quantile,prob=c(0.025,0.50,0.975),na.rm=TRUE)
plot(zgrid,zhat.summary[2,],type='n',xlab='z',
        ylab=expression(paste(hat(E),"[Y(z)]")),col='red')
lines(zgrid,zhat.summary[1,],lty=2,col='red')
lines(zgrid,zhat.summary[3,],lty=2,col='red')
lines(zgrid,zhat.dr.summary[1,],lty=2,col='cyan')
lines(zgrid,zhat.dr.summary[3,],lty=2,col='cyan')
legend(zgrid[1],28,c('True','IPW','AIPW'),
        col=c('black','red','cyan'),lty=c(1,2,2))
title(expression(paste(hat(E),"[Y(z)]")))
lines(zgrid,intercept.val+be[2]*zgrid+0.1*(zgrid-Z0)^2)
fit.true<-lm(zhat.summary[2,]~zgrid+I(zgrid^2))
coef(summary(fit.true))
fit.true<-lm(zhat.dr.summary[2,]~zgrid+I(zgrid^2))
coef(summary(fit.true))


zstab.summary<-apply(zstab.samp,2,quantile,prob=c(0.025,0.50,0.975),na.rm=TRUE)
zstab.dr.summary<-apply(zstab.dr.samp,2,quantile,prob=c(0.025,0.50,0.975),na.rm=TRUE)
plot(zgrid,zstab.summary[2,],type='n',xlab='z',
        ylab=expression(paste(hat(E),"[Y(z)]")),col='red')
lines(zgrid,zstab.summary[1,],lty=2,col='red')
lines(zgrid,zstab.summary[3,],lty=2,col='red')
lines(zgrid,zstab.dr.summary[1,],lty=2,col='cyan')
lines(zgrid,zstab.dr.summary[3,],lty=2,col='cyan')
legend(zgrid[1],28,c('True','IPW','AIPW'),
        col=c('black','red','cyan'),lty=c(1,2,2))
title(expression(paste(hat(E),"[Y(z)] - Stabilized")))
lines(zgrid,intercept.val+be[2]*zgrid+0.1*(zgrid-Z0)^2)

fit.true<-lm(zstab.summary[2,]~zgrid+I(zgrid^2))
coef(summary(fit.true))
fit.true<-lm(zstab.dr.summary[2,]~zgrid+I(zgrid^2))
coef(summary(fit.true))
@


\end{document}