\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 = 1$. Here, $\sigma_Z = 2$ and $\sigma_Y = 10$.  

\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}

}

\begin{figure}[h]
\centering
\begin{tikzpicture}[scale=2]
    % x node set with absolute coordinates
    \node[state] (z) at (0,0) {${Z}$};
    \node[state] (zsq) at (1,1) {${Z^2}$};
    \node[state] (y) at (2,0){${Y}$};
    \node[state] (x) at (1,-2) {${X}$};
    \node[state] (b) at (0.5,-1) {$\scriptstyle b_1(X)$};

    \path (x) edge (y);
    \path (x) edge[line width=1pt, double distance=1pt, -{Classical TikZ Rightarrow[length=1.5mm]}] (b);
    \path (z) edge (y);
    \path (z) edge[line width=1pt, double distance=1pt, -{Classical TikZ Rightarrow[length=1.5mm]}] (zsq);
    \path (b) edge (z);
    \path (zsq) edge (y);

 \end{tikzpicture}
\end{figure}

<<A0,comment='+', fig.width=7.25, fig.height=5.0, fig.align='center', warning=F,cache=FALSE,eval=TRUE,echo=FALSE>>=
#Set-up
set.seed(898177)
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[2:5]<-c(2,-2.0,2.2,3.6)
psi<-c(2,1)
sigZ<-2;sigY<-10
@

\bigskip

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.
\]
We have that the ATE curve is
\[
\E[Y(z)-Y(0)] = \psi_0 z + \psi_1 z^2.
\]
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]
\qquad \qquad
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=FALSE,eval=TRUE,echo=FALSE>>=
#Monte Carlo study
n<-10000
nreps<-2000
gest.samp<-matrix(0,ncol=2,nrow=nreps)
gest2.samp<-matrix(0,ncol=2,nrow=nreps)
or.samp<-matrix(0,ncol=2,nrow=nreps)
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,X)
    muval<-as.vector(Xm %*% be)+Z*psi[1] + Z^2*psi[2]
    Y<-rnorm(n,muval,sigY)


    fit0<-lm(Y~Z+I(Z^2))
    or.samp[irep,]<-coef(fit0)[2:3]
    
    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]
    
    fit2<-lm(Y~Z+I(Z^2)+ps1)
    gest2.samp[irep,]<-coef(fit2)[2:3]

}
@

<<A2,comment='+', fig.width=7.25, fig.height=4.5, fig.align='center', warning=F,cache=FALSE,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)
#hist(Z)
#plot(Z,Y)
pname<-c('OR','G1','G2')
par(mar=c(4,3,3,0))
boxplot(cbind(or.samp[,1],gest.samp[,1],gest2.samp[,1]),names=pname)
abline(h=psi[1],lty=2,col='red')
title(expression(hat(psi)[0]))
boxplot(cbind(or.samp[,2],gest.samp[,2],gest2.samp[,2]),names=pname)
abline(h=psi[2],lty=2,col='red')
title(expression(hat(psi)[1]))
n*apply((cbind(or.samp[,1],gest.samp[,1],gest2.samp[,1])-psi[1])^2,2,mean)
n*apply((cbind(or.samp[,2],gest.samp[,2],gest2.samp[,2])-psi[2])^2,2,mean)
@




\end{document}