Treed distributed lag nonlinear models

Before introducing our proposed method we briefly recap the DLNM framework and standard methodology. Let $y_{i}$ be the continuous outcome for person $i$ from a sample $i=1,\dots,n$. Let $\mathbf{x}_{i}=[x_{i1},\ldots,x_{iT}]^{T}$ denote a vector of exposures observed at equally spaced times $t=1,\ldots,T$. In our case, $y_{i}$ indicates BWGAZ while $x_{it}$ represents the $i^{th}$ mother’s exposure to PM_2.5 in week $t$ of pregnancy. We control for a vector of covariates, denoted $\mathbf{z}_{i}$. The Gaussian DLNM model is

$$y_{i}\sim\mathcal{N}\left[f(\mathbf{x}_{i})+\mathbf{z}_{i}^{T}\boldsymbol{\gamma},\sigma^{2}\right],$$

(1)

where $f(\mathbf{x}_{i})$ is the distributed lag function and $\boldsymbol{\gamma}$ is a vector of regression coefficients.

The distributed lag function $f(\mathbf{x}_{i})$ can take several linear as well as nonlinear forms. The DLNM allows a unique nonlinear association between exposures at each time point and the outcome. In general, the distributed lag function is defined as

$$f(\mathbf{x}_{i})=f(x_{i1},\ldots,x_{iT})=\sum_{t=1}^{T}w(x_{it},t)$$

(2)

where $w(x,t)$ is the exposure-response function relating exposure at week $t$ of gestation to the outcome. Existing frameworks for the DLNM (Gasparrini et al.,, 2010, 2017) utilize a cross-basis where $w$ is represented as a bivariate basis expansion in the exposure concentration and time dimensions. Penalized spline implementations allow for a range of assumptions to be made regarding the structure of the exposure-time-response. For example, varying ridge penalties target shrinkage at specific times, while varying difference penalties control the smoothness along the curve. Basis expansion methods, such as splines, regularize the model to improve stability of the estimated effect in the presence of multicollinearity in the predictor. However, these methods also impose the assumption of smoothness in the DLNM.

We introduce a sum-of-trees model based on the BART framework (Chipman et al.,, 2010) to estimate the exposure-time-response function, $f(\mathbf{x}_{i})$. The general approach is to build dichotomous trees that partition the time-varying exposure $\mathbf{x}_{i}$ in both the exposure concentration and time dimensions. Figure 1 illustrates the approach for a single tree. Figure 1(a) is a diagram of a tree showing binary rules defined on the exposure and time values. These rules divide the exposure-time space into five terminal nodes, denoted $\eta_{1},\ldots,\eta_{5}$. Figure 1(b) shows the exposure-time space partitioned into five regions with each region corresponding to a single terminal node. A tree and corresponding parameters define a piecewise constant exposure-response function,

$$w(x_{it},t)=\mu_{b}\quad\text{if }x_{it}\in\eta_{b}.$$

(3)

The distributed lag function for tree $\mathcal{T}$ takes a form similar to that in (2) and is defined as

$$g(\mathbf{x}_{i},\mathcal{T})=\sum_{t=1}^{T}w(x_{it},t).$$

(4)

In our TDLNM framework, we consider an ensemble of $A$ regression trees. For tree $\mathcal{T}_{a}$, $a\in\{1,\ldots,A\}$, denote the $B_{a}$ terminal nodes as $\eta_{a1},\ldots,\eta_{aB_{a}}$. Each terminal node, $\eta_{ab}$ has a corresponding set of limits in time and exposure concentration given by the rules defined as splits of the tree and a corresponding parameter $\mu_{ab}$. Collectively, the terminal nodes of each tree define a partition of the exposure-time-space and allows for flexible estimation of the exposure-time surface. As in (3), we define the effect of each exposure-time combination in tree $\mathcal{T}_{a}$ to be $w_{a}(x_{it},t)=\mu_{ab}$ if $x_{it}\in\eta_{ab}$. Each regression tree in the ensemble provides a partial estimate of the distributed lag nonlinear function $f$. Formally, the exposure-time-response function for TDLNM is

$$f(\mathbf{x}_{i})=\sum_{a=1}^{A}g(\mathbf{x}_{i},\mathcal{T}_{a}),$$

(5)

where $g(\mathbf{x}_{i},\mathcal{T}_{a})$ represents the partial estimate contributed by tree $a$ given in (4).

