AppReferences

# Bayesian Nonparametrics for Principal Stratification with Continuous Post-
Treatment Variables

Dafne Zorzetto
Data Science Institute, Brown University, Rhode Island, USA
Antonio Canale
Department of Statistics, University of Padova, Italy
Fabrizia Mealli
Department of Economics, European University Institute, Italy
Francesca Dominici
Department of Biostatistics, Harvard School of Public Health, Massachusetts,
USA
Falco J. Bargagli-Stoffi
Department of Biostatistics, Harvard School of Public Health, Massachusetts,
USA  Correspondence can be sent to<EMAIL_ADDRESS>
###### Abstract

Principal stratification provides a causal inference framework that allows
adjustment for confounded post-treatment variables when comparing treatments.
Although the literature has focused mainly on binary post-treatment variables,
there is a growing interest in principal stratification involving continuous
post-treatment variables. However, characterizing the latent principal strata
with a continuous post-treatment presents a significant challenge, which is
further complicated in observational studies where the treatment is not
randomized. In this paper, we introduce the Confounders-Aware SHared atoms
BAyesian mixture (CASBAH), a novel approach for principal stratification with
continuous post-treatment variables that can be directly applied to
observational studies. CASBAH leverages a dependent Dirichlet
process—utilizing shared atoms across treatment levels—to effectively control
for measured confounders and facilitate information sharing between treatment
groups in the identification of principal strata membership. CASBAH also
offers a comprehensive quantification of uncertainty surrounding the
membership of the principal strata. Through Monte Carlo simulations, we show
that the proposed methodology has excellent performance in characterizing the
latent principal strata and estimating the effects of treatment on post-
treatment variables and outcomes. Finally, CASBAH is applied to a case study
in which we estimate the causal effects of US national air quality regulations
on pollution levels and health outcomes.

Keywords: Causal Inference; Dependent Dirichlet Process; Environmental Health;
Principal Strata; Shared Atoms Mixture Model

##  1 Introduction

Principal stratification provides an effective causal inference approach to
deal with post-treatment variables (Frangakis and Rubin, 2002). A post-
treatment variable is defined as a variable that is potentially affected by
treatment and also affects the primary potential outcomes. Thus, a post-
treatment variable can be seen as an intermediate variable between the
treatment and the outcome of interest. Post-treatment variables may pose
inferential challenges for the estimation of treatment effects but also
provide a useful approach to conceptualizing the mechanisms/pathways through
which the treatment affects the outcome through the intermediate variable.

Principal stratification with respect to the post-treatment variable—also
known as intermediate variable—involves categorizing subjects into different
groups, known as principal strata, based on the joint potential values of that
post-treatment variable at various treatment levels (Antonelli et al., 2023).
Principal strata consist of units that share similar potential values of the
post-treatment—which is unaffected by treatment—thus representing an inherent
latent characteristic of the units.

Since the principal strata are, by definition, not influenced by the
assignment of treatment, comparisons of potential outcomes at different
treatment levels within a principal stratum represent well-defined causal
effects (Rubin, 1974). These specific comparisons are referred to as principal
causal effects. Researchers have often referred to the dissociative effect as
the effect of treatment on the outcome for the stratum that includes units
whose potential post-treatment variables do not notably change between
different treatment levels, while labeling as the associative effect the
effect of the treatment on the outcome for the stratum that includes units
whose post-treatment variable do notably change between the different
treatment levels (Frangakis and Rubin, 2002; Mealli and Mattei, 2012).

Most of the literature on principal stratification and principal causal
effects has focused on the case of binary treatment and the binary post-
treatment variable (see, e.g., Angrist et al., 1996). This implies that
researchers must deal with at most four combinations of these two variables
(see the broad literature for binary or discrete post-treatment
variables—e.g., Cheng and Small, 2006; Imai, 2008; Ding et al., 2011; Mattei
and Mealli, 2011; Mealli and Pacini, 2013; Mealli et al., 2016; Jiang et al.,
2016, 2022; Mattei et al., 2024). Considering continuous post-treatment
variables leads to a potentially infinite number of potential principal
strata, which, in turn, makes inference and interpretation very challenging.
Various solutions have been proposed in the causal inference literature to
deal with continuous post-treatments. Generally, it has been proposed to: (i)
dichotomize the post-treatment variable (Sjölander et al., 2009; Jiang et al.,
2022), or model the continuous post-treatment variable conditional to binary
treatment using (ii) fully parametric or semiparametric modeling (Jin and
Rubin, 2008; Conlon et al., 2014; Lu et al., 2023), or (iii) nonparametric
modeling (Schwartz et al., 2011). Recently, Antonelli et al. (2023) has
proposed an extension of nonparametric modeling in the case of continuous
treatment using a Bayesian methodology.

Although these methods provide potential solutions to the challenges of
continuous post-treatment variables, they often require strong assumptions or
subjective choices. On the one hand, dichotomization of the post-treatment
variable requires the choice of the threshold on the post-treatment variable
(e.g., Zigler et al., 2018; Antonelli et al., 2023; Lu et al., 2023). Although
this choice might be informed by the relevant subject matter literature, it
entails a level of arbitrability and requires extensive checking the
robustness of the results with respect to the chosen thresholds.

The method proposed in this paper extends the exciting work of Schwartz et al.
(2011). The authors model a continuous post-treatment outcomes using a
Dirichlet process mixture (DPM) method (Ferguson, 1973). We build upon it to
allow for a novel and more interpretable characterization of the principal
strata that overcomes some challenges. Firstly, our goal is to foster better
information sharing between the marginal distribution of potential post-
treatment variables, so to enhance the ability to uncover the principal strata
and provide a more interpretable definition of them. Secondly, while the DPM
proposed by Schwartz et al. (2011) adheres to the standard stick-breaking
definition by Sethuraman (1994), our approach extends this by incorporating
covariates into the weights of the mixture model. This adjustment improves the
imputation of missing potential outcomes for the post-treatment, consequently
enhancing the estimation performance of treatment effects.

