Accelerating Bayesian inference of dependency between mixed-type biological traits

Abstract

Inferring dependencies between mixed-type biological traits while accounting for evolutionary relationships between specimens is of great scientific interest yet remains infeasible when trait and specimen counts grow large. The state-of-the-art approach uses a phylogenetic multivariate probit model to accommodate binary and continuous traits via a latent variable framework, and utilizes an efficient bouncy particle sampler (BPS) to tackle the computational bottleneck—integrating many latent variables from a high-dimensional truncated normal distribution. This approach breaks down as the number of specimens grows and fails to reliably characterize conditional dependencies between traits. Here, we propose an inference pipeline for phylogenetic probit models that greatly outperforms BPS. The novelty lies in 1) a combination of the recent Zigzag Hamiltonian Monte Carlo (Zigzag-HMC) with linear-time gradient evaluations and 2) a joint sampling scheme for highly correlated latent variables and correlation matrix elements. In an application exploring HIV-1 evolution from 535 viruses, the inference requires joint sampling from an 11,235-dimensional truncated normal and a 24-dimensional covariance matrix. Our method yields a 5-fold speedup compared to BPS and makes it possible to learn partial correlations between candidate viral mutations and virulence. Computational speedup now enables us to tackle even larger problems: we study the evolution of influenza H1N1 glycosylations on around 900 viruses. For broader applicability, we extend the phylogenetic probit model to incorporate categorical traits, and demonstrate its use to study Aquilegia flower and pollinator co-evolution.

Author summary

We aim to learn the relationships between different biological features, or traits, observed in related specimens that have evolved together over time. This is of great scientific interest because by identifying how different traits influence each other, we gain insights into the mechanisms underlying important biological processes. Learning the relationships between traits across a large number of specimens is computationally challenging, particularly when traits have mixed-type values (continuous and discrete). The previous best approach utilizing a method called bouncy particle sampler (BPS) struggles with increasing specimen and trait counts, resulting in unreliable estimates of trait dependencies. We develop a more efficient approach that largely outperforms BPS, reducing the runtime of our large-scale applications from weeks to days. We apply our method to study the evolution of HIV and influenza viruses, as well as flower and pollinator co-evolution. Our work provides an efficient yet general way to understand the connections between mixed-type traits, offering valuable insights into the evolution of complex biological systems.

Introduction

An essential goal in evolutionary biology is to understand the across-trait covariation observed within biological samples, or taxa, ranging from plants and animals to microorganisms and pathogens such as human immunodeficiency virus (HIV) and influenza. This task is difficult because taxa are implicitly correlated through their shared evolutionary history often described with a reconstructed phylogenetic tree. Here, tree tips correspond to the taxa themselves, and internal nodes are their unobserved ancestors. Inferring across-trait covariation requires a highly structured model that can explicitly describe the tree structure and adjust for across-taxa covariation. Phylogenetic models do exactly this but are computationally challenging because one must integrate out unobserved ancestor traits while accounting for uncertainties arising from tree estimation. The computational burden increases when taxon and trait counts grow large and becomes worse when traits include continuous and discrete quantities. Cybis et al. develop the first phylogenetic method that can assess across-trait covariation while controlling for a large, unknown evolutionary tree with hundreds of tips [1]. To jointly model mixed-type traits, this approach assumes discrete traits arise from continuously valued latent variables that follow a Brownian diffusion along the tree [2]. Assuming latent processes is a common strategy for modeling mixed-type data and it finds uses across various fields [3–7]. Subsequent work by [8] solves an essential identifiability issue in [1] by adding specific constraints on the diffusion covariance. The resulting model in particular generalizes the multivariate probit model [9]. The most important contribution of [8], however, is an efficient inference scheme that achieves order-of-magnitudes efficiency gains over [1]. In this work, we significantly advance performance compared to [8] to solve even larger problems.

Here is an intuition on why our new inference scheme, to be formally introduced in Methods section, outperforms the one by [8]. For N taxa each with P continuous or binary traits, Bayesian inference for the phylogenetic probit model involves repeatedly sampling latent variables X from their conditional posterior, an (N × P)-dimensional truncated normal distribution. The (N × P) size of the truncated normal distribution results from having one latent variable for each taxon and each trait. For this task, [8] develop a bouncy particle sampler (BPS) [10] combined with an efficient dynamic programming approach that speeds up the most expensive step in the BPS implementation. Their approach, however, fails to address another source of computational inefficiency in posterior inference under the phylogenetic probit model—a high degree of correlation between X and C. [8] use a separate Hamiltonian Monte Carlo sampler [11, HMC] to sample C and update the two sets of parameters alternately within a random-scan Gibbs scheme [12]. The phylogenetic probit model assumes X to follow a multivariate Gaussian distribution whose covariance matrix incorporates C. By the model’s very design, therefore, the values in C influence the strength and direction of the correlations between elements of X. This correlation between the two parameters slows down convergence and mixing of the Gibbs scheme as each update of X or C is strongly influenced by the current value of the other parameter. To address this issue, our present solution utilizes a state-of-the-art Markov chain Monte Carlo (MCMC) method called Zigzag-HMC [13]. Unlike BPS, this method allows a joint update of X and C through differential operator splitting [13, 14] that generalizes the previously proposed split HMC framework based on Hamiltonian splitting [11, 15]. Zigzag-HMC can further take advantage of the same $\mathcal{O}\left(N\right)$ gradient evaluation strategy developed by [8].

Our sampling scheme greatly improves the mixing of elements in C and thus provides a reliable estimate of the across-trait partial correlation matrix R, the inverse of the correlation matrix normalized to have unit diagonals. The partial correlation between two traits quantifies their conditional dependence that accounts for, and hence removes confounding by, the effects of other traits in the model. Use of partial correlations thus allow us to gain insight into potential causal pathways and help guide further research into underlying biological mechanisms.

We apply our methodology to three real-world examples. First, we re-evaluate the HIV evolution application in [8] and identify HIV-1 gag immune-escape mutations linked with virulence through strong conditional dependence relationships. Our findings closely match with the experimental literature and indicate a general pattern in the immune escape mechanism of HIV. Second, we examine the influenza H1N1 glycosylation pattern across different hosts and detect strong conditional dependencies between glycosylation sites closely related to host switching. Finally, we investigate how floral traits of Aquilegia flower attract different pollinators, for which we generalize the phylogenetic probit model to accommodate a categorical pollinator trait.

Methods

Mixed-type trait evolution

We describe biological trait evolution with the phylogenetic multivariate probit model following [8] and extend it to unordered categorical traits as in [1]. While we do not consider ordered categorical traits in this work and leave it to future work to support such traits, the mapping of latent variables in this case can also be found in [1]. We either know the phylogenetic tree $\mathcal{F}$ a priori or infer it from a molecular sequence alignment S [16]. In our two large-scale HIV and influenza applications (Results section) with available sequence data, we use a continuous-time Markov chain evolutionary model [17] to construct $p(\mathbf{\text{S}}|\mathcal{F})$ and so infer $\mathcal{F}$ simultaneously. We refer interested readers to [16] for more details on tree sampling. When investigating the efficiency gain of our method over [8], we utilize a fixed tree for a more direct comparison and also to reduce the overall run-time. For our third application on flower and pollinator co-evolution, we adopt the same fixed tree as in [18].