TDLNM foregoes the basis-imposed smoothness assumption. However, when different time and exposure breaks are staggered across trees the ensemble of trees can approximate smooth functions. Model regularization is a result of the tree prior, which prefers trees having only a few splits. Smaller trees ensure that the model is stable in the presence of temporal correlation because each terminal node averages across multiple time points.

Most epidemiological studies assume that the exposure-response relationship is smooth in exposure concentration. The TDLNM methods presented above assume a piecewise linear structure that can approximate a smooth function, but it is never truly smooth. In this subsection we propose a TDLNM model that is truly smooth in exposure (TDLNMse). Importantly, TDLNMse does not force smoothness in time to allow for accurate critical window estimation.

To allow for smoothing in the exposure-response, we introduce a weight function on the terminal-node specific effects. A similar idea was introduced by Linero and Yang, (2018), which assigned a node-specific probability to each observation using a gating function at each dichotomous split on a covariate. TDLNMse differs in that we desire smoothing only in the exposure-concentration dimension. To accomplish this, we define smoothing parameter $\sigma_{x}$ and modify (3) to be

$$w_{a}\left(x_{it},t\right)=\sum_{b=1}^{B_{a}}\mu_{ab}\cdot\psi(x_{it};\eta_{ab},\sigma_{x}).$$

(6)

The weight function $\psi(x_{it};\eta_{ab},\sigma_{x})$ allows each observation $x_{it}$ to be distributed across all terminal nodes that contain time point $t$. For the weight function we use a normal kernel with bandwidth $\sigma_{x}$. Hence, the weight for $x_{it}$ assigned to node $\eta_{ab}$ is

$$\psi(x_{it};\eta_{ab},\sigma_{x})=\left\{\Phi\left(\frac{\lceil x_{ab}\rceil-x_{it}}{\sigma_{x}}\right)-\Phi\left(\frac{\lfloor x_{ab}\rfloor-x_{it}}{\sigma_{x}}\right)\right\}\cdot\mathbb{I}(t\in\eta_{ab}),$$

(7)

where $\lceil x_{ab}\rceil$ and $\lfloor x_{ab}\rfloor$ refers to the maximum and minimum exposure concentration limits of node $\eta_{ab}$, respectively, and $\Phi$ is the standard normal cumulative distribution function. The inclusion of the indicator function allows TDLNMse to retain a piecewise constant effect in time at each exposure concentration value. The kernel smoother requires fewer terminal nodes to estimate a smooth effect in exposure-concentration as observations near the boundary of two terminal nodes will have an estimated effect that is in between estimates of observations located centrally in the nodes.

For TDLNMse, we propose fixing the bandwidth $\sigma_{x}$ a priori. Alternatively, we could assign a prior to $\sigma_{x}$ and estimate the bandwidth.

A situation that has not been addressed in the DLNM literature is uncertainty in the exposure. Many exposure models including climate models and spatially kriged exposure models provide measurement of uncertainty. Most commonly these occur in one of two forms—standard errors for the exposure data or multiple realizations from a model such as draws from a posterior predictive distribution or an ensemble method. This uncertainty is not accommodated for in the health effect estimates from standard DLNMs.

Exposure uncertainty can be incorporated into TDLNM by using a weight function to spread the exposure across multiple terminal nodes according to the probability that the exposures is in each of those nodes. The result is similar to TDLNMse, using a weight function corresponding to the uncertainty in each observation. In the case of reported standard errors for the exposure data, we use (7) with observation specific smoothing parameter $\sigma_{xi}$ that is equal to the standard error for each observation. If instead we have multiple draws of exposures from an ensemble or Bayesian model we replace $\boldsymbol{\Phi}$ in (7) with the empirical cumulative distribution function.

To gain some insight into the exposure-response function characterized by TDLNM we consider the DLNM relation at a single time point. The distributed lag function in TDLNM is given by combining (4) and (5), i.e. the sum over trees and the sum of each tree over time. Reversing the order of the summation we get

$$f(\mathbf{x}_{i})=\sum_{t=1}^{T}\sum_{a=1}^{A}w_{a}(x_{it},t).$$

(8)

At time $t$, the exposure-response function, $\sum_{a=1}^{A}w_{a}(x_{it},t)$, is equivalent to the BART model with univariate predictor $x_{it}$. In the case of TDLNM this implies a piecewise-constant exposure-response function across the exposure concentration levels at time $t$. For TDLNMse, the weight function $\psi$ acts as a linear smoother over exposure concentration for the exposure-response function at each time.