To do so, we introduce a novel Bayesian approach to principal stratification.
Specifically, we introduce a Confounders-Aware SHared-atoms BAyesian mixture
model (denoted by the acronym CASBAH, with the intentional alteration of the
letter “H") endowed with innovative estimands for principal causal effects.
CASBAH exploits the dependent Dirichlet process (Quintana et al., 2022) to
establish a highly flexible distribution for the potential post-treatment
variables given the confounders and the treatment variable, similar to the
recent proposal of Zorzetto et al. (2024) in the context of heterogeneity in
the treatment effects.

CASBAH improves the current literature on nonparametric modeling for principal
stratification in many respects. Firstly, it allows for the inclusion of the
confounders in the weights of the infinite Bayesian mixture model, which
provides a way to control for measured confounding and a more precise
estimation of the causal effects. Secondly, the proposed model—which leverages
shared random kernel’s atoms—allows us to identify associative/dissociative
strata borrowing information across the treatment levels. Hence, we do not
have to choose subjective criteria to define the principal strata, but CASBAH
can, simultaneously, (i) identify the principal strata, according to the
nature of the observed data, and (ii) impute the missing post-treatment
variables and outcomes. Quantification of the uncertainty of the membership of
the principal strata and estimation of the principal causal effects are
inherent in CASBAH.

This paper is organized as follows. In Section 2, we introduce the framework
and the definition of the novel principal causal estimands. In Section 3, we
propose CASBAH and how its characteristics allow us to identify the principal
strata using a data-driven method. Theoretical prior properties provides
insights on prior elicitation. An efficient Gibbs sampling algorithm to sample
from the posterior distributions and estimate the model is also derived.
Section 4 illustrates, via simulations, the ability of the proposed model to
impute missing post-treatment outcomes, discover the main strata, and estimate
the principal causal effects. In Section 5, we estimate the principal causal
effects of air pollution policies on the mortality rate in the eastern United
States (Zigler et al., 2018), where long-term fine particles pm2.5 are the
post-treatment variable. In Section 6, we conclude the paper with a discussion
and avenues for future work.

##  2 Setup

We assume that we observe $n$ independent units. For each unit
$i\in\\{1,\dots,n\\}$, let
$\mathbf{X}_{i}\in\mathcal{X}\subseteq\mathbb{R}^{P}$ the set of observed
covariates, $T_{i}\in\\{0,1\\}$ the observed (binary) treatment,
$P_{i}\in\mathcal{P}\subseteq\mathbb{R}$ the post-treatment variable, and
$Y_{i}\in\mathcal{Y}\subseteq\mathbb{R}$ the response. According to Rubin’s
Causal Model (Rubin, 1974), for each unit $i$, the potential outcome of the
post-treatment variable is defined as
$\\{P_{i}(0),P_{i}(1)\\}\in{\mathds{R}}^{2}$ and the potential outcome of the
response is defined as $\\{Y_{i}(0),Y_{i}(1)\\}\in{\mathds{R}}^{2}$, for
$i=1,\dots,n$. The vector $\\{P_{i}(0),P_{i}(1)\\}$ represents the collection
of the two potential outcomes of the post-treatment variable, when the unit
$i$ is assigned to the control or the treatment group, respectively.
Similarly, the vector $\\{Y_{i}(0),Y_{i}(1)\\}$ represents the collection of
the two potential outcomes for the response variable, with $Y_{i}(0)$ when the
unit $i$ is assigned to the control group and $Y_{i}(1)$ when assigned to the
treatment group.

###  2.1 Identifying Assumptions

Following the principal stratification framework (Frangakis and Rubin, 2002;
VanderWeele, 2011; Mealli and Mattei, 2012; Feller et al., 2017; Ding et al.,
2017; Lu et al., 2023), the principal strata are not affected by treatment
assignment; therefore, a principal stratification can be used as any unit
classification to define meaningful causal estimates conditional on the
principal strata and discover the heterogeneous treatment effect (Mealli and
Mattei, 2012).

As highlighted in the Introduction, the definition of the dissociative  and
associative strata is of particular interest. Formally, a dissociative stratum
includes units where the treatment does not change the post-treatment
variable, that is, $P_{i}(0)$ is the same as $P_{i}(1)$. In contrast, the
associative stratum includes units for which treatment substantially affects
the post-treatment variable, increasing it—positive associative stratum—or
decreasing it—negative associative stratum.

Throughout the paper, we invoke the following assumptions:

######  Assumption 1 (Stable Unit Treatment Value Assumption)

$\displaystyle Y_{i}(T_{i})=Y_{i}(T_{1},T_{2},\cdots,T_{i},\cdots,T_{n}),\quad
Y_{i}(T_{i})=Y_{i},\mbox{ for }i=1,\dots,n;$ $\displaystyle
P_{i}(T_{i})=P_{i}(T_{1},T_{2},\cdots,T_{i},\cdots,T_{n}),\quad
P_{i}(T_{i})=P_{i},\mbox{ for }i=1,\dots,n.$

The Stable Unit Treatment Value Assumption (SUTVA) enforces that each unit’s
outcome is a function of its treatment only. This is a combination of (i) no
interference between units, that is, the potential outcome and the potential
values of the post-treatment variable of the unit $i$ do not depend on the
treatment applied to other units, and (ii) consistency, that is, no different
versions of the treatment levels assigned to each unit (Rubin, 1986). Under
the principal stratification framework, this assumption is invoked for both
the outcome variable and the post-treatment variable.

In practice, under SUTVA, for $i=1,\dots,n$, we observe only
$p_{i}\in{\mathds{R}}$ and $y_{i}\in{\mathds{R}}$, that is, the realization of
the random variable $P_{i}$ and $Y_{i}$, respectively, and defined as
$P_{i}=(1-T_{i})\cdot P_{i}(0)+T_{i}\cdot P_{i}(1)$ and $Y_{i}=(1-T_{i})\cdot
Y_{i}(0)+T_{i}\cdot Y_{i}(1)$.

######  Assumption 2 (Strongly Ignorable Treatment Assignment)

Given the observed covariates ${x}_{i}$,

$\displaystyle\\{\bigl{(}Y_{i}(1),Y_{i}(0)\bigl{)},\bigl{(}P_{i}(0),P_{i}(1)\bigl{)}\\}\perp\\!\\!\\!\perp
T_{i}\mid X_{i},$ $0<Pr\left(T_{i}=1\mid
X_{i}=x\right)<1\>\>\forall\>x\in\mathcal{X},$

where $\mathcal{X}$ is the features’ space.

This assumption states that: (i) for each unit, the potential outcome and the
potential values of the post-treatment variable are independent of the
treatment conditional on the set of covariates $X_{i}$; (ii) all units have a
positive chance of receiving the treatment.

##  3 Confounders-Aware Shared-atoms Bayesian Mixture Model

Following the Bayesian paradigm and invoking the De Finetti’s theorem (Rubin,
1978), the joint probability distribution of the involved variables is unit
exchangeable and it is defined as

$\Pr(Y(0),Y(1),P(0),P(1),T,X)=\int_{\Theta}\prod_{i=1}^{n}\Pr(Y_{i}(0),Y_{i}(1),P_{i}(0),P_{i}(1),T_{i},X_{i},\theta)p(\theta)d\theta,$

where $p(\theta)$ is the prior distribution for all the involved parameters
$\theta$ that take values in the parameter space $\Theta$. The inner
probability can be factorized as

$\displaystyle\Pr(T_{i}\mid Y_{i}(0),Y_{i}(1),P_{i}(0),P_{i}(1),X_{i},\theta)$
$\displaystyle\times\Pr(Y_{i}(0),Y_{i}(1)\mid P_{i}(0),P_{i}(1),X_{i},\theta)$
$\displaystyle\times\Pr(P_{i}(0),P_{i}(1)\mid X_{i},\theta)$
$\displaystyle\times\Pr(X_{i}\mid\theta).$ (1)

Following the assumptions in the previous section, the conditional probability
of the treatment variable can be written as $\Pr(T_{i}\mid
Y_{i}(0),Y_{i}(1),P_{i}(0),P_{i}(1),X_{i},\theta)=\Pr(T_{i}\mid
X_{i},\theta)$. Moreover, we condition the empirical distribution of
covariates so that $\Pr(X_{i}\mid\theta)=\Pr(X_{i})$.

However, the remaining two probabilities in equation (1), have to be modeled:
(i) the distribution of potential outcomes of the response conditional on the
potential outcome of the post-treatment variable and covariates, and (ii) the
distribution of the potential outcome of the post-treatment variable
conditional on the covariates.

We focus our attention on the latter, that is, the distribution of the
potential outcome of the post-treatment variable conditional on the covariates
for which we propose a novel Bayesian nonparametric approach. As our primary
focus lies in the post-treatment variable distribution, for the sake of
clarity in our discussion, we make the simplifying assumption of employing a
more straightforward parametric model for the conditional response
distribution following the settings of Schwartz et al. (2011). Although the
outcome model is assumed to be a linear regression model, it can also be
generalized to a more complex and flexible model.

For the post-treatment variable, we previously exploited a dependent
nonparametric mixture, following the approach adopted in the previous section
of this paper and more broadly the dependent Dirichlet process (MacEachern,
2000; Barrientos et al., 2012; Quintana et al., 2022). In particular, our
model also shares some similarities with the hierarchical Dirichlet process
(Teh et al., 2004, 2006; Teh and Jordan, 2009) and the more recent common atom
model of Denti et al. (2023).

Specifically, we assume for each $i=1,\dots,n$:

$\displaystyle\\{P_{i}\mid{x}_{i},t\\}$ $\displaystyle\sim
f^{(t)}(\cdot\mid{x}_{i}),\mbox{ for }t=\\{0,1\\},$ $\displaystyle
f^{(t)}(\cdot\mid{x}_{i})$
$\displaystyle=\int_{\Psi}{\mathcal{K}}(\cdot;{x}_{i},\psi)dG_{{x}_{i}}^{(t)}(\psi),$
$\displaystyle\\{G_{{x}_{i}}^{(0)},G_{{x}_{i}}^{(1)}\\}$
$\displaystyle\sim\Pi,$ (2)

where ${\mathcal{K}}(\cdot;x,\psi)$ is a continuous density function, for
every $\psi\in\Psi$, and $G_{{x}_{i}}^{(t)}$ is a random probability measure
depending on the confounders $x_{i}$ associated to an observation assigned to
treatment level $t$. In particular, we have two treatment levels
$t=\\{0,1\\}$, therefore, the distributions belonging to $G_{{x}_{i}}^{(0)}$
and $G_{{x}_{i}}^{(1)}$ are characterized by specific weights assigned to a
common set of atoms induced by the common process $\Pi$.

Following the characterization of the random probability measures we can
write:

$G_{{x}_{i}}^{(t)}=\sum_{l\geq
1}\pi_{l}^{(t)}({x}_{i})\delta_{\mathbb{\psi}_{l}},$ (3)

where the sequences $\\{\pi_{l}^{(t)}(x_{i})\\}_{l\geq 1}$ for $t=\\{0,1\\}$
represent infinite sequences of random weights, and $\\{\psi_{l}\\}_{l\geq 1}$
is an infinite sequence of random kernel’s parameters—independent and
identically distributed from a base measure $H$—shared among potential outcome
of post-treatment variable under both the treatment level $t$.

Furthermore, the sequences of weights are defined through a stick-breaking
representation (Sethuraman, 1994),

$\displaystyle\pi_{l}^{(t)}({x}_{i})$
$\displaystyle=\omega_{l}^{(t)}({x}_{i})\prod_{r<l}\\{1-\omega_{r}^{(t)}({x}_{i})\\};$
(4)

where $\\{\omega_{l}^{(t)}({x})\\}_{l\geq 1}$ , for $t=\\{0,1\\}$ are
$[0,1]$-valued independent stochastic processes.

The discrete nature of the random probability measure $G_{{x}_{i}}^{(t)}$, for
$t=\\{0,1\\}$, allows us to introduce the latent categorical variables
$S_{i}^{(t)}$, for $t=\\{0,1\\}$, describing the groups of units defined by
heterogeneity in $P_{i}(t)$. Assuming
$\mathds{P}\\{S_{i}^{(t)}=l\\}=\pi_{l}^{(t)}({x}_{i})$, we can write model in
(2), exploiting conditioning on $S_{i}^{(t)}$, as

$\\{P_{i}|{x}_{i},t,\psi,S_{i}^{(t)}=l\\}\sim{\mathcal{K}}(\cdot\mid{x}_{i},\psi_{l}),\quad\psi_{l}\sim
H.$

where $\psi$ represents the infinite sequence $\\{\psi_{l}\\}_{l\geq 1}$.

Among the plethora of dependent non-parametric processes (see the recent
review by Quintana et al., 2022, for a detailed description), we focused on
the dependent probit stick breaking process for its success in applications,
its good theoretical properties and ease of computation (Rodriguez and Dunson,
2011). Consistently with this, we specify:

$\displaystyle\omega_{l}^{(t)}(x_{i})=\Phi(\alpha_{l}^{(t)}(x_{i})),\quad\alpha_{l}^{(t)}(x_{i})\sim\mathcal{N}(\beta_{l0}^{(t)}+x_{i}^{T}\beta_{l}^{(t)},1),$
(5)

where $\Phi(\cdot)$ is the CDF of a standard Normal distribution,
$\\{\alpha_{l}^{(t)}(x_{i})\\}_{l\geq 1}$ has Gaussian distributions with mean
a linear combination of the $p$ covariates $x_{i}$, and
$\\{\beta_{l0}^{(t)},\beta_{l}^{(t)}\\}_{l\geq 1}$ are the probit regression
parameters treatment-specific.

As commonly done, we assume the kernel $\mathcal{K}$ to be a Gaussian, so that
model (2)–(3) specifies to

$\displaystyle\\{P_{i}(t)|{x}_{i},S_{i}^{(t)}=l,\eta,\sigma\\}\sim\mathcal{N}(\eta_{l},\sigma_{l}^{2}).$
(6)

where $\eta$ and $\sigma$ represent infinite sequences of location parameters
$\\{\eta_{l}\\}_{l\geq 1}$ and scale parameters $\\{\sigma_{l}\\}_{l\geq 1}$,
respectively, such that $\psi_{l}=(\eta_{l},\sigma_{l})$.

Prior elicitation is completed by assuming for the parameters $\eta_{l}$ and
$\sigma_{l}$ in (6)

$\eta_{l}\stackrel{{\scriptstyle
iid}}{{\sim}}\mathcal{N}(\mu_{\eta},\sigma_{\eta}^{2}),\mbox{ and
}\sigma_{l}^{2}\stackrel{{\scriptstyle
iid}}{{\sim}}\mbox{InvGamma}(\gamma_{1},\gamma_{2}).$

where InvGamma($\gamma_{1},\gamma_{2}$) represents the inverse gamma
distribution with the shape parameter $\gamma_{1}\in\mathbb{R}^{+}$ and the
scale parameter $\gamma_{2}\in\mathbb{R}^{+}$, and mean equal to
$\frac{\gamma_{2}}{\gamma_{1}-1}$ and variance
$\frac{\gamma_{2}^{2}}{(\gamma_{1}-1)^{2}(\gamma_{1}-2)}$.

For the regression parameters in (5), we assume the following multivariate
Gaussian prior:

$\beta^{(t)}\sim\mathcal{N}_{(p+1)(L-1)}(\xi,\Omega)$ (7)

for $t=0,1$ and $l\geq 1$. According to Fasano et al. (2022), the Gaussian
prior leads to a straightforward posterior computation, as discussed in the
next section.

###  3.1 Causal Estimands

In this paper, we introduce a formal and novel definition of principal strata
that takes into account the continuous nature of the post-treatment variable.

We assume that the units can be marginally—respect to the treatment
level—grouped into clusters with finite first moments. Consequently, each unit
$i$ belongs to either: the dissociative stratum if $S_{i}(1)=S_{i}(0)$, which
is equivalent to
$\mathbb{E}_{S_{i}^{(1)}}[P_{i}(1)]=\mathbb{E}_{S_{i}^{(0)}}[P_{i}(0)]$; or
the associative positive stratum if $S_{i}(1)\not=S_{i}(0)$ and
$\mathbb{E}_{S_{i}^{(1)}}[P_{i}(1)]>\mathbb{E}_{S_{i}^{(0)}}[P_{i}(0)]$ (now
on, we indicate these conditions as $S_{i}(1)\succ S_{i}(0)$); or the
associative negative stratum if $S_{i}(1)\not=S_{i}(0)$ and
$\mathbb{E}_{S_{i}^{(1)}}[P_{i}(1)]<\mathbb{E}_{S_{i}^{(0)}}[P_{i}(0)]$ (now
on, we indicate these conditions as $S_{i}(1)\prec S_{i}(0)$).

######  Definition 1 (Causal Estimands)

Conditionally on strata allocation, we can define the causal estimands for the
expected dissociative effect (EDE) and the expected associative effects (EAE)
as

$\displaystyle EDE=\tau_{0}=\mathbb{E}[Y_{i}(1)-Y_{i}(0)\mid
S_{i}(1)=S_{i}(0)],$ $\displaystyle
EAE_{+}=\tau_{+1}=\mathbb{E}[Y_{i}(1)-Y_{i}(0)\mid S_{i}(1)\succ S_{i}(0)],$
$\displaystyle EAE_{-}=\tau_{-1}=\mathbb{E}[Y_{i}(1)-Y_{i}(0)\mid
S_{i}(1)\prec S_{i}(0)].$

These causal estimands are not directly identifiable from the data due to the
conditioning on unobservable quantities (i.e., $S_{i}(1),S_{i}(0)$). However,
they can be weak identified invoking the causal assumptions introduced
above—i.e., SUTVA and strongly ignorable treatment assignment.

###### Theorem 1

Invoking assumptions (1) and (2), the expected value of each potential outcome
can be rewritten as the following

$\displaystyle\mathbb{E}\big{[}Y_{i}(t)\mid
S_{i}(1)=s_{1},S_{i}(0)=s_{0},X_{i}=x\big{]}$
$\displaystyle\quad=\int_{p_{0}p_{1}}\mathbb{E}\big{[}Y_{i}\mid
T_{i}=t,X_{i}=x,P_{i}(1)=p_{1},P_{i}(0)=p_{0}\big{]}$
$\displaystyle\quad\quad\times Pr\big{(}P_{i}(1)=p_{1},P_{i}(0)=p_{0}\mid
T_{i}=t,S_{i}(1)=s_{1},S_{i}(0)=s_{0},X_{i}=x\big{)}dp_{0}p_{1}dx;$ (8)

where the inner expectations $\mathbb{E}\big{[}Y_{i}\mid
T_{i}=t,X_{i}=x,P_{i}(1)=p_{1},P_{i}(0)=p_{0}\big{]}$ is estimate with the
outcome model $Y_{i}\mid T_{i},X_{i},P_{i}(1),P_{i}(0)$, as well the
probability $\Pr\big{(}P_{i}(1)=p_{1},P_{i}(0)=p_{0}\mid
T_{i}=t,S_{i}(1)=s_{1},S_{i}(0)=s_{0},X_{i}=x\big{)}$, that is defined by the
model for the potential post-treatment variables.

The proof is reported in Appendix B.

###  3.2 Discovery of Principal Strata

The definition of the previous (3) for the post-treatment variable model
allows us to characterize two distinct specific mixtures for each of the two
treatment levels $t=\\{0,1\\}$, taking into account the different
heterogeneity in the potential post-treatment variables induced by the
covariates $X$. Simultaneously, the parameters of the shared random kernel
$\\{\psi_{l}\\}_{l\geq 1}$ induce desirable characteristics in the potential
post-treatment variables, such as (i) the potential post-treatment variables
$\\{P(1),P(0)\\}$ are not independent, even conditionally, of the covariates
$X$—this is clear from the posterior distribution reported in Section 3.3 and
Appendix D—and (ii) there is a nonnull probability that any $i$ belongs to the
defined dissociative stratum—see Theorem 2.

The fundamental element in explaining this is the definition of the latent
categorical variables $\\{S_{i}^{(0)},S_{i}^{(1)}\\}$ for each unit
$i\in\\{1,\dots,n\\}$, which describe the probability that each unit will be
allocated in the components of the mixture. The Bayesian nonparametric
mixtures are well-known for their ability to capture and describe the cluster
allocation of the observations; however, we focus on the causal inference goal
to describe the principal strata.

In our fully Bayesian approach, pairs $\\{S_{i}^{(0)},S_{i}^{(1)}\\}$ are
random variables that define the two marginal distributions of the random unit
partition (Quintana, 2006), based on the two potential post-treatment
variables. However, we want to define the join distribution of them,
specifically the unit allocation in the principal strata. This join
distribution is not directly observed from the data, due to missing data
problems, but we can estimate this thanks to the characteristic structure of
the proposed CASBAH model.

The proposed CASBAH model is defined such that the atoms
$\\{\psi_{l}\\}_{l\geq 1}$ are shared between the two potential outcomes of
the post-treatment variable of the same unit. Therefore, the model informs us
about the posterior probability that $\\{S_{i}^{(0)},S_{i}^{(1)}\\}$ is
associated with the same or different atoms, for each unit $i$. For example,
the probability for the unit $i\in\\{1,\dots,n\\}$ to belong to the
dissociative stratum—i.e., $S_{i}(0)=S_{i}(1)$—is equivalent to the
probability that the atom for the control is the same for the treated. Such a
probability is not null and can be derived, a priori, by means of the
following theorem.

###### Theorem 2

Given a unit $i\in\\{1,\dots,n\\}$, with covariates $x_{i}$, following the (2)
distribution, the probability to belong to the dissociative stratum for each
unit $i$ is a function of its covariates $x_{i}$, defined as following

$\Pr[S_{i}(0)=S_{i}(1)]=\frac{\rho_{2}(x_{i})[(1+\rho_{2}(x_{i})-2\rho_{1}(x_{i}))^{L}-1]}{\rho_{2}(x_{i})-2\rho_{1}(x_{i})}$
(9)

where
$\rho_{1}(x_{i})=\mathbb{E}\big{[}\Phi\big{(}\alpha(x_{i})\big{)}\big{]}$ and
$\rho_{2}(x_{i})=\mathbb{E}\big{[}\Phi\big{(}\alpha(x_{i})\big{)}^{2}\big{]}$
and $0\leq\Pr[S_{i}(0)=S_{i}(1)]\leq 1$.

The proof is reported in Appendix C.

The representation of the function of the probability to belong to the
dissociative stratum (9) is reported in Figure 1. Such probability varies
according with the function $\alpha(x)=\beta_{0}+\beta_{1}x$. Therefore, the
probability of belonging to the dissociative stratum depends on the observed
covariates $x_{i}$ for each unit $i$ and also on the posterior distribution of
the parameters $\text{\boldmath$\beta$}=\\{\beta_{0},\beta_{1}\\}$. Where the
$\beta$ posterior distributions are a function of all the observed data, such
that the probability of belonging to the dissociative strata is directly
informed (or driven or characterized) from the data.

Figure 1: Probability to belong to the dissociative stratum. Function varies
according to different values of $\alpha(x)=\beta_{0}+\beta_{1}x$.

###  3.3 Posterior Inference

As Theorem 1 shows, the principal causal estimand are functions of the models
for potential outcomes and the potential post-treatment variable that involve
parameters in both models. Following the Bayesian paradigm, the estimation is
based on the posterior distribution of the parameters conditional on the data.

Sampling from the posterior joint distribution is straightforward via Gibbs
sampling. In particular, the algorithm—described in detail in Section D—takes
inspiration from those proposed by Teh et al. (2004, 2006); Teh and Jordan
(2009), for the hierarchical structure, and by Fasano et al. (2022) for the
probit regression in the weights.

Assuming a general linear regression for the outcome model that depends on the
parameter $\theta_{0}^{Y}$ for the outcome under control and $\theta_{1}^{Y}$
for the treated outcome, the Gibbs sampler interests the following
distributions:

  * •

The latent variable $S_{i}^{(t)}$ is a multivariate distribution with

$\Pr\\{S_{i}^{(t)}=l\\}\propto\pi_{l}^{(t)}({x}_{i})\Pr(P_{i}(t)\mid
S_{i}^{(t)}=l,X_{i};\eta_{l},\sigma_{l}),$

where
$\pi_{l}^{(t)}({x}_{i})=\Phi(\alpha_{l}^{(t)}({x}_{i}))\prod_{r<l}(1-\Phi(\alpha_{r}^{(t)}({x}_{i})))p$
for $l=1,\dots,L-1$ and $\Phi(\alpha_{L}^{(t)}({x}_{i}))=1$.

  * •

The cluster-specific parameters are drawn from

$\psi_{l}\propto\Pi_{x}\prod_{i:S_{i}^{(0)}=l}\Pr(P_{i}(0)\mid
S_{i}^{(0)}=l,X_{i};\psi_{l})\prod_{i:S_{i}^{(1)}=l}\Pr(P_{i}(1)\mid
S_{i}^{(1)}=l,X_{i};\psi_{l}).$

  * •

The logit regression parameters for conditional-dependent weights are drawn
from

$\beta^{(t)}\propto\Pr(\beta^{(t)})\prod_{i=1}^{n}\Pr(P_{i}(t)\mid
S_{i}^{(t)},X_{i};\beta^{(t)})$

following (Fasano et al., 2022) demonstration—detail in Appendix A.

  * •

The missing post-treatment variable is imputed from

$P_{i}(1-T_{i})\propto\Pr(Y_{i}\mid
P_{i}(0),P_{i}(1),X_{i};\theta_{0}^{Y})^{1-T_{i}}\Pr(Y_{i}\mid
P_{i}(0),P_{i}(1),X_{i};\theta_{1}^{Y})^{T_{i}}\Pr(P_{i}(1-T_{i})\mid
X_{i};\psi).$

  * •

The posterior distribution for the outcome model is built from

$\theta_{t}^{Y}\sim\Pr(\theta_{t}^{Y})\prod_{i=1}^{n}\text{P}(Y_{i}(t)\mid
P_{i}(0),P_{i}(1),X_{i};\theta_{t}^{Y}).$

##  4 Simulation Study

The performance of the proposed CASBAH mixture model is evaluated by means of
a simulation study. Our objective is to investigate the model’s ability to (i)
accurately impute the missing post-treatment and outcome variables, that is,
to assess the potential bias of the expected value $E[P_{i}(1)-P_{i}(0)]$ and
$E[Y_{i}(1)-Y_{i}(0)]$ over the units $i=1\dots,n$—, (ii) correctly identify
the principal strata. To achieve this, we conducted simulations under five
different data-generating mechanisms and analyzed the results to understand
the model’s behavior in different scenarios.

Specifically, we assume a linear regression model for the outcome model,
defined as follows:

$\begin{bmatrix}Y(0)\\\
Y(1)\end{bmatrix}\sim\mathcal{N}_{2}\begin{pmatrix}\begin{bmatrix}\theta_{00}+\theta_{01}P(0)\\\
\theta_{10}+\theta_{11}P(1)+\theta_{12}P(0)+\theta_{13}P(0)P(1)\end{bmatrix},\begin{bmatrix}e^{\lambda_{0}}&0\\\
0&e^{\lambda_{0}+\lambda_{1}P(1)}\end{bmatrix}\end{pmatrix}.$ (10)

Assuming as prior distribution for the parameters

$\displaystyle\theta^{(0)}=(\theta_{00},\theta_{01})\sim\mathcal{N}_{2}(\mu_{\theta},\sigma^{2}_{\theta}\mathbf{I}_{2})$
$\displaystyle\mbox{ and
}\theta^{(1)}=(\theta_{10},\theta_{11},\theta_{12},\theta_{13})\sim\mathcal{N}_{4}(\mu_{\theta},\sigma^{2}_{\theta}\mathbf{I}_{4});$
$\displaystyle\lambda_{0}\sim\mathcal{N}(\mu_{\lambda},\sigma^{2}_{\lambda})$
$\displaystyle\mbox{ and
}\lambda_{1}\sim\mathcal{N}(\mu_{\lambda},\sigma^{2}_{\lambda});$

with $\mu_{\theta}=0$, $\sigma_{\theta}=10$, $\mu_{\lambda}=0$, and
$\sigma_{\lambda}=2$. Clearly, our model can accommodate more complex choices
for the outcome, but we decided to opt for a simpler model following what had
already been done in the literature by Schwartz et al. (2011). The performance
of the proposed approach is compared to those obtained with the Schwartz et
al. (2011)’s model (hereafter referred to as SLM).

For the data generating process, we simulate two Bernoulli confounders
$(X_{1},X_{2})$ for scenarios 1 to 4 and five Bernoulli confounders
$(X_{1},X_{2},X_{3},X_{4},X_{5})$ for scenario 5, and a binary treatment
variable such that $T_{i}\sim\mbox{Be}(\mbox{logit}(0.4X_{1i}+0.6X_{2i}))$ for
the scenarios $1-4$ and
$T_{i}\sim\mbox{Be}(\mbox{logit}(0.4X_{1i}+0.6X_{2i}-0.3X_{3i}+0.2X_{4i}X_{5i}))$
for scenario 5, for $i=1,\dots,n$.

Each scenario assumes a different conformation of the strata for the
continuous post-treatment variable
$P_{i}=(P_{i}(0),P_{i}(1))\in\mathbb{R}^{2}$. Each stratum is obtained by
introducing, for each unit, categorical variables $S_{i}(0)$ and $S_{i}(1)$,
for control and treatment levels, respectively, with the vector of
probabilities that depend on the values of the confounders and allocates the
unit in different clusters. Conditionally on the cluster allocation
$S_{i}^{(t)}=s$, with $s$ that has the same support for control and treatment
levels, we simulate both potential post-treatment variables—under control and
under treatment, respectively—as

$P_{i}(0)|S_{i}(0)=s\sim\mathcal{N}(\eta_{s},\sigma^{2}_{s}),\quad
P_{i}(1)|S_{i}(1)=s\sim\mathcal{N}(\eta_{s},\sigma^{2}_{s}).$

Continuous potential outcomes $(Y_{i}(0),Y_{i}(1))$, for $i=1,\dots,n$, are
simulated following the model (10) with common values for $\lambda_{0}=-0.5$
and $\lambda_{1}=0.1$, while $\theta^{(0)}$ and $\theta^{(1)}$ are different
for each scenario. In each setting, the sample size is set to $n=500$. For
each scenario, we simulate $100$ samples. Below we provide further details for
each one of the five scenarios.

Scenario 1: We investigate a situation in which there are two strata: one with
a dissociative effect and one with a positive associative effect. In
particular, for $S_{i}(0)=S_{i}(1)=1$
$P_{i}(0),P_{i}(1)\sim\mathcal{N}(1,0.05)$, and for $S_{i}(0)=2$ and
$S_{i}(1)=3$ $P_{i}(0)\sim\mathcal{N}(2,0.05)$ and
$P_{i}(1)\sim\mathcal{N}(3,0.05)$. The regression parameters for the Y-model
are $\theta^{(0)}=(1,2)$ and $\theta^{(1)}=(1,2,-1,0.5)$.

Scenario 2: We focus on a case where we have a dissociative stratum—for
$S_{i}(0)=S_{i}(1)=1$ both $P(0)$ and $P(1)$ are simulated from
$\mathcal{N}(2,0.05)$—, an associative stratum with a positive effect—for
$S_{i}(0)=2$ and $S_{i}(1)=3$ $P_{i}(0)\sim\mathcal{N}(2,0.05)$ and
$P_{i}(1)\sim\mathcal{N}(3,0.05)$—, and dissociative stratum with a negative
effect—for $S_{i}(0)=2$ and $S_{i}(1)=1$ $P_{i}(0)\sim\mathcal{N}(2,0.05)$ and
$P_{i}(1)\sim\mathcal{N}(1,0.05)$. The regression parameters for Y-model are
$\theta^{(0)}=(1,2)$ and $\theta^{(1)}=(1,1.2,-1,1)$.

Scenario 3: This scenario corresponds to Scenario 1 with closer atoms for the
strata and different variances. In particular, the dissociative stratum has
$(P_{i}(0)|S_{i}(0)=1)=(P_{i}(1)|S_{i}(1)=1)\sim\mathcal{N}(1.5,0.12)$, and
the associative stratum has $(P_{i}(0)|S_{i}(0)=2)\sim\mathcal{N}(2,0.1)$ and
$(P_{i}(1)|S_{i}(1)=3)\sim\mathcal{N}(2.5,0.08)$. The regression parameters
for Y-model are $\theta^{(0)}=(1,2)$ and $\theta^{(1)}=(1,1.2,-0.8,0.5)$.

Scenario 4: This scenario corresponds to Scenario 2 with closer atoms for the
strata and different variances. In particular, the dissociative stratum has
$(P_{i}(0)|S_{i}(0)=1)=(P_{i}(1)|S_{i}(1)=1)\sim\mathcal{N}(1.5,0.12)$, the
associative positive stratum has $(P_{i}(0)|S_{i}(0)=2)\sim\mathcal{N}(2,0.1)$
and $(P_{i}(1)|S_{i}(1)=3)\sim\mathcal{N}(2.5,0.08)$, and the associative
negative stratum has $(P_{i}(0)|S_{i}(0)=2)\sim\mathcal{N}(2,0.1)$ and
$(P_{i}(1)|S_{i}(1)=1)\sim\mathcal{N}(1.5,0.12)$. The regression parameters
for Y-model are $\theta^{(0)}=(1,2)$ and $\theta^{(1)}=(1,1.2,-0.8,0.5)$.

Scenario 5: We investigate the scenario with the three strata when the number
of confounders increases, in particular the treatment variable and cluster
allocation variables that depend on five confounders. The dissociative stratum
has $(P_{i}(0)|S_{i}(0)=1)=(P_{i}(1)|S_{i}(1)=1)\sim\mathcal{N}(2,0.05)$, the
associative positive stratum has
$(P_{i}(0)|S_{i}(0)=3)\sim\mathcal{N}(3,0.05)$ and
$(P_{i}(1)|S_{i}(1)=4)\sim\mathcal{N}(4,0.05)$, and the associative negative
stratum has $(P_{i}(0)|S_{i}(0)=2)\sim\mathcal{N}(2,0.05)$ and
$(P_{i}(1)|S_{i}(1)=1)\sim\mathcal{N}(1,0.05)$. The regression parameters for
Y-model are $\theta^{(0)}=(1,2)$ and $\theta^{(1)}=(1,1.2,-1,0.5)$.

We choose the same hyperparameters for each setting such that the prior is
non-informative and in common for all the settings. For the regression
parameters in (5) and for the parameters $\eta_{l}$ and $\sigma_{l}$ in (6) we
use the following conjugate priors

$\displaystyle\beta^{(t)}\sim\mathcal{N}_{(p+1)(L-1)}(0,20\times\mathbf{1}_{(p+1)(L-1)}),$
$\displaystyle\eta_{l}\sim\mathcal{N}(0,20),\mbox{ and
}\sigma_{l}\sim\mbox{InvGamma}(2,0.5),$

for $t\in\\{0,1\\}$, $l\in\\{1,\dots,20\\}$, and $p$ according with the
covariates considered in different settings, and where $\mathbf{1}_{q}$ is a
diagonal matrix $q\times q$.

Table 1 reports the median and interquartile range (IQR) of the bias for the
expected value of the posterior distribution of sample average of
$P_{i}(1)-P_{i}(0)$ and $Y_{i}(1)-Y_{i}(0)$, for $i\in\\{1,\dots,n\\}$. The
results for the five scenarios for the proposed model CASBAH show a good
ability of the proposed model to impute the missing variables and capture the
true distribution, indeed the medians of the bias are close to zero and the
interquartile range is reasonable according to the simulated variability. The
comparison with the Schwartz et al. (2011)’s model, SLM, underlines the
superiority of CASBAH, which obtains medians of the bias closer to zero and a
more contained IQR for $\mathbb{E}[P_{i}(1)-P_{i}(0)]$ in all scenarios and
almost of them for $\mathbb{E}[Y_{i}(1)-Y_{i}(0)]$.

|  |  | Scenario 1 | Scenario 2 | Scenario 3 | Scenario 4 | Scenario 5
---|---|---|---|---|---|---|---
Post-treatment variable |  |  |  |  |  |
CASBAH | Median |  | -0.0164 | 0.0027 | -0.0318 | 0.0027 | 0.0011
| IQR |  | 0.0315 | 0.0254 | 0.0462 | 0.0325 | 0.0183
SLM | Median |  | -0.0637 | 0.3388 | 0.1254 | 0.3280 | 0.3572
| IQR |  | 0.0989 | 0.0986 | 0.0956 | 0.0963 | 0.1216
Outcome |  |  |  |  |  |
CASBAH | Median |  | 0.0024 | 0.0256 | 0.0067 | 0.0419 | 0.0344
| IQR |  | 0.2235 | 0.1550 | 0.4661 | 0.1357 | 0.1316
SLM | Median |  | -1.9891 | -1.8265 | -1.4363 | -1.3115 | -1.8856
| IQR |  | 0.2546 | 0.2110 | 0.1904 | 0.1836 | 0.3727
Table 1: Median and interquartile range (IQR) of the bias for the expected
value of the posterior distribution of sample average of $P_{i}(1)-P_{i}(0)$
and $Y_{i}(1)-Y_{i}(0)$, for $i\in\\{1,\dots,n\\}$. Values reported for CASBAH
and SLM.

To evaluate the correct identification of the principal strata, we use the
adjusted Rand index (ARI, introduced by Hubert and Arabie (1985)). The ARI
values for the five scenarios are reported in Table 2. For all scenarios, the
index is close to $1$ and confirms that the proposed CASBAH model can
correctly identify the principal strata and, combined with the good missing
data imputation, allows us to estimate the expected associative and
dissociative effects. The SLM model does not identify the principal strata
according to our definition or simultaneously to the model estimation, in
opposition to CASBAH.

|  | Scenario 1 | Scenario 2 | Scenario 3 | Scenario 4 | Scenario 5
---|---|---|---|---|---|---
Mean |  | 0.9850 | 0.9906 | 0.9706 | 0.9717 | 0.9154
SD |  | 0.0997 | 0.0597 | 0.1021 | 0.0892 | 0.1577
Table 2: Adjusted rand index for the five simulated scenarios computed on the
point estimated partitions obtained with the proposed model CASBAH. mean and
empirical standard deviation (sd) are reported.

##  5 Assessing the Impact of Environmental Policies on Particulate Pollution
and Health

As recently as 7 February 2024 the US Environmental Protection Agency (EPA)
has proposed a revision of the National Ambient Air Quality Standards (NAAQS)
(U.S. Environmental Protection Agency, 2024), which will result in a more
stringent environmental policy for the concentration of air pollutants in the
US that cause health hazards. To gather information and better understand the
implementation of this revised policy, it is critically important to evaluate
how past air quality regulations benefitted health outcomes and how these
effects varied between different communities.

In our application, we explore this relevant research question by
investigating the repercussions of a previous revision of NAAQS, leveraging
insights gleaned from historical data and health data. First, we evaluate the
direct and indirect effects of the 2005 NAAQS revision on the mortality rate
in the period 2010-2016, considering the variation of levels pm2.5, under the
principal stratification framework. Second, we take advantage of the flexible
Bayesian non-parametric mixture model for the post-treatment variable (that
is, pm2.5) to understand how these effects can vary within strata. Third, we
characterize the different identified groups of individuals.

###  5.1 Data and Study Design

To answer our research question, we have merged two datasets: (i) one dataset
containing the information about the NAAQS designations, as well as air
pollution levels and the demographic and socioeconomic characteristics in the
counties in the eastern U.S., used in the Zigler et al. (2018)’s analysis
(data can be found at Zigler, 2017), and (ii) one dataset containing the
information on the age-adjusted mortality rate in these counties available by
the Center for Disease Control and Prevention of U.S. (Friede et al., 1993).

The initial dataset from Zigler (2017) is made up of 482 counties in the
eastern U.S., where national monitoring networks have detected the
concentration of pm2.5. In 2005, the EPA designated as non-attainment of the
NAAQS these counties where the average concentration of pm2.5 was above $15\mu
g/m^{3}$, or otherwise attainment (note that the revision occurred in 1997,
but became effective just in 2005 due to legal disputes). States containing
counties designated as non-attainment were required to develop or revise State
Implementation Plans (SIP) outlining how a non-attainment area will achieve
standards with strategies to reduce ambient concentrations of pm2.5.

As noted in the previous literature, there is a large diversity in local
actions in response to the non-attainment designation and in the enactment of
subsequent SIPs (Greenstone, 2004; Zigler et al., 2018). Due to this
diversity, the direct analysis of the effect of the designation is very
similar to an intention-to-treat analysis. This highlights the importance of
conducting a principal stratification analysis to compare the effects on
health outcomes across different locations: specifically, it contrasts the
effects in areas where pollution was effectively reduced by the non-attainment
designation (associative negative effects) with those areas where pollution
was not measurably affected (dissociative effect), or where it actually
increased (associative positive effects).

The study design specifies a baseline period from 2000 to 2005. During this
time, each county, $i$, in the dataset is classified according to its
attainment status under the EPA 2005 NAAQS revisions. We define a binary
treatment variable, $T_{i}$, where $T_{i}=1$ indicates that the i-th county
was designated as non-attainment and therefore was required to develop or
revise its State Implementation Plans. In contrast, $T_{i}=0$ indicates that
the i-th county was designated as attainment and did not have these
requirements.

For this period, we also have the average ambient concentration of pm2.5(from
the EPA monitoring locations within each county), as well as a number of
counfounders such as census variables like: the percentage of Hispanic and
black residents; the average household income; the percentage of females; the
average house value; the proportion of residents in poverty; the proportion of
residents with a high school diploma; the smoking rate; the population; the
percentage of residents in urban area; the employment rate (the percentage of
the workforce employed); the percentage of the move in the last 5 years. We
also have meteorological variables such as the averages of daily temperatures
and the relative average of humidity; the dew point (the temperature at which
air becomes saturated with moisture, leading to the formation of dew or
condensation).

Consistently with Zigler et al. (2018), we identify as the follow-up period a
period after 5 years from the implementation of the 2005 NAAQS revision. For
this period, we have data on the levels of pm2.5(2010-2012) from Zigler
(2017). The levels of pm2.5in the follow-up period are compared with the
levels in the baseline period to establish the decrease / increase in air
pollution following the implementation of NAAQS. This variable serves as the
post-treatment variable $P_{i}$. Furthermore, we gather the publicly available
age-adjusted mortality rate of all-cause mortality (2010-2016) from the
website of the US Center for Disease Control and Prevention (CDC) (Friede et
al., 1993). Mortality rates in the follow-up period are also contrasted with
the rates in the baseline period to assess the impact of NAAQS on mortality
rates accounting for baseline rates. This variable serves as our outcome
variable $Y_{i}$.

The final dataset, obtained by merging the previously mentioned data sources
and cleaning the data set to avoid missing data, comprises 384 counties, 270
of which were designed as attainment and 114 as non-attainment. In Appendix,
Figure F.1 maps this final analysis data set.

###  5.2 Results

We apply the proposed CASBAH model to the 384 counties in the eastern US
described above, including all covariates —- census and meteorological
variables —- in the weights of the post-treatment variable mixture, while for
the outcome model we use the linear model in (10). The model identifies the
three strata: the dissociative stratum with 115 counties ($30\%$ of the total
counties analyzed), the associative positive stratum with 48 counties
($12.5\%$), and the associative negative stratum with 221 counties ($57.5\%$).
According to the definition of the three strata and as visualized in the image
on the left in Figure 2, the dissociative stratum—identified with the color
yellow—is composed of counties where the NAAQS revision does not substantially
affect the level of pm2.5 , in fact, the expected value of
$\mathbb{E}[P_{i}(1)-P_{i}(0)|S_{i}(1)=S_{i}(0)]$ for the counties allocated
to this strata has a median close to zero and the first and third quantiles of
$-0.29\mu g/m^{3}$ and $0.07\mu g/m^{3}$. The associative negative
stratum—identified with the color green—is made up of counties where the
implementation of environmental plans significantly decreases the levels of
pm2.5 . Specifically, in these counties, the NAAQS revision reduced by $0.9\mu
g/m^{3}$ the median levels pm2.5 , with first and third quantiles of $-1.46\mu
g/m^{3}$ and $-0.38\mu g/m^{3}$. The associative positive stratum—identified
with the color red—is made up of counties where the revision of NAAQS was
associated with increases in pm2.5 levels by $0.8\mu g/m^{3}$ in the median.
However, this latter stratum has the highest uncertainty for
$\mathbb{E}[P_{i}(1)-P_{i}(0)|S_{i}(1)\prec S_{i}(0)]$, with values of the
first and third quantiles equal to $0.13\mu g/m^{3}$ and $1.34\mu g/m^{3}$,
respectively.

The corresponding distributions of the expected dissociative/associative
effects are reported in the right image in Figure 2 and show the effect of the
attainment or non-attainment designations on the mortality rate conditional to
the three strata—i.e., conditional to the heterogeneity in the causal effects
of the NAAQS revision on the level of pollution. The EDE and the EAE+ have a
similar increase in the median
$2.7{}^{\text{o}}\mkern-5.0mu/\mkern-3.0mu_{\text{oooo}}$ of the mortality
rate when environmental plans are applied. The means of these two effects are
$0.8{}^{\text{o}}\mkern-5.0mu/\mkern-3.0mu_{\text{oooo}}$ and
$1.3{}^{\text{o}}\mkern-5.0mu/\mkern-3.0mu_{\text{oooo}}$, respectively. For
both principal causal effects, the interquartile interval includes the zero
value. This means that in counties where the NAAQS revision does not affect
the level pm2.5 or where this level increases, the mortality rate is also not
affected. The EAE-—in green color on the right image of Figure 2—assumes
negative values, indicating that the implementation of environmental plans, in
the counties where the regulations affect the level of pm2.5 , significantly
reducing it, also reduces the median age adjusted mortality rate.
Specifically, the mortality rate decreases by
$11.9{}^{\text{o}}\mkern-5.0mu/\mkern-3.0mu_{\text{oooo}}$ in the median and
with an interquartile interval of
$[-15.2{}^{\text{o}}\mkern-5.0mu/\mkern-3.0mu_{\text{oooo}},-9.1{}^{\text{o}}\mkern-5.0mu/\mkern-3.0mu_{\text{oooo}}]$.

Figure 2: Boxplots of the three identified strata for (left) the conditional
average of the difference of the post-treatment variables, and (right) the
expected associative/dissociative effects. The light-blue vertical lines show
the value zero, identifying the null effects.

In addition, it is our interest to characterize the three different strata
(Figure F.3 in the Supplementary Materials visualizes the average of observed
covariates within each stratum—reported as colored lines—and compares them
with the average of covariates between the full 384 counties). The associative
positive stratum and the dissociative stratum are composed of rural counties
with a smaller population density, where the population has lower income with
respect to the mean of the overall counties, and with a small employment rate.
In contrast, in countries of the associative positive stratum, the percentage
of Hispanics is lower than the average, in favor of a higher percentage of
whites and Asians, while the dissociative stratum is characterized by a higher
percentage of Hispanics and a smaller community of black people, with a higher
level of education. Furthermore, the associative positive stratum is
distinguished by a high smoke rate. In contrast, the associative negative
stratum identifies urban counties with a high population density. This
population is characterized by higher income, but also a higher percentage of
poor people and lower levels of education. The meteorological variable, such
as the averages of daily temperatures, the relative average of humidity, and
the dew point, appears to play an important role in the distinction of
strata—that is, in the characterization of the different effects of the
implementation of environmental plans on the level of pm2.5. We refer to the
Supplementary Materials for more details about the probability of each county
to belong in each of the three strata.

##  6 Conclusion

In this paper, we proposed the Confounders-Aware SHared Atoms Bayesian Mixture
Model (CASBAH). CASBAH uses the dependent Dirichlet process to create an
highly flexible and adaptable distribution for the potential post-treatment
variables, taking into account both the confounders and the treatment
variable. This approach extends the recent BNP model for causal inference
proposed by Zorzetto et al. (2024) to the principal stratification settings.
CASBAH also provides an estimate of causal effects for a set of innovative
estimates of principal causal effects.

The proposed model improves nonparametric modeling for principal
stratification, e.g., the model proposed by Schwartz et al. (2011), by
offering a more flexible approach to handling post-treatment variables. CASBAH
does so by integrating the confounders into the weights of an infinite
Bayesian mixture model, improving control over measured confounding, and
enhancing the precision of causal effect estimates. Additionally, by utilizing
shared random kernel atoms, the model facilitates the identification of
associative or dissociative strata without predetermined criteria, effectively
learning from the data across different treatment levels.

CASBAH dynamically identifies the main strata and imputes missing post-
treatment variables and outcomes without the need to set predefined thresholds
on the post-treatment values, as is done in most of the principal
stratification literature in the context of continuous post-treatment.
Moreover, the model quantifies the uncertainty in membership in the principal
strata and estimates the principal causal effects. Theorem 1 provides a useful
factorization of the expected value for each potential outcome, while Theorem
2 illustrates how the probability of a unit belonging to a dissociative
stratum is non-zero and depends on its observed covariates.

The performance of the proposed model is evaluated in a battery of Monte Carlo
simulations where the performance of CASBAH is compared and contrasted with
the performance of previous BNP modeling for principal stratification proposed
by Schwartz et al. (2011). Furthermore, the proposed model is used to assess
the effectiveness of a previous revision of NAAQS in reducing levels of
pm2.5and, in turn, reducing the mortality rate in treated counties.
Specifically, we find significant effects on the reduction of mortality rate
in areas where pollution was effectively reduced by the non-attainment
designation (associative negative effects). In contrast, those areas where
pollution was not measurably affected (dissociative effect) or where it
actually increased (associative positive effects) did not show any significant
effect on the mortality rate.

Future research could extend our methodology by adopting more flexible
approaches to model the final outcome. Although CASBAH currently employs a
linear model for this outcome, future work could move beyond linearity to
embrace more sophisticated modeling techniques. In this contribution, we focus
on innovations in modeling the post-treatment variable and decided to avoid
overcomplicating the outcome model. However, exploring more complex options
for outcome models remains an opportunity for further research. Furthermore,
our methodology could be adapted to accommodate continuous treatments, similar
to the approach suggested by Antonelli et al. (2023). We also leave the
exploration of this extension for future studies.

## References

  * Angrist et al. (1996) Angrist, J. D., G. W. Imbens, and D. B. Rubin (1996).  Identification of causal effects using instrumental variables.  Journal of the American Statistical Association 91(434), 444–455.
  * Antonelli et al. (2023) Antonelli, J., F. Mealli, B. Beck, and A. Mattei (2023).  Principal stratification with continuous treatments and continuous post-treatment variables.  arXiv preprint arXiv:2309.14486.
  * Barrientos et al. (2012) Barrientos, A. F., A. Jara, and F. A. Quintana (2012).  On the support of MacEachern’s dependent Dirichlet Processes and extensions.  Bayesian Analysis 7(2), 277–310.
  * Cheng and Small (2006) Cheng, J. and D. S. Small (2006).  Bounds on causal effects in three-arm trials with non-compliance.  Journal of the Royal Statistical Society Series B: Statistical Methodology 68(5), 815–836.
  * Conlon et al. (2014) Conlon, A. S., J. M. Taylor, and M. R. Elliott (2014).  Surrogacy assessment using principal stratification when surrogate and outcome measures are multivariate normal.  Biostatistics 15(2), 266–283.
  * Denti et al. (2023) Denti, F., F. Camerlenghi, M. Guindani, and A. Mira (2023).  A common atoms model for the Bayesian nonparametric analysis of nested data.  Journal of the American Statistical Association 118(541), 405–416.
  * Ding et al. (2011) Ding, P., Z. Geng, W. Yan, and X.-H. Zhou (2011).  Identifiability and estimation of causal effects by principal stratification with outcomes truncated by death.  Journal of the American Statistical Association 106(496), 1578–1591.
  * Ding et al. (2017) Ding, P., J. Lu, et al. (2017).  Principal stratification analysis using principal scores.  Journal of the Royal Statistical Society Series B 79(3), 757–777.
  * Fasano et al. (2022) Fasano, A., D. Durante, et al. (2022).  A class of conjugate priors for multinomial probit models which includes the multivariate normal one.  Journal of Machine Learning Research 23(30).
  * Feller et al. (2017) Feller, A., F. Mealli, and L. Miratrix (2017).  Principal score methods: Assumptions, extensions, and practical considerations.  Journal of Educational and Behavioral Statistics 42(6), 726–758.
  * Ferguson (1973) Ferguson, T. S. (1973).  A Bayesian analysis of some nonparametric problems.  The Annals of Statistics, 209–230.
  * Frangakis and Rubin (2002) Frangakis, C. E. and D. B. Rubin (2002).  Principal stratification in causal inference.  Biometrics 58(1), 21–29.
  * Friede et al. (1993) Friede, A., J. A. Reid, and H. W. Ory (1993).  CDC WONDER: A comprehensive on-line public health information system of the Centers for Disease Control and Prevention.  American Journal of Public Health 83(9), 1289–1294.
  * Greenstone (2004) Greenstone, M. (2004).  Did the clean air act cause the remarkable decline in sulfur dioxide concentrations?  Journal of Environmental Economics and Management 47(3), 585–611.
  * Hubert and Arabie (1985) Hubert, L. and P. Arabie (1985).  Comparing partitions.  Journal of Classification 2, 193–218.
  * Imai (2008) Imai, K. (2008).  Sharp bounds on the causal effects in randomized experiments with “truncation-by-death”.  Statistics & Probability Letters 78(2), 144–149.
  * Jiang et al. (2016) Jiang, Z., P. Ding, and Z. Geng (2016).  Principal causal effect identification and surrogate end point evaluation by multiple trials.  Journal of the Royal Statistical Society Series B: Statistical Methodology 78(4), 829–848.
  * Jiang et al. (2022) Jiang, Z., S. Yang, and P. Ding (2022).  Multiply robust estimation of causal effects under principal ignorability.  Journal of the Royal Statistical Society Series B: Statistical Methodology 84(4), 1423–1445.
  * Jin and Rubin (2008) Jin, H. and D. B. Rubin (2008).  Principal stratification for causal inference with extended partial compliance.  Journal of the American Statistical Association 103(481), 101–111.
  * Lu et al. (2023) Lu, S., Z. Jiang, and P. Ding (2023).  Principal stratification with continuous post-treatment variables: Nonparametric identification and semiparametric estimation.  arXiv preprint arXiv:2309.12425.
  * MacEachern (2000) MacEachern, S. N. (2000).  Dependent Dirichlet processes.  Technical Report, Department of Statistics, The Ohio State University, Columbus, OH.
  * Mattei et al. (2024) Mattei, A., P. Ding, V. Ballerini, and F. Mealli (2024).  Assessing causal effects in the presence of treatment switching through principal stratification.  Bayesian Analysis 1(1), 1–28.
  * Mattei and Mealli (2011) Mattei, A. and F. Mealli (2011).  Augmented designs to assess principal strata direct effects.  Journal of the Royal Statistical Society Series B: Statistical Methodology 73(5), 729–752.
  * Mealli and Mattei (2012) Mealli, F. and A. Mattei (2012).  A refreshing account of principal stratification.  The International Journal of Biostatistics 8(1).
  * Mealli and Pacini (2013) Mealli, F. and B. Pacini (2013).  Using secondary outcomes to sharpen inference in randomized experiments with noncompliance.  Journal of the American Statistical Association 108(503), 1120–1131.
  * Mealli et al. (2016) Mealli, F., B. Pacini, and E. Stanghellini (2016).  Identification of principal causal effects using additional outcomes in concentration graphs.  Journal of Educational and Behavioral Statistics 41(5), 463–480.
  * Quintana (2006) Quintana, F. A. (2006).  A predictive view of Bayesian clustering.  Journal of Statistical Planning and Inference 136(8), 2407–2429.
  * Quintana et al. (2022) Quintana, F. A., P. Müller, A. Jara, and S. N. MacEachern (2022).  The dependent Dirichlet process and related models.  Statistical Science 37(1), 24–41.
  * Rodriguez and Dunson (2011) Rodriguez, A. and D. B. Dunson (2011).  Nonparametric Bayesian models through probit stick-breaking processes.  Bayesian Analysis 6(1).
  * Rubin (1974) Rubin, D. B. (1974).  Estimating causal effects of treatments in randomized and nonrandomized studies.  Journal of Educational Psychology 66(5), 688.
  * Rubin (1978) Rubin, D. B. (1978).  Bayesian inference for causal effects: The role of randomization.  The Annals of Statistics, 34–58.
  * Rubin (1986) Rubin, D. B. (1986).  Comment: Which ifs have causal answers.  Journal of the American Statistical Association 81(396), 961–962.
  * Schwartz et al. (2011) Schwartz, S. L., F. Li, and F. Mealli (2011).  A Bayesian semiparametric approach to intermediate variables in causal inference.  Journal of the American Statistical Association 106(496), 1331–1344.
  * Sethuraman (1994) Sethuraman, J. (1994).  A constructive definition of Dirichlet priors.  Statistica Sinica, 639–650.
  * Sjölander et al. (2009) Sjölander, A., K. Humphreys, S. Vansteelandt, R. Bellocco, and J. Palmgren (2009).  Sensitivity analysis for principal stratum direct effects, with an application to a study of physical activity and coronary heart disease.  Biometrics 65(2), 514–520.
  * Teh et al. (2004) Teh, Y., M. Jordan, M. Beal, and D. Blei (2004).  Sharing clusters among related groups: Hierarchical Dirichlet processes.  Advances in Neural Information Processing Systems 17.
  * Teh and Jordan (2009) Teh, Y. W. and M. I. Jordan (2009).  Hierarchical Bayesian nonparametric models with applications.  Bayesian Nonparametrics 28(158), 42.
  * Teh et al. (2006) Teh, Y. W., M. I. Jordan, M. J. Beal, and D. M. Blei (2006).  Hierarchical Dirichlet processes.  Journal of the American Statistical Association 101(476), 1566–1581.
  * U.S. Environmental Protection Agency (2024) U.S. Environmental Protection Agency (2024).  PM NAAQS Implementation Fact Sheet.  https://www.epa.gov/system/files/documents/2024-02/pm-naaqs-implementation-fact-sheet.pdf.
  * VanderWeele (2011) VanderWeele, T. J. (2011).  Principal stratification–uses and limitations.  The International Journal of Biostatistics 7(1), 1–14.
  * Zigler (2017) Zigler, C. (2017).  (Partial) Replication Data for: An empirical evaluation of the causal impact of NAAQS nonattainment designations on particulate pollution and health.
  * Zigler et al. (2018) Zigler, C. M., C. Choirat, and F. Dominici (2018).  Impact of National Ambient Air Quality Standards nonattainment designations on particulate pollution and health.  Epidemiology 29(2), 165.
  * Zorzetto et al. (2024) Zorzetto, D., F. J. Bargagli-Stoffi, A. Canale, and F. Dominici (2024).  Confounder-dependent Bayesian mixture model: Characterizing heterogeneity of causal effects in air pollution epidemiology.  Biometrics 80(2).

SUPPLEMENTARY MATERIAL TO
“BAYESIAN NONPARAMETRICS FOR PRINCIPAL STRATIFICATION WITH CONTINUOUS POST-
TREATMENT VARIABLES”
DAFNE ZORZETTO, ANTONIO CANALE, FABRIZIA MEALLI, FRANCESCA DOMINICI, AND FALCO
J. BARGAGLI-STOFFI

##  Appendix A Conjugate prior distribution for logit regression parameters

The choice to assume a SUN distribution arellano2006unification as prior for
regression parameters in the probit regression—Eq. (7)—is due to it is the
conjugate prior for probit regression and consequently, it allows us (i) to
obtain an efficient step in the Gibbs sampler and (ii) to avoid data
augmentation that has usually same drawbacks.

The SUN density distribution for $\beta^{(t)}$ in Eq. (7) is defined, for
$q=(p+1)(L-1)$, as

$\mathds{P}(\beta^{(t)})=\phi_{q}(\beta^{(t)}-\xi;\Omega)\frac{\Phi_{h}(\gamma+\Delta^{T}\bar{\Omega}^{-1}\omega^{-1}(\beta^{(t)}-\xi);\Gamma-\Delta^{T}\bar{\Omega}^{-1}\Delta)}{\Phi_{h}(\gamma;\Gamma)},$
(A.1)

where $\phi_{q}(\beta^{(t)}-\xi;\Omega)$ is a $q$-variate Gaussian
distribution with $\xi$ vector of location parameters and $\Omega$ the
covariance matrix, such that $\Omega=\omega\bar{\Omega}\omega$ where
$\bar{\Omega}$ is the correlation matrix and
$\omega=(\Omega\bigodot\mathbf{1}_{q})^{1/2}$ where $\bigodot$ is the element-
wise Hadamard product.. The second part of the formula introduces a skewness
mechanism, driven by the cumulative distribution function, computed at
$\gamma+\Delta^{T}\bar{\Omega}^{-1}\omega^{-1}(\beta^{(t)}-\xi)\in\mathbb{R}^{h}$
of an $h$-variate Gaussian with mean vector $0$ and $h\times h$ covariance
matrix $\Gamma-\Delta^{T}\bar{\Omega}^{-1}\Delta$. The quantity
$\Phi_{h}(\gamma;\Gamma)$ is the normalizing constant, which coincides with
the cumulative distribution function, evaluated at $\gamma\in\mathbb{R}^{h}$,
of an $h$-variate Gaussian with mean vector $0$ and $h\times h$ covariance
matrix $\Gamma$.

The amount of skewness in the prior is mainly controlled by the $q\times h$
matrix $\Delta$, and when all the entries in $\Delta$ are $0$, the prior for
$\beta_{t}$ coincides with the density of a $q$-variate Gaussian distribution
with $\xi$ vector of location parameters and $\Omega$ the covariance matrix
fasano2022class.

arellano2006unification show that if
$\beta^{(t)}\sim\textsc{SUN}_{q,h}(\xi,\Omega,\Delta,\gamma,\Gamma)$ then

$\displaystyle\beta^{(t)}$ $\displaystyle\stackrel{{\scriptstyle
d}}{{=}}\xi+\omega(B_{0}^{(t)}+\Delta\Gamma^{-1}B_{1}^{(t)}),$ $\displaystyle
B_{0}^{(t)}$
$\displaystyle\sim\mathcal{N}_{q}(0,\bar{\Omega}-\Delta\Gamma^{-1}\Delta^{T}),$
$\displaystyle B_{1}^{(t)}$ $\displaystyle\sim TN_{h}(-\gamma;0,\Gamma),$

where $TN_{h}(-\gamma;0,\Gamma)$ denotes an $h$-variate Gaussian with zero
mean, covariance matrix $\Gamma$ and truncation below $-\gamma$. A simple
mechanism that helps in the simulation of SUN variables.

The multinomial probit distribution for the weights
$\pi^{(t)}=\\{\pi_{l}^{(t)}\\}_{l=1}^{L}$ can be rewritten as

$\mathds{P}(S_{i}(t)=l|\beta^{(t)},X_{i})=\Phi(x_{i}^{T}\beta_{l}^{(t)})\prod_{k=1}^{l-1}[1-\Phi(x_{i}^{T}\beta_{k}^{(t)})]=\prod_{k=1}^{l}\Phi\left((2\bar{s}_{ik}^{(t)}-1)x_{i}^{T}\beta_{k}^{(t)}\right)=\Phi_{l}(x_{i}^{T}\beta^{(t)};\mathbf{I}_{l})$

for $t=\\{0,1\\}$ and $l=1,\dots,L-1$, and where
$x_{i}=(1,x_{i1},\dots,x_{ip}^{T})$ is the vector of the $p$ covariates and
intercept for the unit $i$, $\bar{s}_{i}^{(t)}=(0_{S_{i}(t)-1}^{T},1)^{T}$ if
$S_{i}(t)\leq L-1$ and $\bar{s}_{i}^{(t)}=0_{L-1}$ if $S_{i}(t)=L$, and
$\mathbf{I}_{l}$ refers to the $l\times l$ identity matrix.

Consequently, the probability over the observation $i=1,\dots,n_{t}$ for
$t=\\{0,1\\}$ is

$\mathds{P}(S^{(t)}|\beta^{(t)},X)=\prod_{i=1}^{n^{(t)}}\mathds{P}(S^{(t)}|\beta^{(t)},X_{i})=\Phi_{\bar{n}_{t}}(\bar{X}^{(t)}\beta^{(t)},\mathbf{I}_{\bar{n}^{(t)}})$
(A.2)

where $\bar{n}^{(t)}=n_{1}^{(t)}+\dots+n_{n}^{(t)}$ with
$n_{i}^{(t)}=\min(S_{i}(t),L-1)$, $\bar{X}^{(t)}$ is a
$\bar{n}^{(t)}\times[(p+1)(L-1)]$ matrix with row blocks
$\bar{X}^{(t)}_{[i]}=X_{i}^{(t)}$ and
$X_{i}^{(t)}=(diag(2\bar{s}_{i}^{(t)}-1)\bigotimes
x_{i}^{T},0_{(n_{i}^{(t)}\times[(p+1)(L-1-n_{i}^{(t)})])})$.

Considering with the prior (A.1) and the likelihood (A.2), the posterior
distribution for $\beta^{(t)}$ is

$\mathds{P}(\beta^{(t)}|S^{(t)},X)=\phi_{q}(\beta^{(t)}-\xi;\Omega)\frac{\Phi_{h+{\bar{n}_{t}}}(\gamma_{pst}+\Delta_{pst}^{T}\bar{\Omega}^{-1}\omega^{-1}(\beta^{(t)}-\xi);\Gamma_{pst}-\Delta_{pst}^{T}\bar{\Omega}^{-1}\Delta_{pst})}{\Phi_{h+{\bar{n}_{t}}}(\gamma_{pst};\Gamma_{pst})},$
(A.3)

where $\Delta_{pst}=(\Delta,\bar{\Omega}\omega(\bar{X}^{(t)})^{T}d^{-1})$,
$\gamma_{pst}=(\gamma^{T},\xi^{T}(\bar{X}^{(t)})^{T}d^{-1})$, $\Gamma_{pst}$
is an $(h+{\bar{n}_{t}})\times(h+{\bar{n}_{t}})$ covariance matrix with blocks
$\Gamma_{pst[11]}=\Gamma$,
$\Gamma_{pst[22]}=d^{-1}(\bar{X}^{(t)}\Omega(\bar{X}^{(t)})^{T}+\mathbf{I}_{\bar{n}^{(t)}})d^{-1}$,
and $\Gamma_{pst[12]}=\Gamma_{pst[21]}=d^{-1}\bar{X}^{(t)}\omega\Delta$, where
$d=[(\bar{X}^{(t)}\Omega(\bar{X}^{(t)})^{T}+\mathbf{I}_{\bar{n}^{(t)}})\bigodot\mathbf{I}_{\bar{n}^{(t)}}]^{1/2}$.

In the particular case in which the prior for $\beta^{(t)}$ is a multivariate
Gaussian distribution—i.e.$h=0$—, then the posterior is still the SUN
distribution in Eq. (A.3) with
$\Delta_{pst}=\bar{\Omega}\omega(\bar{X}^{(t)})^{T}d^{-1}$,
$\gamma_{pst}=d^{-1}\bar{X}^{(t)}\xi$, and
$\Gamma_{pst}=d^{-1}(\bar{X}^{(t)}\Omega(\bar{X}^{(t)})^{T}+\mathbf{I}_{\bar{n}^{(t)}})d^{-1}$.

Moreover, a reasonable assumption for $\beta^{(t)}$ prior is the independence
among the $q$ elements, such that $\Omega=\omega^{2}\cdot\mathbf{I}_{q}$—i.e.
the correlation matrix $\bar{\Omega}=\mathbf{I}_{q}$. Therefore, following
again the arellano2006unification’s results, the posterior distribution of
$\beta^{(t)}$ can be drawn from

$\displaystyle\beta^{(t)}$ $\displaystyle\stackrel{{\scriptstyle
d}}{{=}}\xi+\omega(B_{0,pst}^{(t)}+\Delta_{pst}\Gamma_{pst}^{-1}B_{1,pst}^{(t)}),$
$\displaystyle B_{0,pst}^{(t)}$
$\displaystyle\sim\mathcal{N}_{q}(0,\mathbf{I}_{q}-\Delta_{pst}\Gamma_{pst}^{-1}\Delta_{pst}^{T}),$
$\displaystyle B_{1,pst}^{(t)}$ $\displaystyle\sim
TN_{h+{\bar{n}^{(t)}}}(-\gamma_{pst};0,\Gamma_{pst}),$

with $\Delta_{pst}=\omega(\bar{X}^{(t)})^{T}d^{-1}$,
$\gamma_{pst}=d^{-1}\bar{X}^{(t)}\xi$, and
$\Gamma_{pst}=d^{-1}(\omega^{2}\bar{X}^{(t)}(\bar{X}^{(t)})^{T}+\mathbf{I}_{\bar{n}^{(t)}})d^{-1}$
where
$d=[(\omega^{2}\bar{X}^{(t)}(\bar{X}^{(t)})^{T}+\mathbf{I}_{\bar{n}^{(t)}})\bigodot\mathbf{I}_{\bar{n}^{(t)}}]^{1/2}$.

##  Appendix B Proof Theorem 1

Theorem 1 in Section 3.1 wants to rewrite the causal estimands in (1) that are
not identifiable.

For a clearer notation, we use the notation $\mathbb{E}_{A}\mathbb{E}[B\mid
A]$ to indicate $\int_{\mathcal{A}}\mathbb{E}[B\mid A=a]\Pr(A=a)da$ where
$\mathcal{A}$ is the support of variable A. The general causal estimand (1)
can be rewritten as following

$\displaystyle\mathbb{E}\big{[}Y_{i}(1)-Y_{i}(0)\mid
S_{i}(1)=s_{1},S_{i}(0)=s_{0}\big{]}$
$\displaystyle\quad=\mathbb{E}\big{[}Y_{i}(1)\mid
S_{i}(1)=s_{1},S_{i}(0)=s_{0}\big{]}-\mathbb{E}\big{[}Y_{i}(0)\mid
S_{i}(1)=s_{1},S_{i}(0)=s_{0}\big{]}$
$\displaystyle\quad=\mathbb{E}_{X_{i}\mid
S_{i}(1),S_{i}(0)}\mathbb{E}\big{[}Y_{i}(1)\mid
S_{i}(1)=s_{1},S_{i}(0)=s_{0},X_{i}=x\big{]}$
$\displaystyle\quad\quad-\mathbb{E}_{X_{i}\mid
S_{i}(1),S_{i}(0)}\mathbb{E}\big{[}Y_{i}(0)\mid
S_{i}(1)=s_{1},S_{i}(0)=s_{0},X_{i}=x\big{]}$
$\displaystyle\quad=\mathbb{E}_{X_{i}\mid
S_{i}(1),S_{i}(0)}\mathbb{E}\big{[}Y_{i}\mid
T_{i}=1,S_{i}(1)=s_{1},S_{i}(0)=s_{0},X_{i}=x\big{]}$
$\displaystyle\quad\quad-\mathbb{E}_{X_{i}\mid
S_{i}(1),S_{i}(0)}\mathbb{E}\big{[}Y_{i}\mid
T_{i}=0,S_{i}(1)=s_{1},S_{i}(0)=s_{0},X_{i}=x\big{]},$

where the first and second equality invokes the properties of expectation and
the third equality is due to the no unmeasured confounding and SUTVA
assumption. The inner expectation of each integral, for the treatment value
$t\in\\{0,1\\}$, can be decomposed as

$\displaystyle\mathbb{E}\big{[}Y_{i}\mid
T_{i}=t,S_{i}(1)=s_{1},S_{i}(0)=s_{0},X_{i}=x\big{]}$
$\displaystyle\quad=\mathbb{E}_{P_{i}(1),P_{i}(0)\mid
T_{i},S_{i}(1),S_{i}(0),X_{i}}\mathbb{E}\big{[}Y_{i}\mid
T_{i}=t,S_{i}(1)=s_{1},S_{i}(0)=s_{0},X_{i}=x,P_{i}(1)=p_{1},P_{i}(0)=p_{0}\big{]}$
$\displaystyle\quad=\mathbb{E}_{P_{i}(1),P_{i}(0)\mid
T_{i},S_{i}(1),S_{i}(0),X_{i}}\mathbb{E}\big{[}Y_{i}\mid
T_{i}=t,X_{i}=x,P_{i}(1)=p_{1},P_{i}(0)=p_{0}\big{]},$

where the second equality invokes the property that the outcome $Y_{i}$ is
independent to the latent cluster-allocation $(S_{i}(0),S_{i}(1))$ given the
potential post-treatment variable $(P_{i}(0),P_{i}(1))$, by definition of the
post-treatment variable model.

Therefore, we can rewritten it explicitly, as reported in Theorem 1,

$\displaystyle\mathbb{E}\big{[}Y_{i}\mid
T_{i}=t,S_{i}(1)=s_{1},S_{i}(0)=s_{0},X_{i}=x\big{]}$
$\displaystyle\quad=\int_{p_{0}p_{1}}\mathbb{E}\big{[}Y_{i}\mid
T_{i}=1,X_{i}=x,P_{i}(1)=p_{1},P_{i}(0)=p_{0}\big{]}$
$\displaystyle\quad\quad\times Pr\big{(}P_{i}(1)=p_{1},P_{i}(0)=p_{0}\mid
T_{i}=t,S_{i}(1)=s_{1},S_{i}(0)=s_{0},X_{i}=x\big{)}dp_{0}p_{1}dx$ (B.1)

The inner expectations $\mathbb{E}\big{[}Y_{i}\mid
T_{i}=t,X_{i}=x,P_{i}(1)=p_{1},P_{i}(0)=p_{0}\big{]}$ is estimate with the
outcome model $Y_{i}\mid T_{i},X_{i},P_{i}(1),P_{i}(0)$, as well the
probability $\Pr\big{(}P_{i}(1)=p_{1},P_{i}(0)=p_{0}\mid
T_{i}=t,S_{i}(1)=s_{1},S_{i}(0)=s_{0},X_{i}=x\big{)}$, that is defined by the
model for the potential post-treatment variables.

##  Appendix C Proof Dissociative Stratum Probability

As defined in Section 3.1, a unit $i$ is allocated in the dissociative stratum
when the two latent cluster allocation variables have same value, i.e.,
$S_{i}(0)=S_{i}(1)$. Given the prior distribution
$G^{(0)}_{x_{i}},G^{(1)}_{x_{i}}$, the probability of unit $i$ to be allocated
in the dissociative stratum is the following:

$\displaystyle\Pr[S_{i}(0)=S_{i}(1)]$
$\displaystyle=\mathbb{E}\big{[}\Pr[S_{i}(0)=S_{i}(1)\mid
G^{(0)}_{x_{i}},G^{(1)}_{x_{i}}]\big{]}$
$\displaystyle=\mathbb{E}\bigg{[}\sum_{l=1}^{L}\pi_{l}^{(0)}(x_{i})\pi_{l}^{(1)}(x_{i})\bigg{]}$
$\displaystyle=\sum_{l=1}^{L}\mathbb{E}\big{[}\pi_{l}^{(0)}(x_{i})\pi_{l}^{(1)}(x_{i})\big{]}$
$\displaystyle=\sum_{l=1}^{L}\mathbb{E}\big{[}\pi_{l}(x_{i})^{2}\big{]},$
(C.1)

where the first and third equalities invokes the properties of expectation and
the second equality invokes the definition of the latent variables
$S_{i}^{(t)}$, for $t=\\{0,1\\}$, and the independence of the sequence of the
random weights $\\{\pi^{(0)}(x_{i})\\}_{l\geq 1}$ and
$\\{\pi^{(1)}(x_{i})\\}_{l\geq 1}$. Moreover the two sequence of the random
weights are equivalent in distribution such that we can write
$\pi_{l}^{(0)}(x)\stackrel{{\scriptstyle
d}}{{=}}\pi_{l}^{(1)}(x)\stackrel{{\scriptstyle d}}{{=}}\pi_{l}(x)$, notation
used in the fourth equality.

By the definition of the probit stick-breaking, the expected value of the
$l$-th weight squared is the follows

$\displaystyle\mathbb{E}\big{[}\pi_{l}(x_{i})^{2}\big{]}$
$\displaystyle=\mathbb{E}\bigg{[}\Phi\big{(}\alpha_{l}(x_{i})\big{)}^{2}\prod_{r<l}\bigg{(}1-\Phi\big{(}\alpha_{r}(x_{i})\big{)}\bigg{)}^{2}\bigg{]}$
$\displaystyle=\mathbb{E}\bigg{[}\Phi\big{(}\alpha_{l}(x_{i})\big{)}^{2}\prod_{r<l}\bigg{(}1-\Phi\big{(}\alpha_{r}(x_{i}))^{2}-2\Phi\big{(}\alpha_{r}(x_{i})\big{)}\bigg{)}\bigg{]}$
$\displaystyle=\mathbb{E}\big{[}\Phi\big{(}\alpha_{l}(x_{i})\big{)}^{2}\big{]}\prod_{r<l}\bigg{(}1+\mathbb{E}\big{[}\Phi\big{(}\alpha_{r}(x_{i}))^{2}\big{]}-2\mathbb{E}\big{[}\Phi\big{(}\alpha_{r}(x_{i})\big{)}\big{]}\bigg{)}\bigg{]}$

