\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{Propensity Score Regression with 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}

}

\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 + \psi_1 z^2 = \beta_0 + \sum_{l=1}^k \beta_l x_l + \psi_0 z + \psi_1 z^2.
\]
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
\]
with treatment effect parameters $\psi_0 = 1$, $\psi_1 = 0.2$.

<<A0,comment='+', fig.width=7.25, fig.height=5.0, fig.align='center', warning=F,cache=TRUE,eval=TRUE,echo=TRUE>>=
#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+1)
be[1:5]<-c(2,-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)])
psi<-c(1,0.2)
Xm<-cbind(1,X)
muval<-as.vector(Xm %*% be) + Z*psi[1] + Z^2*psi[2]
Y<-rnorm(n,muval,sigY)
summary(Z)
@

The quantities $\psi_0$ and $\psi_1$ measure the unconfounded effect of treatment, as, marginally,
\[
\E[Y(z)] = \psi_0 z + \psi_1 z^2 + \mu_0 \beta.
\]
where $\mu_0 = (1,\mu^\top)$ so that the intercept is $\mu_0 \beta = \Sexpr{intercept.val}$. Here, $\sigma_Z = 5$ and $\sigma_Y = 1$.  We have that the ATE curve is
\[
\E[Y(z)-Y(0)] = \psi_0 z + \psi_1 z^2.
\]


\pagebreak

Propensity score regression can be carried out using balancing scores that block the backdoor paths.  For up to quadratic terms, we may need to model
\[
b_1(x) = \E_{Z|X}[Z | X = x]
\]
and
\[
b_2(x) = \E_{Z|X}[Z^2 | X = x]
\]
and then fit a model such as
\[
m(x,z) = \beta_0 + \psi_0 z + \psi_1 z^2 + \phi_0 b_1(x) + \phi_1 b_2(x)
\]
and so on.

\medskip

We carry out a simulation study of 2000 replicates with $n=1000$.
<<A1,comment='+', fig.width=7.25, fig.height=4.5, fig.align='center', warning=F,cache=TRUE,eval=TRUE,echo=TRUE>>=
#Monte Carlo study
n<-1000
nreps<-2000
gest.samp<-matrix(0,ncol=2,nrow=nreps)
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,X)
    muval<-as.vector(Xm %*% be)+Z*psi[1] + Z^2*psi[2]
    Y<-rnorm(n,muval,sigY)

    ps1<-fitted(lm(Z~X))
    ps2<-fitted(lm(Z^2~X))

    fit1<-lm(Y~Z+I(Z^2)+ps1+ps2)
    gest.samp[irep,]<-coef(fit1)[2:3]
}
@

<<A2,comment='+', fig.width=7.25, fig.height=4.5, fig.align='center', warning=F,cache=TRUE,eval=TRUE,echo=FALSE>>=
par(mar=c(4,3,3,0),mfrow=c(1,2))
hist(gest.samp[,1],nclass=40,main=expression(psi[0]),xlab='');box()
abline(v=psi[1],col='red',lty=2)
hist(gest.samp[,2],nclass=40,main=expression(psi[1]),xlab='');box()
abline(v=psi[2],col='red',lty=2)
@

Now with only one term in the propensity score:
\[
m(x,z) = \beta_0 + \psi_0 z + \psi_1 z^2 + \phi_0 b_1(x)
\]

<<A3,comment='+', fig.width=7.25, fig.height=4.5, fig.align='center', warning=F,cache=TRUE,eval=TRUE,echo=TRUE>>=
#Monte Carlo study
n<-1000
nreps<-2000
gest2.samp<-matrix(0,ncol=2,nrow=nreps)
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,X)
    muval<-as.vector(Xm %*% be)+Z*psi[1] + Z^2*psi[2]
    Y<-rnorm(n,muval,sigY)

    ps1<-fitted(lm(Z~X))

    fit1<-lm(Y~Z+I(Z^2)+ps1)
    gest2.samp[irep,]<-coef(fit1)[2:3]
}
@

<<A4,comment='+', fig.width=7.25, fig.height=4.5, fig.align='center', warning=F,cache=TRUE,eval=TRUE,echo=FALSE>>=
par(mar=c(4,3,3,0),mfrow=c(1,2))
hist(gest2.samp[,1],nclass=40,main=expression(psi[0]),xlab='');box()
abline(v=psi[1],col='red',lty=2)
hist(gest2.samp[,2],nclass=40,main=expression(psi[1]),xlab='');box()
abline(v=psi[2],col='red',lty=2)
@