We apply a tree-specific, horseshoe-like prior to the effects at the terminal nodes $\mu_{ab}$ (Carvalho et al.,, 2010). The prior for terminal node $b$ on tree $a$ is

$$\mu_{ab}|\sigma^{2},\omega^{2},\tau_{a}^{2}\sim\mathcal{N}\left(0,\sigma^{2}\omega^{2}\tau_{a}^{2}\right).$$

(9)

Here $\tau_{a}\sim\mathcal{C}^{+}(0,1)$ and $\omega\sim\mathcal{C}^{+}(0,1)$ defines the horseshoe prior on trees. We specify prior $\sigma\sim\mathcal{C}^{+}(0,1)$ and $\boldsymbol{\gamma}\sim\mathcal{MVN}(\boldsymbol{0},\sigma^{2}c\mathbf{I})$, where $c$ is fixed at a large value.

For the half-Cauchy priors on all variance parameters we adopt the hierarchical framework of Makalic and Schmidt, (2015), where $r^{2}|s\sim\mathcal{IG}(1/2,1/s)$ and $s\sim\mathcal{IG}(1/2,1)$ gives that marginally $r\sim\mathcal{C}^{+}(0,1)$. This allows for Gibbs sampling of all variance components.

The tree-specific shrinkage prior on $\mu_{ab}$ results in better mixing throughout the MCMC sampler. This occurs by allowing shrunken trees with a small variance component and smaller effects $\mu_{a1},\ldots,\mu_{aB_{a}}$ to more easily explore splitting locations in the exposure-time space. After reconfiguration, these trees have the ability to contribute larger partial estimates.

Our stochastic tree generating process largely follows Chipman et al., (1998). The probability a tree splits at node $\eta$ with depth $d_{\eta}$ equals $p_{\text{split}}(\eta)=\alpha(1+d_{\eta})^{-\beta}$, where hyperpriors $\alpha\in(0,1)$ and $\beta\geq 0$. In our data setting the number of potential splits in the time direction is $T-1$ while the number of potential split points in the exposure direction is equal to the number of unique exposure values minus one, which is substantially larger than $T-1$. To address this imbalance we limit the potential exposure split points a priori and propose an alternative prior on potential split points. By limiting the exposure split points we also avoid situations where a split in one dimension limits future splits in another dimension due to empty nodes. For example, if TDLNM has a tree that splits on an extreme value in the exposure-concentration dimension, it may be unable to further split on the time dimension due to lack of data in one partition of the exposure-time space. We restrict the potential split locations in the exposure dimension to a predefined set of quantiles or values. Specification of potential splitting values also improves computational efficiency by allowing for precalculation of counts or weights for the limited number of potential splits.

We assign prior probabilities uniformly across potential time splits and uniformly across potential exposure splits such that there is a $0.5$ probability of selecting either a time or exposure split as the first splitting rule in a tree. Hence, for $s_{x}$ and $s_{t}$ total potential exposure and time splitting values, respectively, the probability of the first split in a tree being on exposure equals $1/(2s_{x})$ or $1/(2s_{t})$ if the split is on time. For a splitting decision further down the tree, the splitting rule probability is proportional to the probabilities of the potential remaining splits in a selected node. Following a split in time, there are fewer potential remaining splits in time, increasing the probability that the next split will take place in the exposure dimension.

We estimate TDLNM using MCMC. The MCMC approaches used for BART do not apply to the current model for two reasons. First, the algorithm of Chipman et al., (2010) relied on the fact that any specific vector of predictors $\mathbf{x}_{i}$ is contained in a single terminal node on each tree, whereas TDLNM divides the exposures related to each observation across the terminal nodes. Second, we modify the algorithm to allow for parametric control of the confounding variables $\mathbf{z}$. Due to these differences we propose an alternative MCMC approach for the TDLNM model. In particular, we integrate out $\boldsymbol{\gamma}$ using standard analytical techniques. Then, we apply Bayesian backfitting (Hastie and Tibshirani,, 2000) to simultaneously estimate the effects of the partial exposure-time-response based on the partition defined by each tree, $\mathcal{T}_{a}$. Full details of the MCMC sampler can be found in Supplemental Materials Section B.