For simplify the notation we will indicate
$\rho_{1}(x_{i})=\mathbb{E}\big{[}\Phi\big{(}\alpha_{l}(x_{i})\big{)}\big{]}$
and
$\rho_{2}(x_{i})=\mathbb{E}\big{[}\Phi\big{(}\alpha_{l}(x_{i})\big{)}^{2}\big{]}$,
for each $l\in\\{1,\dots,L\\}$. Then (C.1) is rewritten as

$\displaystyle\Pr[S_{i}(0)=S_{i}(1)]$
$\displaystyle=\sum_{l=1}^{L}\mathbb{E}\big{[}\pi_{l}(x_{i})^{2}\big{]}$
$\displaystyle=\sum_{l=1}^{L}\bigg{[}\rho_{2}(x_{i})\prod_{r<l}\big{(}1+\rho_{2}(x_{i})-2\rho_{1}(x_{i})\big{)}\bigg{]}$
$\displaystyle=\sum_{l=1}^{L}\bigg{[}\rho_{2}(x_{i})(1+\rho_{2}(x_{i})-2\rho_{1}(x_{i}))^{l-1}\bigg{]}$
$\displaystyle=\frac{\rho_{2}(x_{i})[(1+\rho_{2}(x_{i})-2\rho_{1}(x_{i}))^{L}-1]}{\rho_{2}(x_{i})-2\rho_{1}(x_{i})},$