That is, we still get unbiased estimation of $\psi_0$ and $\psi_1$ even if we only use the balancing score $b_1(x)$.

\pagebreak

\textbf{EXAMPLE:} We now consider the model with interactions with the first confounder:
\[
\mu(x,z) = \psi_0 z + \psi_1 z^2 + \psi_2 z x_1 + \psi_3 z^2 x_1  + \bx_0 \beta.
\]
We first fit the propensity score model with augmenting term
\[
\phi_0 b_1(x) + \phi_1 b_1(x) x_1 + \phi_2 b_2(x) + \phi_3 b_2(x) x_1
\]
<<A5,comment='+', fig.width=7.25, fig.height=5.0, fig.align='center', warning=F,cache=TRUE,eval=TRUE,echo=TRUE>>=
#Monte Carlo study
n<-1000
nreps<-2000
gest3.samp<-matrix(0,ncol=4,nrow=nreps)
psi<-c(1,0.2,2,-3)

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)
    Zsq<-Z^2
    Xm<-cbind(1,X)
    X1<-X[,1]
    muval<-as.vector(Xm %*% be)+Z*(psi[1]+psi[3]*X1) + Z^2*(psi[2]+psi[4]*X1)
    Y<-rnorm(n,muval,sigY)

    ps1<-fitted(lm(Z~X))
    ps2<-fitted(lm(Z^2~X))

    fit1<-lm(Y~Z+Zsq+Z:X1+Zsq:X1+ps1+ps1:X1+ps2+ps2:X1)
    gest3.samp[irep,]<-coef(fit1)[c(2:3,6:7)]
}
@

<<A6,comment='+', fig.width=7.25, fig.height=4.5, fig.align='center', warning=F,cache=TRUE,eval=TRUE,echo=FALSE>>=
par(mar=c(4,3,3,0),mfrow=c(2,2))
hist(gest3.samp[,1],nclass=40,main=expression(psi[0]),xlab='');box()
abline(v=psi[1],col='red',lty=2)
hist(gest3.samp[,2],nclass=40,main=expression(psi[1]),xlab='');box()
abline(v=psi[2],col='red',lty=2)
hist(gest3.samp[,3],nclass=40,main=expression(psi[2]),xlab='');box()
abline(v=psi[3],col='red',lty=2)
hist(gest3.samp[,4],nclass=40,main=expression(psi[3]),xlab='');box()
abline(v=psi[4],col='red',lty=2)

@

We now fit the propensity score model with augmenting term
\[
\phi_0 b_1(x) + \phi_1 b_1(x) x_1
\]
<<A7,comment='+', fig.width=7.25, fig.height=5.0, fig.align='center', warning=F,cache=TRUE,eval=TRUE,echo=TRUE>>=
#Monte Carlo study
n<-1000
nreps<-2000
gest4.samp<-matrix(0,ncol=4,nrow=nreps)
psi<-c(1,0.2,2,-3)

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)
    Zsq<-Z^2
    Xm<-cbind(1,X)
    X1<-X[,1]
    muval<-as.vector(Xm %*% be)+Z*(psi[1]+psi[3]*X1) + Z^2*(psi[2]+psi[4]*X1)
    Y<-rnorm(n,muval,sigY)

    ps1<-fitted(lm(Z~X))
    ps2<-fitted(lm(Z^2~X))

    fit1<-lm(Y~Z+Zsq+Z:X1+Zsq:X1+ps1+ps1:X1)
    gest4.samp[irep,]<-coef(fit1)[c(2:3,5:6)]
}
@

<<A8,comment='+', fig.width=7.25, fig.height=5, fig.align='center', warning=F,cache=TRUE,eval=TRUE,echo=FALSE>>=
par(mar=c(4,3,3,0),mfrow=c(2,2))
hist(gest4.samp[,1],nclass=40,main=expression(psi[0]),xlab='');box()
abline(v=psi[1],col='red',lty=2)
hist(gest4.samp[,2],nclass=40,main=expression(psi[1]),xlab='');box()
abline(v=psi[2],col='red',lty=2)
hist(gest4.samp[,3],nclass=40,main=expression(psi[2]),xlab='');box()
abline(v=psi[3],col='red',lty=2)
hist(gest4.samp[,4],nclass=40,main=expression(psi[3]),xlab='');box()
abline(v=psi[4],col='red',lty=2)
@

In this case we do get bias as the path via the $Z^2 X_1$ interaction is not blocked by conditioning only on $b_1(X)$.

\end{document}