Tree splitting hyperpriors were set to the defaults used in Chipman et al., (2010) with $\alpha=0.95$ and $\beta=2$; different settings did not improve results. Trees in TDLNM explore only two dimensions, which requires fewer trees to adequately explore the predictor space. In preliminary work, we found that 10 to 20 trees was sufficient. Results did not change using more than 20 trees. We used $A=20$ trees for our simulation and data analysis. We assigned the stochastic tree process to grow or prune with probability 0.3 each and change with probability 0.4. The fixed smoothing parameter in TDLNMse, $\sigma_{x}$, is data dependent: too large and the estimated effect will appear linear, too small and the model reverts to TDLNM (no smooth effect). For our simulation and data analysis, we set $\sigma_{x}$ to half the standard deviation of the exposure data. We found this setting to balance a smooth effect while also clarifying nonlinearity in the exposure-concentration effects.

The distributed lag nonlinear function $f$ includes the model intercept. To ease interpretation we remove the intercept by centering $f$ at a reference exposure value, $x_{0}$, at each time. As $f$ is estimated as a sum of tree, we center each tree at the reference value and we use the centered trees for posterior inference.

TDLNM and TDLNMse used the prior settings described in section 4.3. Thirty evenly spaced values ranging between the 0.01 percentile and the 99.9 percentile of all log-exposure value were designated as potential splits in the exposure-dimension. After a burn-in period of 5,000 iterations, we ran each model for 15,000 iterations, thinning to every tenth draw.

We compare TDLNM and TDLNMse to several spline-based penalized and unpenalized DLNM models. The models are described as follows with more detail given in Gasparrini et al., (2017).

• •

  GAM: base model defined by penalized cubic $B$-spline smoothers of rank 10 in both exposure and time dimensions, with second-order penalties, estimated with REML;

• •

  DLM: using GAM with a linear assumption in exposure concentration;

• •

  GLM-AIC: optimal number of unpenalized, quadratic $B$-splines in both exposure and time dimensions (df 1 to 10) selected by minimizing AIC;

• •

  GAMcr: defined by replacing the cubic $B$-spline basis in GAM with cubic regression splines and penalties on the second derivatives;

• •

  GAM-exp: GAM, replacing the second-order penalties with a varying ridge penalty.

To assess model performance, we center the DLNM for each model at log-exposure value $1$ and evaluate the estimated DLNM over a grid of points.

In each model we include all 10 simulated covariates as well as indicators for year and month to control for the additional seasonal trend. We log-transform the exposure concentration values to reduce skew in the exposure data and allow for equally spaced knots in the spline basis models. The decision to log-transform the response has no impact on TDLNM as the model will produce identical results with or without a log-transform; it does impact the smoothing with TDLNMse.

Summary measures of model performance are shown in Table 1. Here, we compare each model by the root mean square error (RMSE) of the entire exposure-time surface and broken down to the RMSE within and outside the simulated critical windows. We also show the empirical coverage of 95% confidence intervals along with average confidence interval width. In addition, the models are compared on the probability of identifying a non-zero effect across grid points inside the simulated critical window (TP), the probability of incorrectly placing a non-zero effect across grid points outside the simulated critical window (FP), and the precision of correct identification of a non-zero effect: TP/(TP+FP). We designate a non-zero effect in the true exposure-time surface as any effect outside of the interval from $-0.005$ to $0.005$ to account for scenario B and C, which have a non-zero effect everywhere between weeks 11-15 and scenario D, which has a non-zero effect everywhere. Figures 2 and 3 show cross-sections of the exposure-time-response surface using estimates from models TDLNM, TDLNMse, and GAMcr. A non-zero estimate in the plots indicates a change in the response for any observation with that particular time and exposure-concentration value.

TDLNM and TDLNMse have as good or better overall RMSE in scenarios A, C and D. In all scenarios, the tree-based methods have the lowest RMSE in areas of zero effect in the exposure-time-response surface. Figure 2 highlights the ability of TDLNM and TDLNMse to find a sharp distinction between times with or without effects. The shrinkage prior on the tree-specific parameters reduces variance leading to lower RMSE in areas of no effect. In areas of non-zero effect, our models have lower RMSE than spline based models in scenario A and are comparable in scenarios C and D. In scenario B the RMSE in areas of non-zero effect is higher for TDLNM and TDLNMse, as the spline based models do a better job interpolating into the extreme exposure values where few data points reside. Figure 3 contrasts how tree-based models attenuate the effect at the boundaries of exposure values, while GAMcr continues the trend linearly.