where the third and fourth equalities are following by properties of
mathematical series, and it is reported as Theorem 2.

##  Appendix D Posterior Computation

rodriguez2011nonparametric proves that the finite truncation of the dependent
Probit Stick-Breaking process is a good approximation; therefore, we can
rewrite (3) as a finite mixture to $L<\infty$ components with $L$ a reasonable
conservative upper bound. rodriguez2011nonparametric’s proof is a key point
that allows us to provide a simpler algorithm without losing the robustness of
the model.

In this section, we describe the Gibbs sampling algorithm for model fitting
that allows us to draw from the posterior distribution. Following the steps in
the algorithm 1, in each iteration $r=1,\dots,R$, we use the observed data
$(y,p,t,x)$ to update the parameters and the augmented variables and impute
the missing post-treatment variable $P^{mis}$ and missing outcome $Y^{mis}$.

The Gibbs sampling algorithm is divided into three parts: the estimation of
the shared atoms mixture model for the post-treatment variables (divided in
the estimation of cluster allocation, cluster-specific parameters, and
confounder-dependent weights), the imputation of the missing post-treatment
variables, and the estimation of the outcome model.

As already discussed, the outcome model is not our main concern, therefore we
assume a linear model. In particular, in the following Gibbs sampler, we
consider the outcome model used in the simulation study—i.e., (10)—where only
the potential post-treatment variables are included. This choice is driven by
the purpose of focusing attention on the essential definition of the relation
between the post-treatment variable and outcome, which is crucial to impute
the missing post-treatment variable $P^{mis}$. However, the algorithm can be
easily modified to include the covariates $X$ in the linear regression or to
consider a more complex and flexible model.