Consider N taxa on a tree $\mathcal{F}=(\mathbb{V},\mathbf{\text{b}})$ that is a directed, bifurcating acyclic graph. The node set $\mathbb{V}$ of size 2N − 1 contains N tip nodes, N − 2 internal nodes and one root node. The branch lengths b = (b₁, …, b_2N−2) denote the child-parent distance in real time. We observe P mixed-type traits for each taxon. The trait data Y = {y_ij} = (Y^cont, Y^disc) partition as Y^cont, an N × P_cont matrix of continuous traits and Y^disc, an N × P_disc matrix of discrete ones. We associate with each trait a latent variable $x_{ij}\in \mathbb{R}$, if the j-th trait is continuous or binary, and a (m_j − 1)-dimensional latent vector ${\mathbf{\text{X}}}_{ij}=\left\{x_{ij,k}\right\}\in {\mathbb{R}}^{m_j-1}$, if the trait is categorical, where m_j denotes the number of categorical classes. Continuous traits y_ij can be seen as as latent variables that are directly observed so x_ij = y_ij. To relate latent variables to observed discrete traits, we assume a threshold model for binary traits and a choice model for traits with more than two classes. For a binary trait y_ij,

$$\begin{array}{c}\hfill y_{ij}=g_b\left(x_{ij}\right)=\{\begin{array}{cc}0\hfill & \text{if}{x}_{ij}\le 0,\hfill \\ 1,\hfill & \text{if}{x}_{ij}>0.\hfill \end{array}\end{array}$$ (1)

For a categorical trait y_ij, the possible classes are $\{c_1,\dots ,c_{m_j}\}$ with the reference class being c₁. We have

$$\begin{array}{c}\hfill y_{ij}=g_c(x_{ij,1},\dots ,x_{ij,m_j-1})=\{\begin{array}{cc}{c}_1\hfill & \text{if}{x}_{ij,\text{max}}\le 0,\hfill \\ c_k\hfill & \text{if}{m}_j>1\text{and}{x}_{ij,\text{max}}=x_{ij,k-1}>0,\hfill \end{array}\end{array}$$ (2)

where $x_{ij,\text{max}}=\text{max}(x_{i,j},\dots ,x_{i,j+m_j-2})$. This data augmentation strategy is a common choice to model categorical data [19].

After concatenating all the latent variables, for each node i = 1, …, 2N − 1 in $\mathcal{F}$ we have P_lat-dimensional latent variable ${\mathbf{\text{X}}}_i\in {\mathbb{R}}^{P_{\text{lat}}}$ with $P_{\text{lat}}=P_{\text{cont}}+{\sum }_{j=P_{\text{cont}}+1}^{P_{\text{cont}}+P_{\text{disc}}}(m_j-1)$. As a side note, for continuous y_ij the corresponding x_ij is observed, and so X_i is actually a partially latent vector. Since in our applications only a small fraction of y_ij is continuous, we omit “partial” to ease the notation.

The latent variables follow a multivariate Brownian diffusion process along $\mathcal{F}$ such that X_i distributes as a multivariate normal (MVN)

$$\begin{array}{c}\hfill {\mathbf{\text{X}}}_i\sim \mathcal{N}({\mathbf{\text{X}}}_{\text{pa}\left(i\right)},b_i\mathbf{\Omega }),i=1,\dots ,2N-2,\end{array}$$ (3)

where X_pa(i) is the parent node value and the P_lat × P_lat covariance matrix Ω describes the across-trait association. The intuition behind b_iΩ is that the further away a child node is from its parent node (larger b_i), the bigger difference between their node values. Assuming a conjugate root prior ${\mathbf{\text{X}}}_{2N-1}\sim \mathcal{N}({\mathit{\mu }}_0,{\omega }^{-1}\mathbf{\Omega })$ with prior mean μ₀ and prior variance ω⁻¹Ω, we can analytically integrate out latent variables on all internal nodes. Marginally, then, the N × P_lat tip latent variables X have the matrix normal distribution

$$\begin{array}{c}\hfill \mathbf{\text{X}}\sim {\text{MTN}}_{N{P}_{\text{lat}}}(\mathbf{\text{M}},\mathbf{\Upsilon },\mathbf{\Omega }),\end{array}$$ (4)

where M = (μ₀, …, μ₀)^T is an N × P_lat mean matrix and the across-taxa covariance matrix Υ equals $\mathbf{\text{V}}(\mathcal{F})+{\omega }^{-1}\mathbf{\text{J}}$ [20]. The diffusion matrix $\mathbf{\text{V}}(\mathcal{F})$ is a function of branch lengths such that its diagonal elements represent the sum of branch lengths from a tip to the root, while the off-diagonal elements are the branch length from the root to the most recent common ancestor of two tips. The augmented likelihood of X and Y factorizes as

$$\begin{array}{c}\hfill p(\mathbf{\text{Y}},\mathbf{\text{X}}|\mathbf{\Upsilon },\mathbf{\Omega },{\mathit{\mu }}_0,\omega )=p\left(\mathbf{\text{Y}}\right|\mathbf{\text{X}})p\left(\mathbf{\text{X}}\right|\mathbf{\Upsilon },\mathbf{\Omega },{\mathit{\mu }}_0,\omega ),\end{array}$$ (5)

where p(Y|X) = 1 if X are consistent with Y according to Eqs (1) and (2) and 0 otherwise. Following [8], we decompose Ω = DCD where C is a P_lat × P_lat correlation matrix. The diagonal entries of C are all equal to 1, while the off-diagonal entries lie in the range of [−1, 1] and represent the correlations between pairs of latent variables and, hence, their corresponding traits. The P_lat × P_lat diagonal matrix D = {σ_ii} for i = 1, …, P_lat contains the marginal standard deviation of each latent variable. Importantly, since discrete traits only inform the sign or ordering of their underlying latent variables, certain elements of D must be set as a fixed value to ensure that the model is parameter-identifiable [8]. Without loss of generality, we fix σ_ii = 1 for σ_ii corresponding to discrete traits. For continuous traits, the square of the corresponding element (${\sigma }_{ii}^2$) multiplied by a branch length is the marginal variance for the Brownian diffusion process along that branch (Eq 3). In other words, this product reports the amount of trait variation that accumulates along a branch. [8] demonstrate the necessity of this DCD decomposition, which also allows a non-informative prior [21, LKJ] on C. For goodness-of-fit of the phylogenetic probit model we refer interested readers to [8] where the explicit tree modeling leads to a significantly better fit.

A novel inference scheme

We sample from the joint posterior to learn the across-trait correlation C

$$\begin{array}{cc}\hfill p(\mathbf{\text{C}},\mathbf{\text{D}},\mathbf{\text{X}},\mathcal{F}|\mathbf{\text{Y}},\mathbf{\text{S}})& \propto p\left(\mathbf{\text{Y}}\right|\mathbf{\text{X}})\times p(\mathbf{\text{X}}|\mathbf{\text{C}},\mathbf{\text{D}},\mathcal{F})\times \hfill \\ & p(\mathbf{\text{C}},\mathbf{\text{D}})\times p\left(\mathbf{\text{S}}\right|\mathcal{F})\times p(\mathcal{F}),\hfill \end{array}$$ (6)

where we drop the dependence on hyper-parameters (Υ, μ₀, ω) to ease notation. We fix μ₀ to be a P_lat-dimensional zero vector and ω to be 1. We then specify the priors p(C, D) and $p(\mathcal{F})$ as in [8] where $p(\mathcal{F})$ is a typical coalescent tree prior on $\mathcal{F}$ [22] and p(C, D) = p(C)p(D). We set independent log normal priors on D diagonals that correspond to continuous traits. We assume an LKJ prior on the Cholesky factor of C to ensure that C and Ω are positive definite and invertible. [8] use a random-scan Gibbs [12] scheme to alternately update X, {C, D} and $\mathcal{F}$ from their full conditionals [16]. They sample X from an NP_lat-dimensional truncated normal distribution with BPS and deploy the standard HMC based on Gaussian momentum [23] to update {C, D}. Instead, we simulate the joint Hamiltonian dynamics on {X, C, D} by combining novel Hamiltonian zigzag dynamics on X [24] and traditional Hamiltonian dynamics on {C, D}. This strategy enables an efficient joint update of the two highly-correlated sets of parameters. The improved efficiency allow us to focus on the across-trait partial correlation matrix R = {r_ij}. After collecting the MCMC samples of Ω, we obtain R by the standard transformation [25]:

$$\begin{array}{c}\hfill {\mathbf{\Omega }}^{-1}=\mathbf{\text{P}}=\left\{p_{ij}\right\},r_{ij}=-{\displaystyle \frac{p_{ij}}{\sqrt{p_{ii}{p}_{jj}}}}.\end{array}$$ (7)

Since R measures the linear relationship between pairs of variables after controlling for effects of all other variables in the model, R usually lies in a more-constrained space than C and is more difficult for the sampler to effectively explore its posterior distribution. We demonstrate the improved efficiency of our method in inferring R in Results section. In the subsequent sections, we first describe how Zigzag-HMC samples X from a truncated normal and then detail the joint update of {X, C, D}.

Zigzag-HMC for truncated multivariate normals

We outline the main ideas behind HMC [11] before describing Zigzag-HMC as a version of HMC based on Hamiltonian zigzag dynamics [13, 24]. In order to sample a d-dimensional parameter x = (x₁, …, x_d) from the target distribution π(x), HMC introduces an auxiliary momentum variable $\mathit{p}=(p_1,\dots ,p_d)\in {\mathbb{R}}^d$ and samples from the product density π(x, p) = π(x)π(p) by numerically discretizing the Hamiltonian dynamics

$$\begin{array}{c}\hfill \frac{\text{d}\mathit{x}}{\text{d}t}=\nabla K\left(\mathit{p}\right),\frac{\text{d}\mathit{p}}{\text{d}t}=-\nabla U\left(\mathit{x}\right),\end{array}$$ (8)

where U(x) = −log π(x) and K(p) = −log π(p) are the potential and kinetic energy. In each HMC iteration, we first draw p from its marginal distribution $\pi \left(\mathit{p}\right)\sim \mathcal{N}(\mathbf{0},\mathbf{\text{I}})$, a standard Gaussian and then approximate (8) from time t = 0 to t = τ by L = ⌊τ/ϵ⌋ steps of the leapfrog update with step size ϵ [26]:

$$\begin{array}{c}\hfill \mathit{p}\leftarrow \mathit{p}+\frac{ϵ}{2}{\nabla }_{\mathit{x}}\text{log}\pi \left(\mathit{x}\right),\mathit{x}\leftarrow \mathit{x}+ϵ\mathit{p},\mathit{p}\leftarrow \mathit{p}+\frac{ϵ}{2}{\nabla }_{\mathit{x}}\text{log}\pi \left(\mathit{x}\right).\end{array}$$ (9)

The end state is a valid Metropolis proposal that one accepts or rejects according to the standard acceptance probability formula [27, 28].

Zigzag-HMC differs from standard HMC insofar as it posits a Laplace momentum π(p) ∝ ∏_i exp(−|p_i|), i = 1, …, d. The Hamiltonian differential equations now become

$$\begin{array}{c}\hfill \frac{\text{d}\mathit{x}}{\text{d}t}=\text{sign}\left(\mathit{p}\right),\frac{\text{d}\mathit{p}}{\text{d}t}=-\nabla U\left(\mathit{x}\right),\end{array}$$ (10)

and the velocity v ≔ dx/dt ∈ {±1}^d depends only on the sign of p and thus remains constant until one of p_i’s undergoes a sign change (an “event”). To understand how the Hamiltonian zigzag dynamics (10) evolve over time, one must investigate when such events happen. Before moving to the truncated MVN, we first review the event time calculation for a general π(x) following [24]. Let τ^(k) be the kth event time and (x (τ⁽⁰⁾), v (τ⁽⁰⁾), p (τ⁽⁰⁾)) is the initial state at time τ⁽⁰⁾. Between τ^(k) and τ^(k+1), x follows a piecewise linear path and the dynamics evolve as

$$\begin{array}{c}\hfill \mathit{x}({\tau }^{\left(k\right)}+t)=\mathit{x}\left({\tau }^{\left(k\right)}\right)+t\mathit{v}\left({\tau }^{\left(k\right)}\right),\mathit{v}({\tau }^{\left(k\right)}+t)=\mathit{v}\left({\tau }^{\left(k\right)}\right),t\in [0,{\tau }^{(k+1)}-{\tau }^{\left(k\right)}),\end{array}$$ (11)

and

$$\begin{array}{c}\hfill p_i({\tau }^{\left(k\right)}+t)=p_i\left({\tau }^{\left(k\right)}\right)-{\int }_0^t{\partial }_{i}U\left[\mathit{x}\right({\tau }^{\left(k\right)})+s\mathit{v}({\tau }^{\left(k\right)}\left)\right]ds\text{for}i=1,\dots ,d.\end{array}$$ (12)

Therefore we can derive the (k + 1)th event time

$$\begin{array}{c}\hfill {\tau }^{(k+1)}={\tau }^{\left(k\right)}+\underset{i}{\text{min}}{t}_i,t_i=\underset{t>0}{\text{min}}\left\{p_i\right({\tau }^{\left(k\right)})={\int }_0^t{\partial }_{i}U[\mathit{x}\left({\tau }^{\left(k\right)}\right)+s\mathit{v}\left({\tau }^{\left(k\right)}\right)\left]ds\right\},\end{array}$$ (13)

and the dimension causing this event is i* = argmin_i t_i. At the moment of τ^(k+1), the i*th velocity component flips its sign

$$\begin{array}{c}\hfill v_{i^{*}}\left({\tau }^{(k+1)}\right)=-v_{i^{*}}\left({\tau }^{\left(k\right)}\right),v_j\left({\tau }^{(k+1)}\right)=v_j\left({\tau }^{\left(k\right)}\right)\text{for}j\ne i^{*}.\end{array}$$ (14)

Then the dynamics continue for the next interval [τ^(k+1), τ^(k+2)).

We now consider simulating the Hamiltonian zigzag dynamics for a truncated MVN arising from the phylogenetic probit model.

$$\begin{array}{c}\hfill \mathit{x}\sim \mathcal{N}(\mathit{\mu },\mathbf{\Sigma })\text{subject}\text{to}\mathit{x}\in \{\text{map}\left(\mathit{x}\right)=\mathbf{\text{y}}\},\end{array}$$ (15)

where μ and Σ are the mean vector and covariance matrix for the MVN and map(⋅) is the mapping from the vectorized latent variables x to y as in Eqs (1) and (2). In other words, y is the NP-dimensional vectorized discrete data such that $\mathit{x}\in {\mathbb{R}}^d$ for d = NP_lat. Since vectorizing the random variables under a matrix normal distribution (3) results in a MVN distribution, we have Σ = Ω ⊗ Υ where ⊗ denotes the Kronecker product. The mean vector μ is N copies of the pre-specified root prior mean vector μ₀ concatenated together.

In the setting of Eq (15), we have ∇U(x) = Σ⁻¹x whenever x ∈ {map(x) = y}. Importantly, this structure allows us to simulate the Hamiltonian zigzag dynamics exactly and efficiently [24]. We handle the constraint map(x) = y with a technique from [11] where the constraint boundaries embody “hard walls” that the Hamiltonian zigzag dynamics “bounce” against upon impact. To distinguish different types of events, we define gradient events arising from solutions of Eq (13), binary events arising from hitting binary data boundaries and categorical events arising from hitting categorical data boundaries.

We first consider how to find the gradient event time. Starting from a state (x, v, p), by plugging in ∇U(x) = Σ⁻¹x to Eq (13), we can calculate the gradient event time t_g by first solving d quadratic equations

$$\begin{array}{c}\hfill \mathit{p}=t{\mathbf{\Sigma }}^{-1}(\mathit{x}-\mathit{\mu })+\frac{t^2}{2}{\mathbf{\Sigma }}^{-1}\mathit{v},\end{array}$$ (16)

and then taking the minimum among all positive roots of Eq (16). When π(x) is a truncated MVN arising from the phylogenetic probit model, we exploit the efficient gradient evaluation strategy in [8] to obtain Σ⁻¹(x − μ) and Σ⁻¹v without the notorious $\mathcal{O}\left(d^3\right)$ cost to invert Σ. In our application, μ is a vector of all zeros since we set the root prior mean μ₀ to be all zero. If there is prior knowledge about μ₀, we can use another fixed value without increasing the computational cost.

Next, we focus on the binary and categorical events. We partition x into two sets: S_bin = {x_i : x_i is for binary data} and S_cat = {x_i : x_i is for categorical data}. Starting from a state (x, v, p), a binary event happens at time t_b when the trajectory first reaches a binary boundary at dimension i_b

$$\begin{array}{c}\hfill t_{\text{b}}=|x_{i_{\text{b}}}/v_{i_{\text{b}}}|,i_{\text{b}}={\text{argmin}}_{i\in I_{\text{bin}}}|x_i/v_i|\text{for}{I}_{\text{bin}}=\{i:x_i{v}_i<0\text{and}{x}_i\in S_{\text{bin}}\}.\end{array}$$ (17)

Here, we only need to check the dimensions satisfying x_iv_i < 0, i.e., those for which the trajectory is heading towards the boundary. At time t_b, the trajectory bounces against the binary boundary, and so the i_bth velocity and momentum element both undergo an instantaneous flip $v_{i_{\text{b}}}\leftarrow -v_{i_{\text{b}}}$, $p_{i_{\text{b}}}\leftarrow -p_{i_{\text{b}}}$, while other dimensions stay unchanged.

Finally, we turn to categorical events. Suppose that a categorical trait y_j = c_k belongs to one of m possible classes, and x₁, x₂, …, x_m−1 the underlying latent variables. Eq (2) specifies the boundary constraints. If k = 1, the m − 1 latent variables must be all negative, which poses the same constraint as if they were for n − 1 binary traits, therefore we can solve the event time using Eq (17). If k > 1, we must check when and which two dimensions first violate the order constraint x_k−1 = max(x₁, …, x_m−1) > 0. With the dynamics starting from (x, v, p), the categorical event time $t_{\text{c}}^j$ is given by

$$\begin{array}{c}\hfill \begin{array}{cc}\hfill t_{\text{c}}^j=|(x_{k-1}-x_{i_{\text{c}}})/(v_{k-1}-v_{i_{\text{c}}})|,& i_{\text{c}}={\text{argmin}}_{i\in I_{\text{cat}}}|(x_{k-1}-x_i)/(v_{k-1}-v_i)|,\hfill \\ & \text{for}{I}_{\text{cat}}=\{i:v_{k-1}<v_i\text{and}{x}_i\in S_{\text{cat}}\},\hfill \end{array}\end{array}$$ (18)

when $x_{i_{\text{c}}}$ reaches x_k−1 and violates the constraint. To identify i_c we only need to check dimensions with v_k−1 < v_i where the distance x_k−1 − x_i is decreasing. At $t_{\text{c}}^j$, the two dimensions involved (k − 1 and i_c) bounce against each other such that v_k−1 ← −v_k−1, $v_{i_{\text{c}}}\leftarrow -v_{i_{\text{c}}}$, p_k−1 ← −p_k−1, $p_{i_{\text{c}}}\leftarrow -p_{i_{\text{c}}}$. Note $t_{\text{c}}^j$ is for a single y_j and we need to consider all categorical data to find the actual categorical event time $t_{\text{c}}={\text{min}}_j{t}_{\text{c}}^j$.

We now present the dynamics simulation with all three event types included, starting from a state (x, v, p) with x ∈ {map(x) = y}:

1. Solve t_g, t_b, t_c using Eqs (16), (17) and (18) respectively.

2. Determine the actual (first) event time t = min{t_g, t_b, t_c} and update x and p as in Eqs (11) and (12) for a duration of t.

3. Make instantaneous velocity and momentum sign flips according to the rules of the actual event type, then go back to Step 1.

Based on the above discussion, Algorithm 1 describes one iteration of Zigzag-HMC on truncated MVNs where we simulate the Hamiltonian zigzag dynamic for a pre-specified duration t_total. For a truncated MVN arising from the phylogenetic probit model, the most computationally expensive step is the gradient evaluation in Line 3, where a matrix-vector multiplication by the precision matrix Φ = Σ⁻¹ is involved. A matrix inversion to evaluate Φ directly is expensive since Φ = Ω⁻¹ ⊗ Υ⁻¹ and computing Υ⁻¹ has a cost of $\mathcal{O}\left(N^3\right)$. We adopt the dynamic programming strategy of [8] to reduce the cost of Line 3 from either $\mathcal{O}(N^2{P}_{\text{lat}}+N{P}_{\text{lat}}^2)$ when $\mathcal{F}$ is fixed, or $\mathcal{O}(N^3+P_{\text{lat}}^3)$ when $\mathcal{F}$ is random, to $\mathcal{O}\left(N{P}_{\text{lat}}^2\right)$. We refer interested readers to [8] for details on the dynamic programming strategy. In brief, this strategy avoids explicitly inverting Υ by recursively traversing the tree [20] to obtain N conditional densities that directly translate to the desired gradient φ_x.

Algorithm 1 Zigzag-HMC for multivariate truncated normal distributions

1: function HzzTMVN(x, p, t_total)

2:  v ← sign(p)

3:  φ_x ← Φ(x − μ)

4:  t_remain ← t_total

5:  while t_remain > 0 do

 ▹ find gradient event time t_g

6:   a ← φ_v/2, b ← φ_x, c ← −p

7:   t_g ← min_i {minPositiveRoot(a_i, b_i, c_i)} ▹ “minPositiveRoot” defined below

 ▹ find binary boundary event time

8:   t_b ← min_i x_i/v_i, for i with x_iv_i < 0 and x_i ∈ S_bin

 ▹ find categorical boundary event time, n_c = number of categorical traits

9:   for j = 1, …, n_c do

10:    $t_{\text{c}}^j\leftarrow {\text{min}}_i|(x_{k-1}-x_{i_{\text{c}}})/(v_{k-1}-v_i)|\text{for}i\text{with}{v}_{k-1}<v_i$ and x_i ∈ S_cat

11:   end for

12:   $t_{\text{c}}\leftarrow {\text{min}}_j{t}_{\text{c}}^j$

 ▹ the actual event happens at time t

13:   t ← min {t_g, t_b, t_c, t_remain}

14:   x ← x + tv, p ← p − tφ_x − t²φ_v/2, φ_x ← φ_x + tφ_v

15:   if a gradient event happens at i_g then

16:    $v_{i_g}\leftarrow -v_{i_g}$

17:   else if a binary boundary event happens at i_b then

18:    $v_{i_{\text{b}}}\leftarrow -v_{i_{\text{b}}}$, $p_{i_{\text{b}}}\leftarrow -p_{i_{\text{b}}}$

19:   else if a categorical boundary event happens at i_c1, i_c2 then

20:    $v_{i_{\text{c}1}}\leftarrow -v_{i_{\text{c}1}},v_{i_{\text{c}2}}\leftarrow -v_{i_{\text{c}2}},p_{i_{\text{c}1}}\leftarrow -p_{i_{\text{c}1}},p_{i_{\text{c}2}}\leftarrow -p_{i_{\text{c}2}}$

21:   end if

22:   φ_v ← φ_v + 2v_iΦe_i

23:   t_remain ← t_remain − t

24:  end while

25: return x, p

26: end function

* minPositiveRoot(a_i, b_i, c_i) returns the minimal positive root of the equation a_ix² + b_ix + c = 0, or else returns +∞ if no positive root exists.

Jointly updating latent variables and across-trait covariance

The N × P_lat latent variables and P_lat × P_lat across-trait covariance are highly correlated with each other, so individual Gibbs updates can be inefficient. The posterior conditional of X is truncated normal and thus allows for the efficient Hamiltonian zigzag simulation. The conditional distribution for covariance components C and D has no such special structure, so we map them to an unconstrained space and deploy Hamiltonian dynamics based on Gaussian momentum. We use a standard mapping of C elements to real numbers [29] that first transforms C to canonical partial correlations (CPC) that fall in [−1, 1] and then apply the Fisher transformation to map CPC to the real line. We then construct the joint update of latent variables and covariance via differential operator splitting [13, 14] to approximate the joint dynamics of Laplace-Gauss mixed momenta.

We denote the two concatenated sets of parameters X and {C, D} as x = (x_G, x_L) with momenta p = (p_G, p_L), where indices G and L refer to Gaussian or Laplace momenta. The joint sampler updates (x_G, p_G) first, then (x_L, p_L), followed by another update of (x_G, p_G). This symmetric splitting ensures that the simulated dynamics is reversible and hence constitutes a valid Metropolis proposal mechanism [13]. The LG-STEP function in Algorithm 2 describes the process of simulating the joint dynamics for time duration 2ϵ via the analytical Hamiltonian zigzag dynamics for (x_L, p_L) and the approximate leapfrog dynamics (9) for (x_G, p_G). Because x_G and x_L can have very different scales, we incorporate a tuning parameter, the step size ratio r, to allow different step sizes for the two dynamics. To approximate a trajectory of the joint dynamics from t = 0 to t = τ, we apply the function LG-STEP m = ⌊τ/2ϵ⌋ times, and accept or reject the end point following the standard acceptance probability formula [27, 28]. We call this version of HMC based on Laplace-Gauss mixed momenta as LG-HMC and describe one iteration of LG-HMC in Algorithm 2 where the inputs include the joint potential function U(x_G, x_L). We use LG-HMC to update {X, C, D} as a Metropolis-within-Gibbs step of our random-scan Gibbs scheme. The overall sampling efficiency strongly depends on m, the step size ϵ and the step size ratio r, so it is preferable to auto-tune all of them. We provide an empirical method to automatically tune r in S1 File. We provide another option utilizing the no-U-turn algorithm to automatically decide the trajectory length m [23] and call the resulting algorithm LG No-U-Turn Sampler (LG-NUTS). We adapt the step size ϵ with primal-dual averaging to achieve an optimal acceptance rate [23].

Algorithm 2 One LG-HMC iteration

1: function LG-HMC(x_G, x_L, p_G, p_L, U, m, ϵ, r)

 ▹ Record the initial state

2:   ${\mathit{x}}_G^0\leftarrow {\mathit{x}}_G,{\mathit{x}}_L^0\leftarrow {\mathit{x}}_L,{\mathit{p}}_G^0\leftarrow {\mathit{p}}_G,{\mathit{p}}_L^0\leftarrow {\mathit{p}}_L$

3:   for i = 1, …, m do

4:    x_G, x_L, p_G, p_L ← LG-STEP(x_G, x_L, p_G, p_L, ϵ, r)

5:   end for

 ▹ Calculate the acceptance probability a, where K_G and K_L denote the kinetic energy based on Gaussian or Laplace momentum and ‖⋅‖₁, ‖⋅‖₂ are the L¹ and L² norm.

6:   $K_G^0\leftarrow ({\Vert {\mathit{p}}_G^0\Vert }_2{)}^2/2$, $K_L^0\leftarrow {\Vert {\mathit{p}}_L^0\Vert }_1$

7:   K_G ← (‖p_G‖₂)²/2, K_L ← ‖p_L‖₁

8:   $a\leftarrow \text{min}\{1,\text{exp}[U({\mathit{x}}_G^0,{\mathit{x}}_L^0)-U({\mathit{x}}_G,{\mathit{x}}_L)+K_G^0+K_L^0-K_G-K_L\left]\right\}$

 ▹ Accept or reject

9:   u ← one draw from uniform(0, 1)

10:   if u < a then

11:    return x_G, x_L, p_G, p_L

12:   else

13:    return ${\mathit{x}}_G^0,{\mathit{x}}_L^0,{\mathit{p}}_G^0,{\mathit{p}}_L^0$

14:   end if

15: end function

16: function LG-STEP(x_G, x_L, p_G, p_L, ϵ, r)

17:   x_G, p_G ← LeapFrog(x_G, p_G, ϵ)

18:   x_L, p_L ← HzzTMVN(x_G, p_G, rϵ)

19:   x_G, p_G ← LeapFrog(x_G, p_G, ϵ)

20:   return x_G, x_L, p_G, p_L

21: end function

22: function LeapFrog(x_G, p_G, ϵ)

23:   ${\mathit{p}}_G\leftarrow {\mathit{p}}_G+\frac{ϵ}{2}{\nabla }_{{\mathit{x}}_G}\text{log}p\left(\mathit{x}\right)$

24:   x_G ← x_G + ϵp_G

25:   ${\mathit{p}}_G\leftarrow {\mathit{p}}_G+\frac{ϵ}{2}{\nabla }_{{\mathit{x}}_G}\text{log}p\left(\mathit{x}\right)$

26:   return x_G, x_L

27: end function

Results

To illustrate the broad applicability of our method, we detail three real-world applications and discuss the scientific findings. We first apply our method to the HIV virulence application of [8]. The improved efficiency allows us to estimate the across-trait partial correlation with adequate effective sample size (ESS) and to reveal the conditional dependence among traits of scientific interest. We use the same HIV data set to demonstrate that LG-HMC and LG-NUTS outperform BPS (Section “Efficiency gain from the new inference scheme”), followed by two more LG-NUTS applications on influenza and Aquilegia flower evolution. We conclude this section with MCMC convergence criteria and timing results.

HIV immune escape

In the HIV evolution application of [8], a main scientific focus lies on the association between HIV-1 immune escape mutations and virulence, the pathogen’s ability to cause disease. The human leukocyte antigen (HLA) system is predictive of the disease course as it plays an important role in the immune response against HIV-1. Through its rapid evolution, HIV-1 can acquire mutations that aid in escaping HLA-mediated immune response, but the escape mutations may reduce its fitness and virulence [30, 31]. [8] identify HLA escape mutations associated with virulence while controlling for the unknown evolutionary history of the viruses. However, [8] interpret their results based on the across-trait correlation C which only informs marginal associations that can remain confounded. Now armed with a more efficient inference method, we direct our attention towards the across-trait partial correlation matrix R.

The data contain N = 535 aligned HIV-1 gag gene sequences collected from 535 patients between 2003 and 2010 in Botswana and South Africa [31]. Each sequence is associated with 3 continuous and 21 binary traits. The continuous virulence measurements are replicative capacity (RC), viral load (VL) and cluster of differentiation 4 (CD4) cell count. The binary traits include the existence of HLA-associated escape mutations at 20 different amino acid positions in the gag protein and another trait for the sampling country (Botswana or South Africa). Fig 1 depicts across-trait correlations and partial correlations with posterior medians > 0.2 (or < −0.2). Compared to correlations (Fig 1A), we observe more partial correlations with greater magnitude (Fig 1B). They indicate conditional dependencies among traits after removing effects from other variables in the model, helping to explore the causal pathway. For example, we only detect a negative conditional dependence between RC and CD4. In other words, holding one of CD4 and RC as constant, the other does not affect VL, suggesting that RC increases VL via reducing CD4. The fact that RC is not found to share a strong conditional dependence with VL may be explained by the strong modulatory role of immune system on VL. Only when viruses with higher RC also lead to more immune damage, as reflected in the CD4 count, higher VL may be observed as a consequence of less suppression of viral replication. As such, our findings are in line with the demonstration that viral RC impacts HIV-1 immunopathogenesis independent of VL [32].

Fig 1.

(A) Across-trait correlation and (B) partial correlation with a posterior median > 0.2 or < −0.2 (in color). HIV gag mutation names start with the wild type amino acid state, followed by the amino acid site number according to the HXB2 reference genome and end with the amino acid as a result of the mutation (‘X’ means a deletion). Country = sample region: 1 = South Africa, -1 = Botswana; RC = replicative capacity; VL = viral load; CD4 = CD4 cell count. (C) Conditional dependencies between HIV-1 immune escape mutations that affect RC or VL. Node and edge color indicates whether the dependence is positive (orange) or negative (blue).

The partial correlation also helps to decipher epistatic interactions and how the escape mutations and potential compensatory mutations affect HIV-1 virulence. For example, we find a strong positive partial correlation between T186X and T190X. Studies have shown that T186X is highly associated with reduced VL [33, 34] and it requires T190I to partly compensate for this impaired fitness so the virus stays replication competent [35]. The negative conditional dependence between T186X and RC and the positive conditional dependence between T190I and RC are consistent with this experimental observation. In contrast, with the strong positive association between T186X and T190, the marginal association fails to identify their opposite effects on RC. Another pair of mutations that potentially shows a similar interaction is H28X and M30X, which have a positive and negative partial correlation with VL, respectively. These mutations have indeed been observed to co-occur in gag epitopes from longitudinally followed-up patients [36]. Fig 1B keeps all the other compensatory mutation pairs in Fig 1A such as A146X-I147X and A163X-S165X that find confirmation in experimental studies [37, 38].

More generally, when considering the viral trait RC and the infection trait VL, for which their variation are to a considerable extent attributable to viral genetic variation [39], we reveal an intriguing pattern. As in Fig 1C, when two escape mutations impair virulence, and there is a conditional dependence between them, it is always negative. When two mutations have opposing effects on these virulence traits, the conditional dependence between them (if present) is almost always positive, with one exception of the negative effect between V168I and S357X. For example, T186X and I61X both have a negative impact on RC and the negative effect between them suggests that their additive, or even potentially synergistic, impact on RC is inhibited. Moreover, they appear to benefit from a compensatory mutation, T190X, which has been corroborated for the T186X-T190X pair at least as reported above. Also for VL, the conditional dependence between mutations that both have a negative impact on this virulence trait is consistently negative. Several of these individual mutations may benefit from H28X as a compensatory mutation, as indicated by the positive effect between pairs that include this mutation, and as suggested above for H28X—M30X. This illustrates the extent to which escape mutations may have a negative impact on virulence and the need to evolve compensatory mutations to restore it. We note that our analysis is not designed to recover compensatory mutations at great length as we restrict it to a limited set of known escape mutations, while mutations on many other sites may be compensatory. In fact, our analysis suggests that some of the considered mutations may be implicated in immune escape due to their compensatory effect rather than a direct escape benefit.

Efficiency gain from the new inference scheme

We demonstrate that the joint update of latent variables X and the covariance matrix Ω significantly improve inference efficiency. For this purpose we use the large HIV dataset from Section “HIV immune escape” with N = 535, P_disc = 21, P_cont = 3, where the efficiency gain becomes significant. Our implementations of the algorithms have been validated on smaller truncated MVNs, on which simple rejection sampling can provide the ground truth up to quantifiable Monte Carlo errors. The Zigzag-HMC implementation has also been validated through the standalone implementation in an R package “htdg” [40].

We consider 4 sampling schemes BPS, Zigzag-HMC, LG-HMC, and LG-NUTS. To enable a more direct comparison while saving computational time, we separate tree inference from the inference for Ω and X and fix $\mathcal{F}$ as the maximum clade credibility tree from the HIV immune escape application. BPS and Zigzag-HMC only update X and we use the standard NUTS transition kernel (i.e. standard HMC combined with no-U-turn algorithm) for the Ω elements. LG-HMC employs the joint update of X and Ω described in Section “Jointly updating latent variables and across-trait covariance”. LG-NUTS additionally employs the No-U-Turn algorithm to decide the number of steps and a primal-dual averaging algorithm to calibrate the step size. We set the same t_total for BPS and Zigzag-HMC for a fair comparison. To tune LG-HMC, we first supply it with an optimal step size ϵ learned by LG-NUTS, then decide the number of steps m = 100 as it gives the best performance among the choices (10, 100, 1000). We conduct 3 independent simulations for each sampling scheme and report the per-run-time ESS for 5 parameters—the across-trait correlation C, partial correlation R, latent variable X, log joint density log p(X, Ω) and log likelihood l(X, Ω). C and R are of primary scientific interest as they provide insights into correlation structure among the traits. Examining ESS of the highest dimensional parameter X is also important for diagnostic purposes. ESS’s of log p(X, Ω) and l(X, Ω) help us additionally evaluate how well the samplers explore the target distribution overall. As reported in Table 1, BPS is outperformed by the three other samplers in terms of efficiency for all five parameters. While a formal theoretical analysis is beyond the scope of this work, we provide an empirical explanation for the different performances of BPS and Zigzag-HMC in S2 File. LG-HMC achieves the highest per run-time ESS for R, resulting in a 5× speed-up compared to BPS. The result also highlights that inferring R is more challenging than inferring C, with the elements of R generally having lower ESS, but the difficulty can be largely eliminated by jointly updating X and Ω through LG-HMC and LG-NUTS. Although Zigzag-HMC achieves much higher ESS for X than LG-HMC, the latter performs best in the most difficult and critical task of updating R. Compared to LG-HMC, LG-NUTS exhibits lower efficiency and higher variance across the 3 runs, likely due to the No-U-Turn algorithm’s tendency to require some extraneous leapfrog steps [41, 42]. We also provide the histograms for the per run-time ESS of R elements in S1 Fig. Based on our findings, we recommend using LG-HMC with multiple choices of hyper parameters (m, ϵ), with a good starting point being (100, 0.01), or the auto-tuned LG-NUTS.

Table 1. Efficiency comparison among different sampling schemes (BPS, Zigzag-HMC, LG-HMC, LG-NUTS).

We calculate effective sample size (ESS) per hour run-time for the elements of C, R, X, log joint density log p(X, Ω), and likelihood l(X, Ω). For the three multivariate parameters (C, R, X) with dimensions 276, 276, and 11,235, respectively, we report the minimal ESS across all dimensions. We conduct three independent simulations for each method and report the ESS values in the first three rows. We include the mean and standard deviation in the last row for each method to provide a summary of its overall performance. The bold number indicates the highest value in each of the five columns. For BPS, given the larger number of iterations required to achieve convergence, we record one sample of X every 1,000 iterations to comply with storage limitations, and report upper bounds of the actual ESS by multiplying the ESS from thinned samples by 1,000.

ESS/hour	C(276d)	R(276d)	X (11,235d)	log p(X, Ω)	l(X, Ω)
BPS	6.05	1.46	< 760*	0.56	0.56
	5.86	2.41	< 670	0.52	0.52
	0.55	0.49	< 100	0.42	0.43
	4.16(3.13)	1.45(0.96)	-	0.5(0.07)	0.5(0.07)
Zigzag-HMC	13.75	2.23	1480	4.42	4.44
	7.79	2.36	1057	5.38	5.38
	14.9	2.53	927	5.16	5.2
	12.15(3.82)	2.37(0.15)	1155(289)	4.99(0.5)	5.01(0.5)
LG-HMC	8.26	7.33	4.92	4.79	4.81
	7.11	8.59	7.76	5.09	5.1
	7.44	6.49	4.46	5.33	5.34
	7.6(0.59)	7.47(1.06)	5.71(1.79)	5.07(0.27)	5.08(0.26)
LG-NUTS	1.31	1.29	1.69	0.7	0.7
	11.93	7.52	6.37	1	1.06
	7.77	6.09	2.64	2.71	2.72
	7.01(5.35)	4.97(3.26)	3.57(2.47)	1.47(1.09)	1.49(1.07)

* The ESS estimates after 1/1000 thinning are 0.76, 0.67, 0.10

Glycosylation of Influenza A virus H1N1

Influenza A viruses of the H1N1 subtype currently circulate in birds, humans, and swine [43–45], where they are responsible for substantial morbidity and mortality [46, 47]. The two surface glycoproteins hemagglutinin (HA) and neuraminidase (NA) interact with a cell surface receptor and so their characteristics largely affect virus fitness and transmissibility. Mutations in the HA and NA, particularly in their immunodominant head domain, sometimes produce glycosylations that shield the antigenic sites against detection by host antibodies and so help the virus evade antibody detection [48–51]. On the other hand, glycosylation may interfere with the receptor binding and also be targeted by the innate host immunity to neutralize viruses. Therefore there must be an equilibrium between competing pressures to evade immune detection and maintain virus fitness [52, 53]. The number of glycosylations that leads to this balance is expected to vary in host species experiencing different strengths of immune selection. Despite decades of tracking IAVs evolution in humans for vaccine strain selection and recent expansions of zoonotic surveillance, the evolvability and selective pressures on the HA and NA have not been rigorously compared across multiple host species. Here, we examine the conditional dependence between host type and multiple glycosylation sites by estimating the posterior distribution of across-trait partial correlation while jointly inferring the IAVs evolutionary history.

We use hemagglutinin (H1) and neuraminidase (N1) sequence data sets for influenza A H1N1 produced by Trovão et al. as described in [54]. We scan all H1 and N1 sequences to identify potential N-linked glycosylation sites, based on the motif Asn-X-Ser/Thr-X, where X is any amino acid other than proline (Pro) [55]. We then set a binary trait for each sequence encoding for the presence or absence of glycosylations at a particular amino acid site. We keep sites with a glycosylation frequency between 20% and 80% for our analysis. This gives six sites in H1 and four sites in N1. We include another binary trait for the host type being mammalian (human or swine) or avian, so the sample sizes are N = 964, P = 7 (H1) and N = 896, P = 5 (N1).

The six H1 glycosylation sites consist of three pairs that are physically close (63/94, 129/163, and 278/289, see Fig 2). Sites 63 and 94 are particularly close to each other, though distances will vary slightly with sequence. A negative conditional dependence suggests glycosylation at two close sites may be harmful for the virus (63/94 and 278/289) while a positive effect between two sites suggests a potential benefit (63/129 and 94/278). We detect a negative conditional dependence between mammalian host and glycosylation site 94 and 289. Avian viruses have a stronger tendency to have site 289 glycosylated (Fig 2).

Fig 2.

(A) Across-trait partial correlation among H1 glycosylation sites and host type with a posterior median > 0.2 or < −0.2 (in color and number). (B) HA structure of a 2009 H1N1 influenza virus (PDB entry 3LZG) with six glycosylation sites highlighted. Site 278 and 289 are in the stalk domain and all others are in the head domain. (C) The maximum clade credibility (MCC) tree with branches colored by the posterior median of the latent variable underlying H1 glycosylation site 289. The heatmap on the right indicates the host type of each taxon.

In N1, glycosylations are more strongly correlated than H1 (Fig 3). Two pairs of glycosylation sites have a positive conditional dependency in between (50/68 and 50/389) and two pairs (44/68 and 68/389) have a negative one. We omit a structural interpretation since all sites but 389 are located in the NA stalk, for which no protein structure is available. There is a positive conditional dependence between mammalian host and glycosylations at sites 44 and 68. None of the avian lineages has glycosylation site 44 while most swine and some human lineages have it. Similarly, glycosylation at site 68 is present in most swine and human lineages but only in avian lineages circulating in wild birds, not those in poultry.

Fig 3.

(A) Across-trait partial correlation among N1 glycosylation sites and host type with a posterior median > 0.2 or < −0.2 (in color and number). (B)(C) The maximum clade credibility (MCC) tree with branches colored by the posterior median of the latent variable underlying N1 glycosylation site 44 and 68.

Aquilegia flower and pollinator co-evolution

Reproductive isolation allows two groups of organisms to evolve separately, eventually forming new species. For plants, pollinators play an important role in reproductive isolation [56]. We examine the relationship between floral phenotypes and the three main pollinators for the columbine genus Aquilegia: bumblebees, hummingbirds, and hawk moths [18]. Here, the pollinator species represents a categorical trait with three classes and we choose bumblebee with the shortest tongue as the reference class. Fig 4 provides the across-trait correlation and partial correlation. Compared to a similar analysis on the same data set that only looks at correlation or marginal association [1], partial correlation controls confounding and indicates the conditional dependencies between pollinators and floral phenotypes that can bring new insights.

Fig 4. Across-trait correlations (A) and partial correlations (B) with posterior medians > 0.2 or < −0.2 (in color).

BB = bumblebee.

For example, we observe a positive marginal association between hawk moth pollinator and spur length but no conditional dependence between them. The marginal association matches with the observation that flowers with long spur length have pollinators with long tongues [18, 57]. The absence of a conditional dependence makes intuitive sense because hawk moth’s long tongue is not likely to stop them from visiting a flower with short spurs when the other floral traits are held constant. In fact, researchers observe that shortening the nectar spurs does not affect hawk moth visitation [58]. Similarly, the positive partial correlation between orientation and hawk moth also finds experimental support. The orientation trait is the angle of flower axis relative to gravity, in the range of (0, 180). A small orientation value implies a pendent flower whereas a large value represents a more upright flower [59]. Due to their different morphologies, hawk moths prefer upright flowers while hummingbirds tend to visit pendent ones. Making the naturally pendent Aquilegia formosa flowers upright increases hawk moth visitation [59]. These results suggest that partial correlation may have predictive power for results from carefully designed experiments with controlled variables.

MCMC setup and convergence assessment

We run all simulations on a node equipped with AMD EPYC 7642 server processors which possess 48 cores and 96 threads, with a base clock speed of 2.3 GHz. For every MCMC run, the minimal effective sample size (ESS) across all dimensions of X and R after burn-in is above 100. As another diagnostic, for our two large-scale applications on HIV-1 and H1N1 influenza, we run three independent chains and confirm the potential scale reduction statistic $\widehat{R}$ for all partial correlation elements falls between [1, 1.03], below the common criterion of 1.1 [60]. To reach a minimal ESS = 100 across all R elements, the post burn-in run-time and number of MCMC transition kernels applied for the joint inference are 21 hours and 1.3 × 10⁶ (HIV-1), 113 hours and 7.9 × 10⁷ (H1), 76 hours and 1.4 × 10⁸ (N1). These run-times suggest the difficulty of our large-scale inference tasks where besides the main challenge of sampling {X, C, D}, updating the many tree parameters with Metropolis-Hastings transition kernels also takes a large number of iterations. To reduce the computational burden associated with tree inference, one practical approach is to utilize a set of pre-computed trees and incorporate tree swaps within the MCMC transition kernel.

Discussion

Learning how different biological traits interact with each other from many evolutionarily related taxa is a long-standing problem of scientific interest that sheds light on various aspects of evolution. Towards this goal, we develop a scalable solution that significantly improves inferential efficiency compared to established state-of-the-art approaches [1, 8]. Our novel strategy enables learning across-trait conditional dependencies that are more informative than the previous marginal association based analyses. This approach provides reliable estimates of across-trait partial correlations for large problems, on which the established BPS-based method struggles. In two large-scale analyses featuring HIV-1 and H1N1 influenza, the improved efficiency allows us to infer conditional dependencies among traits of scientific interest and therefore investigate some of the most important molecular mechanisms underlying the disease. In addition, our approach incorporates automatic tuning, so that the most influential tuning parameters automatically adapt to the specific challenge the target distribution presents. Finally, we extend the phylogenetic probit model to include categorical traits and illustrate its use in examining the co-evolution of Aquilegia flower and pollinators.

We leverage the cutting-edge Zigzag-HMC [13] to tackle the exceedingly difficult computational task of sampling from a high-dimensional truncated normal distribution in the context of the phylogenetic probit model. Zigzag-HMC proves to be more efficient than the previously optimal approach that uses the BPS, especially when combined with differential operator splitting to jointly update two sets of parameters X and Ω that are highly correlated. The improved efficiency allows us to obtain reliable estimates of the conditional dependencies among traits. In our applications, we find that these conditional dependencies better describe trait interactions than do the marginal associations. It is worth mentioning that another closely related sampler, the Markovian zigzag sampler [61], or MZZ, may also be appropriate for this task but provides lower efficiency than Zigzag-HMC [24]. While Zigzag-HMC is a recent and less explored version of HMC, BPS and MZZ are two central methods within the piecewise deterministic Markov process literature that have attracted growing interest in recent years [62, 63]. Intriguingly, the most expensive step of all three samplers is to obtain the log-density gradient, and the same linear-order gradient evaluation method [8] largely speeds it up.

We now consider limitations of this work and the future directions to which they point. First, the phylogenetic probit model does not currently accommodate a directional effect among traits since it only describes pairwise and symmetric correlations. However, the real biological processes are often not symmetric but directional, where it is common that one reaction may trigger another but not the opposite way. A model allowing directed paths is preferable since it better describes the complicated causal network among multiple traits. Graphical models with directed edges [64] are commonly used to learn molecular pathways [65, 66], but challenges remain to integrate these methods with a large and randomly distributed phylogenetic tree. Toward this goal, one may construct a continuous-time Markov chain to describe how discrete traits evolve [67, 68], but with P binary traits the transition rate matrix grows to the astronomical size 2^P. Second, though our method achieves the current best inference efficiency under the phylogenetic probit model, there is still room for improvement. In the influenza glycosylation example, we use a binary trait indicating the host being either avian or mammal (human or swine), instead of setting a categorical trait for host type. In fact, we choose not to use a three-class host type trait because it causes poor mixing for the partial correlation elements. We suspect two potential reasons for this. First, according to our model assumptions for categorical traits (Eq 2), the latent variables underneath the same trait are very negatively correlated, leading to a more correlated and challenging posterior. Second, in our specific data sets, the glycosylation sites tend to be similar in human and swine viruses, further increasing the correlation among posterior dimensions. One potential solution is to de-correlate some latent variables by grouping them into independent factors using phylogenetic factor analysis [69, 70]. Finally, one may consider a logistic or softmax function to map latent variables to the probability of a discrete trait. This avoids the hard truncations in the probit model but also adds another layer of noise. It requires substantial effort to develop an approach that overcomes the above limitations while supporting efficient inference at the scale of applications in this work.

Supporting information

S1 File. Auto-tuning of r.

A heuristic to auto-tune the step size ratio r.

(PDF)

S2 File. Zigzag-HMC explores the energy space more efficiently than BPS.

An intuition for BPS’s slow movement in energy space.

(PDF)

S1 Fig. Histograms of per run-time ESS for r_ij.

(PDF)

Data Availability

We implement our algorithms within BEAST software and provide the data sets and instructions at https://github.com/suchard-group/hzz_data_supplementary.

Funding Statement

ZZ, PL and MAS are partially supported by National Institutes of Health grant R01 AI153044. MAS and PL acknowledge support from the European Research Council under the European Union’s Horizon 2020 research and innovation programme (grant agreement no. 725422 - ReservoirDOCS) and from the Wellcome Trust through project 206298/Z/17/Z (The Artic Network). JLC is supported by the intramural research program of the National Library of Medicine, National Institutes of Health. AH is supported by NIH grant K25 AI153816. The funders had no role in study design, data collection and analysis, decision to publish, or preparation of the manuscript.