The tree-based models have near nominal coverage, except in scenario B. All models show below nominal coverage in scenario B, however, TDLNM and TDLNMse perform best, each having 87% surface coverage. In addition, our models have the smallest average confidence interval width, which is particularly notable at the boundaries in time or extreme exposure concentration where the ‘wiggliness’ of spline-based models becomes more pronounced (Figures 2 and 3). The lack of ‘wiggliness’ in the tree-based model estimates contributes to narrow confidence intervals as well as decreased RMSE, especially in areas of zero-effect. Furthermore, the variation between simulation replicates is much smaller for TDLNM and TDLNMse.

Scenario B, while seemingly natural for a DLM, poses several difficult situations. First, a proper estimate by TDLNM would require trees with many breaks spanning the exposure concentration during the correct critical window. Second, TDLNM attenuates the effect when data is sparse (e.g. high and low concentrations in this scenario). Third, at high concentration, there is a jump from zero to a large effect that smooth methods cannot accommodate; in particular, DLM extends the critical window well beyond the true period of effect as a result of the smoothness assumption.

Precision with TDLNM and TDLNMse is the highest across all simulation scenarios (Table 1). The high precision is a result of near zero FP, but with a tradeoff of lower TP in scenarios B and D. The cross-sectional plots in Figure 2 shows the ability of TDLNM and TDLNMse to adapt to non-smooth exposure-time response surfaces. Supplemental Figure 3 indicates the probability detecting a non-zero effect in at least one exposure value in each week. These results shows that the spline-based methods have a much higher probability of misclassifying weeks just outside of the true critical windows. On the other hand, the tree-based models adapt to changing smoothness in the exposure-time-response surface and rarely detect non-zero effects outside of the true critical window. The key takeaway is that the critical windows detected by TDLNM and TDLNMse have a high probably of being correct.

The posterior mean exposure-time-response estimates for TDLNMse is shown in Figure 4(a). PM_2.5 exposure below the median is associated with an increase in BWGAZ. Exposure concentration above the median value indicate a slight decrease in BWGAZ, but the 95% credible intervals do not give reason to believe this is different from zero. This pattern is present across all gestational weeks. Cross-sections of the exposure-time-response surface at weeks 5, 15, 25, and 35 are shown in Figure 4(b) and indicate a critical window spanning the entire pregnancy.

Based on TDLNMse, a change from median (7.0 $\mu g/m^{3}$) to the 25th percentile of PM_2.5 exposure (5.89 $\mu g/m^{3}$) across the pregnancy would result in a cumulative mean increase in BWGAZ of 0.0132 (95% CI $[0.0003,0.0354]$) or approximately 5.74g (95% CI $[0.11,15.41]$) when translated to actual birth weight (this is approximate because BWGAZ accounts for gestational age and fetal sex). The nonlinear association shows that a further decrease in PM_2.5 exposure to the 10th percentile (5.02 $\mu g/m^{3}$) would result in a 0.055 (95% CI $[0.016,0.090]$) mean increase in BWGAZ, or an approximate increase of 24.1g (95% CI $[7.155,39.10]$). These results suggest that decreasing PM_2.5 below the current national ambient air quality standards would result in higher average birth weights in this population.

The mean exposure-time-response estimate for GAMcr, shown in Figure 4(a), closely resembles the estimates of TDLNMse. As in our simulations, we see a difference in the tail behavior. GAMcr continues the trend in the effect with large intervals. Despite the large point estimate with GAMcr at low exposure levels the larger confidence intervals include zero. In contrast, TDLNMse tapers off and estimates a smaller effect with substantially smaller intervals that do not contain zero. The smaller intervals found in TDLNMse near the boundaries are a result of these boundary regions being grouped in terminal nodes that also contain internal regions and therefore receive the same estimates.

Our findings of an association between increased PM_2.5 and decreased BWGAZ are consistent with previous literature. A meta-analysis by Sun et al., (2016) found a 10 $\mu$g/$m^{3}$ increase in PM_2.5 across pregnancy to be associated with 15.9g decrease in birth weight (95% CI $[-26.8,-5]$); increased exposures in the second and third trimesters were also determined to have a nonzero negative association with birth weight. Zhu et al., (2015) reported similar results in a separate meta-analysis. Strickland et al., (2019) found that the magnitude of associations between PM_2.5 and birth weight increased for higher percentiles of the birth weight distribution across all trimesters. Finally, a study investigating individual chemical components of PM_2.5 found non-zero increased risk of low birth weight for maternal exposures during each trimester of pregnancy (Ebisu and Bell,, 2012).

For comparison, we fit TDLNM, a DLM and several linear models to compare results. Each of these models was consistent with the TDLNMse results. More details on these methods can be found in Supplemental Materials Section D.