Cluster Allocation. The latent variables $S_{i}(t)$ identifies the cluster
allocation for each units $i\in\\{1,\dots,n\\}$ at the treatment level $t$.
Its posterior distribution is a multinomial distribution where

$\Pr\\{S_{i}(t)=l\\}\propto\pi_{l}^{(t)}({x}_{i})\mathcal{N}(p_{i};\eta_{l},\sigma_{l}^{2}),$

for $i=1,\dots,n$ and $l=1,\dots,L$, with $\pi^{(t)}_{l}$ defined as:

$\pi_{l}^{(t)}({x}_{i})=\Phi(\alpha_{l}^{(t)}({x}_{i}))\prod_{r<l}(1-\Phi(\alpha_{r}^{(t)}({x}_{i}))),$

for $l=1,\dots,L-1$ and with $\Phi(\alpha_{L}^{(t)}({x}_{i}))=1$.

Cluster Specific Parameters. Thanks to the latent variables
$\\{S_{i}(0),S_{i}(1)\\}$, that cluster the units by the value of their
outcome, we know for each cluster $l\in\\{1,\dots,L\\}$, the allocated units
and we can update the values of the parameters from their posterior
distributions:

$\displaystyle\eta_{l}$
$\displaystyle\sim\mathcal{N}\left(V_{l}^{-1}\times\left(\frac{\sum_{\\{i:S_{i}(0)=l\\}}p_{i}(0)+\sum_{\\{i:S_{i}(1)=l\\}}p_{i}(1)}{\sigma_{l}^{2}}+\frac{\mu_{\eta}}{\sigma^{2}_{\eta}}\right),V_{l}^{-1}\right);$
$\displaystyle\sigma_{l}^{2}$
$\displaystyle\sim\mbox{InvGamma}\left(\gamma_{1}+\frac{n_{l}}{2},\gamma_{2}+\frac{\sum_{\\{i:S_{i}(0)=l\\}}(p_{i}(0)-\eta_{l})^{2}+\sum_{\\{i:S_{i}(1)=l\\}}(p_{i}(1)-\eta_{l})^{2}}{2}\right);$

for $l=1,\dots,L$ and where $V_{l}=n_{l}/\sigma_{l}^{2}+1/\sigma^{2}_{\eta}$
and $n_{l}$ is the number of units allocated in the $l$-th cluster.

Confounder-Dependent Weights. The
$\\{\beta_{ql}^{(t)}\\}_{q=0}^{p}=(\beta_{0l}^{(t)},\beta_{l}^{(t)})$, for
$l=1,\dots,\max(S_{i}(t),L-1)$, are updated for the posterior distribution:

$\displaystyle\beta^{(t)}$ $\displaystyle\stackrel{{\scriptstyle
d}}{{=}}\xi+\omega(B_{0,pst}^{(t)}+\Delta_{pst}\Gamma_{pst}^{-1}B_{1,pst}^{(t)}),$
$\displaystyle B_{0,pst}^{(t)}$
$\displaystyle\sim\mathcal{N}_{q}(0,\mathbf{I}_{q}-\Delta_{pst}\Gamma_{pst}^{-1}\Delta_{pst}^{T}),$
$\displaystyle B_{1,pst}^{(t)}$ $\displaystyle\sim
TN_{h+{\bar{n}^{(t)}}}(-\gamma_{pst};0,\Gamma_{pst}),$

with $\Delta_{pst}=\omega(\bar{X}^{(t)})^{T}d^{-1}$,
$\gamma_{pst}=d^{-1}\bar{X}^{(t)}\xi$, and
$\Gamma_{pst}=d^{-1}(\omega^{2}\bar{X}^{(t)}(\bar{X}^{(t)})^{T}+\mathbf{I}_{\bar{n}^{(t)}})d^{-1}$
where
$d=[(\omega^{2}\bar{X}^{(t)}(\bar{X}^{(t)})^{T}+\mathbf{I}_{\bar{n}^{(t)}})\bigodot\mathbf{I}_{\bar{n}^{(t)}}]^{1/2}$.
More details for $\bar{X}^{(t)}$ and $\omega$ definitions in Section A.

Imputation Missing Post-Treatment Variables. For each unit
$i\in\\{1,\dots,n\\}$, we impute the missing post-treatment variable
$P_{i}^{mis}$. Firstly, drawing the relative cluster-allocation variable
$S_{i}(1-t)$—where $t$ is the observed treatment of the unit $i$—from a
multinomial distribution with

$\mathbb{P}\\{S_{i}(1-t)=l\\}\propto\pi_{l}^{(1-t)}({x}_{i}),$

for $l=1,\dots,L$. Where
$\pi_{l}^{(1-t)}({x}_{i})=\Phi(\alpha_{l}^{(1-t)}({x}_{i}))\prod_{r<l}(1-\Phi(\alpha_{r}^{(1-t)}({x}_{i})))$,
for $l=1,\dots,L-1$ and with $\Phi(\alpha_{L}^{(1-t)}({x}_{i}))=1$.

Successively, drawing the missing post-treatment variable $P_{i}^{mis}$,
conditioned to the allocation to the cluster $l$ and the observed outcome
variables $Y_{i}(t)$. For each $i$ such that the observed treatment level is
$T=1$, $P_{i}(1-t)$ is drown from

$\\{P_{i}^{mis}|S_{i}(1-t)=l,\eta,\sigma^{2},P_{i},Y_{i}\\}\sim\mathcal{N}\left(v^{-1}\left(\frac{\eta_{l}}{\sigma^{2}_{l}}+\frac{m_{1}}{v_{1}}\right),v^{-1}\right);$

where

$v=\frac{1}{\sigma^{2}_{l}}+\frac{1}{v_{1}},\quad
m_{1}=\frac{Y_{i}(1)-\theta_{10}-\theta_{11}P_{i}(1)}{\theta_{12}+\theta_{13}P_{i}(1)},\quad
v_{1}=\frac{e^{\lambda_{0}+\lambda_{1}P_{i}(1)}}{(\theta_{12}+\theta_{13}P_{i}(1))^{2}}.$

While for each $i$ such that the observed treatment level is $T=0$,
$P_{i}(1-t)$ is drown from

$\\{P_{i}^{mis}|S_{i}(1-t)=l,\eta,\sigma^{2},P_{i},Y_{i}\\}\sim\mathcal{N}(\eta_{l},\sigma^{2}_{l}).$

Outcome Model. The $\theta^{(t)}$ parameters are independent for the treatment
level $t$, therefor the posterior distributions are, respectively for $t=0,1$:

$\displaystyle\theta^{(t)}$
$\displaystyle\sim\mathcal{N}_{q^{(t)}}((V^{(t)})^{-1}M^{(t)},V^{(t)})^{-1});$
$\displaystyle M^{(t)}$
$\displaystyle=(\tilde{P}^{(t)})^{T}\Phi^{(t)}\tilde{P}^{(t)}+(\sigma^{2}_{\theta})^{-1}\mathbf{I}_{q^{(t)}};$
$\displaystyle V^{(t)}$
$\displaystyle=(\tilde{P}^{(t)})^{T}\Phi^{(t)}Y(t)+\frac{\mu_{\theta}}{\sigma^{2}_{\theta}}.$

For the treatment level $t=0$: $q^{(0)}=2$; $\tilde{P}^{(0)}$ is a matrix
$n_{0}\times n_{0}$ such that $\tilde{P}^{(0)}=[1_{n_{0}},P(0)]$ with
$1_{n_{0}}$ a vector of $1$ and $P(0)$ the vector of observed values of post-
treatment variable for the units $n_{0}$ assigned at the control group—i.e.
$t=0$; and $\Phi^{(0)}$ is a diagonal matrix $n_{0}\times n_{0}$ with value
$\exp(\lambda_{0})$ in the diagonal. In similar way, for the treatment level
$t=1$: $q^{(1)}=4$; $\tilde{P}^{(1)}$ is a matrix $n_{1}\times n_{1}$ such
that $\tilde{P}^{(1)}=[1_{n_{0}},P(1),P(0),P(1)\cdot P(0)]$ with $1_{n_{1}}$ a
vector of $1$, $P(1)$ the vector of observed values of post-treatment variable
for the units $n_{1}$ assigned at the treated group—i.e. $t=0$—and $P(0)$ the
vector of imputed values of post-treatment variable; and $\Phi^{(1)}$ is a
diagonal matrix $n_{1}\times n_{1}$ with the values
$\exp(\lambda_{0}+\lambda_{1}P(1))$ in the diagonal.

The parameters in the variance of the $Y$-model, $\lambda_{0}$ and
$\lambda_{1}$, do not have conjugate priors, therefore a independent
Metropolis proposal step is necessary. At each iteration
$r\in\\{1,\dots,R\\}$, $\lambda_{0}^{*}$ and $\lambda_{1}^{*}$ are drown from
the proposal distribution
$\mathcal{N}(\mu_{\lambda_{0}},\sigma^{2}_{\lambda_{0}})$ and
$\mathcal{N}(\mu_{\lambda_{1}},\sigma^{2}_{\lambda_{1}})$ respectively. Then,
at iteration $r$ the value of the parameter are updated as following:
$\lambda_{0}^{(r)}=\lambda_{0}^{*}$ with probability

$\frac{\prod_{i\in
n}\mathcal{N}(Y_{i}|\mu_{Y}^{(t)},\exp(\lambda_{0}^{*}+\mathbb{I}_{(T_{i}=1)}\lambda_{1}^{(r-1)}P_{i}(1))}{\prod_{i\in
n}\mathcal{N}(Y_{i}|\mu_{Y}^{(t)},\exp(\lambda_{0}^{(r-1)}+\mathbb{I}_{(T_{i}=1)}\lambda_{1}^{(r-1)}P_{i}(1))},$

otherwise $\lambda_{0}^{(r)}=\lambda_{0}^{(r-1)}$; and
$\lambda_{1}^{(r)}=\lambda_{1}^{*}$ with probability

$\frac{\prod_{i\in
n_{1}}\mathcal{N}(Y_{i}|\mu_{Y}^{(t)},\exp(\lambda_{0}^{(r-1)}+\lambda_{1}^{*}P_{i}(1))}{\prod_{i\in
n_{1}}\mathcal{N}(Y_{i}|\mu_{Y}^{(t)},\exp(\lambda_{0}^{(r-1)}+\lambda_{1}^{(r-1)}P_{i}(1))},$

otherwise $\lambda_{1}^{(r)}=\lambda_{1}^{(r-1)}$; where
$\mu_{Y}^{(0)}=\theta_{00}+\theta_{01}P_{i}(0)$ and
$\mu_{Y}^{(1)}=\theta_{10}+\theta_{11}P_{i}(1)+\theta_{12}P_{i}(0)+\theta_{13}P_{i}(0)P_{i}(1)$.

Algorithm 1 Confounders-Aware Shared-atoms Bayesian Mixture Model

Inputs:

\- the observed data $(y,p,t,x)$.

Outputs:

\- posterior distributions of parameters: $\eta$, $\sigma$, $\beta$, $\theta$,
and $\lambda$;

\- imputed values for $P^{mis}$;

\- posterior distribution over the space of partitions of the units.

Procedure:
Initialization of all parameters and latent variables.
For $r\in\\{1,\dots,R\\}$ :

$\xrightarrow{}$ Estimation of Shared Atoms Mixture Model:

Compute $\omega^{(t)}(x_{i})$ for $i=1,\dots,n$ and $t=0,1$;

Draw $S_{i}(t)$ for $i=1,\dots,n$ and $t=0,1$;

Draw $\eta$ and $\sigma$;

Compute $\alpha^{(t)}({x}_{i})$ for $i=1,\dots,n$ and $t=0,1$;

Draw $\beta^{(t)}$ for$t=0,1$.

$\xrightarrow{}$ Imputation of Missing Post-Treatment Variables:

Draw $P_{i}^{mis}$ for $i=1,\dots,n$ and $t=0,1$.

$\xrightarrow{}$ Estimation of Outcome Model:

Draw $\theta^{(t)}$ for $t=0,1$;

Draw $\lambda_{0}$ and $\lambda_{1}$.
End

##  Appendix E More simulations details

Figure E.1 reports the distribution of the simulated post-treatment variables
and outcomes for the five simulated scenarios. The complexity of the
distributions increases over the simulated scenarios, in particular in
scenarios 4 and 5, it is not unidentifiable the groups in the distributions of
the potential post-treatment variables and the potential outcomes.

Figure E.1: Distributions of (right) the post-treatment variables and (left)
the outcomes, for the five simulated scenarios. In light blue the distribution
of the variables given the treatment and in yellow given the control.

Figure E.2 reports the results for the five simulated scenarios obtained with
CASBAH. On the left, the boxplots show the distribution over the simulated
samples of the expected values of the difference of the post-treatment
variables under treatment and control in each stratum. The graphics confirm
the ability of our proposed model to (i) identify correctly the number of
strata—two strata in the simulated scenario 1 and 3, and three in the
others—and (ii) capture the definition of the associative/dissociative strata
without an a priori criteria—the dissociative stratum is always around zero
for $\mathbb{E}[P(1)-P(0)]$, while the dissociative stratum does not include
the zero. In the boxplots on the right, there is the distribution of the
principal causal effect: EAE-, EDE, and EAE+. The different strata identify
different treatment effects on the outcome, allowing us to characterize the
heterogeneity in the causal effects. Few outliers are observed, however, they
are found in particular in Scenario 5 which describes a more complex relation
among variables and strata.

Figure E.2: Representations of the five simulated scenarios. (Left) The
expected value of the difference of the post-treatment variables under
treatment and control given the strata allocation. (Right) The expected
associative/dissociative causal effects. In green and indicated with $V=-1$,
the associative negative stratum—i.e. corresponding to $S_{i}(1)\prec
S_{i}(0)$—, in yellow and indicated with $V=0$, the dissociative stratum—i.e.,
$S_{i}(1)=S_{i}(0)$—, and in red and indicated with $V=+1$ the associative
positive stratum—i.e., $S_{i}(1)\succ S_{i}(0)$.

##  Appendix F More application details and results

Figure F.1 maps of the final analysis dataset described in Section 5.1 On the
top map in Figure F.1, the red points visualize the treated counties, which
appear to be closer to the main cities, such as Chicago, New York City,
Washington DC or Cleveland. The continuous post-treatment variable is the
difference in level pm2.5 between the follow-up period (2010-2016) and the
baseline period (2000-2005), reported on the map to the left of the bottom of
Figure F.1. The outcome variable is defined as the difference between the age-
adjusted mortality rate between the follow-up period and the baseline period,
and is visualized in the map at the bottom right of Figure F.1. As reported in
the maps, both the post-treatment variable and the outcome have almost all
negative values, highlighting a general trend to decrease the level of pm2.5
and a reduction in mortality rates in the last decade.

Figure F.1: Considered counties in the Easter U.S. (top) Attainment of the air
pollution standard ($15\mu m$ for pm2.5) and consequent application of air
pollution regulations: $0$ if the county was in attainment, $1$ if the county
was in non-attainment. (bottom left) Difference between the average long-term
exposure pm2.5 between baseline and the follow-up period (measured in $\mu
m$). (bottom right) Difference of the age-adjusted mortality rate between
baseline and the follow-up period, value per 100,000.

In addition to the analysis reported in Section 5.2, our proposed approach
allows us to quantify the uncertainty of strata allocation. In fact, for each
county, we know the probability of being allocated in each of the three
strata, in addition to the estimation of its allocation. In addition, we can
visualize this information on the US map. Specifically, the first three maps
in Figure F.2 visualize the probability that each county is assigned to the
three different strata. As already underlined in Figure F.3, counties with a
higher probability of being allocated in the associative negative stratum and
the dissociative stratum are far from the largest cities, different from the
associative positive stratum. In addition, western countries seem to have a
small probability of being assigned to the associative negative stratum. The
fourth map—bottom right in Figure F.2—reports the estimation of the partition
point of the strata, a partition that is used to estimate the principal causal
effects.

Figure F.2: Considered counties in the Easter U.S.. (top left) Probability to
be allocated in the associative positive stratum. (top right) Probability to
be allocated in the dissociative stratum. (bottom left) Probability to be
allocated in the associative negative stratum. (bottom right) Point estimation
of strata allocation.

Figure F.3: Representation of the characteristics of the identified strata.
Each spider plot reports in the colored area the strata-specific
characteristics—the mean of the analyzed covariates—and in the gray area the
collective characteristics—the mean of the covariates among all the analyzed
counties in the Eastern U.S.. We can consider the gray area as the benchmark
to understand how the characteristics of each stratum differ from the
collective characteristics of the analyzed population.

dcu biblio