Development as drift-diffusion with birth-death.

We model cells as points in a space , which we take to be a representation of the space of possible cellular molecular states (for instance, in the case of scRNA-seq data, represents the space of gene expression profiles). Typically, we will take to be the ambient state space. We regard cells as evolving following a drift-diffusion process [18, 23, 24] described by the stochastic differential equation (SDE) (1) where is the state of a cell at time t, v is a vector field, the diffusivity σ² captures the noise level and dB_t denotes the increments of a Brownian motion. Over time a cell traces out a path, or trajectory, X_t in the state space. We illustrate these concepts in Fig 1A. In addition, cells are subject to division and death events at exponential rates β(x) and δ(x) respectively, which may vary in spatial location in . That is, in an infinitesimal time interval dt, a cell X_t may divide with probability β(X_t) dt or die with probability δ(X_t) dt. The underlying assumption in our framework is that the evolution of cell states can be well-described by a Markov process. While in reality this property may not be truly satisfied, we note that many other methods [18, 20–22] also make this assumption for the sake of analytical tractability.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 1. Conceptual illustration of inference problem.

(a) Vector field (grey), sample trajectory (green) and observed cell state (red) drawn from ground truth process. (b) Sampled snapshot with labelled source (red) and sink (blue) states. (c) Inferred state transition probabilities. (d) Fate probabilities calculated from Markov chain.

https://doi.org/10.1371/journal.pcbi.1009466.g001

Population-level model.

At the population level, the drift-diffusion process with birth and death can be described by a population balance partial differential equation (PDE) [18, 25] (2) where ρ(x, t) is a continuous population density, and R is a spatially varying flux rate defined as R(x) = β(x) − δ(x) that captures creation and destruction of cells due to birth and death, as well as entry and exit from the system. For R(x) > 0 we refer to x as a source state, and similarly for R(x) < 0 we refer to x as a sink state.

Observation model.

As we discussed earlier in this introduction, many biological processes exist approximately in an equilibrium or steady state. In this setting, a snapshot at a single instant in time will capture all stages of cellular development in the system [2, 18], and relative proportions of various cell types remain unchanged over time. Mathematically at the population level, this assumption amounts to demanding that ∂_tρ(x, t) = 0 in Eq (2), that is, the population level cell density does not change. We will write ρ_eq(x) for this steady-state solution. Experimental observation of such a system is therefore equivalent to sampling a collection of N cellular states from the population, i.e. We may describe this finite sample as an empirical distribution supported on the discrete space , (3) Fig 1B shows an example of such a sample dataset drawn from the equilibrium distribution, where we have identified also source and sink states.

Laws on paths.

The process described by (1) and (2) is a superposition of a birth-death process and a drift-diffusion process. The drift-diffusion component of this process, described by Eq (1), governs the evolution of individual cell states. Therefore, we seek to learn something about the drift-diffusion dynamics from observed snapshot data. In the framework of SDEs, Eq (1) (equipped also with an initial condition) induces a probability distribution over the space of possible cell trajectories. As we also argue in [23], this law on paths is the natural object we seek to estimate since it directly encodes the trajectories that cells may follow. In practice, since the process is time-homogeneous, it is characterised by its time-Δt transition densities P[X_Δt ∈ ⋅|X₀ = x]. We illustrate in Fig 1C the concept of a transition density in the discrete setting of Eq (3). As we will find, an approximation of this transition density for small enough Δt can be obtained as the solution to a strictly convex minimisation problem. Conveniently, we do not have issues of multiple local minima which may be the case if we attempt to recover the drift field v or potential landscape Ψ directly, such as in [26, 27].

Identifiability.

For the sake of making inferences about the law on paths induced by Eq (1), we must necessarily have estimates of the flux rate R(x) and the noise level σ². As discussed at length by Weinreb et al. [18], when only a single snapshot (i.e. ) is available, in general more than one drift field v can give rise to the same steady-state density profile ρ_eq. To ensure uniqueness of the solution, we must restrict to the case where the drift is given by the gradient of a scalar potential [18, 23], i.e. v = −∇Ψ. We note that Ψ can be thought of a kind of Waddington’s epigenetic landscape [1]. For completeness, we provide an example illustrating the issue of identifiability in S1 Appendix.

Formulation of the discrete problem.

In practice, we have access to an empirical distribution (see Eq (3)) supported on the discrete set that can be thought of as approximating the true continuous density ρ_eq discussed previously. We also assume for each observed cell x_i that we have an estimate of the corresponding flux rate . In a practical biological setting, cell states which are expected to divide or die should therefore have or respectively, and those states which do neither should have . In addition to division and death, terminally differentiated cells expected to shortly exit the system may be regarded as representing sinks, and therefore assigned . The numerical values for flux rates may be estimated from cell-cycle signatures [3] or prior biological knowledge [18].

In this discrete setting, the growth step Eq (4) is local in space and thus its analogue can be directly written for a chosen small value of Δt to obtain μ₀: (10) Next, the effect of the transport step Eq (5) is to rearrange mass via diffusion and drift so that we return to the steady state distribution . We cannot take μ₁ to be exactly, since a single step of the splitting scheme introduced in Eqs (4) and (5) is only accurate up to . Therefore, a straightforward application of the growth step Eq (10) will result in a slight change in the total mass of the system. Additionally, in practice we have only estimates of the true flux rates, further contributing to this effect. Consequently, we must instead re-normalise μ₁ so that it has the same mass as μ₀: (11) With μ₀ and μ₁ constructed in this way, we may compute the solution to the discrete Schrödinger problem Eq (8).

Choice of ε and Δt.

The key parameters for the scheme we describe are Δt, the time step introduced in the growth splitting, and the regularisation parameter ε for entropy-regularised optimal transport. In the theoretical framework of the Schrödinger problem, these parameters have a proportional relationship ε = σ²Δt. The accuracy of the scheme should improve in the limit as ε → 0, Δt → 0 and ε = σ²Δt, since the splitting approximation becomes exact. However, in practice where we have discrete samples we find that allowing ε and Δt to deviate from this relationship often leads to better results.

In the discrete setting, a key limitation is the value of ε, which controls the level of diffusion in the reference process in the Schrödinger problem, and consequently influences the level of diffusion in the inferred process. In practice when we are dealing with a limited number of samples in a potentially high-dimensional space, taking ε too small may lead to an ill-conditioned problem. The reason for this is that the distance between points in the set of samples may be quite large compared to ε = σ²Δt. That is, may be exceedingly small, resulting in a reference process that mixes extremely slowly. On the other hand, if we pick a reasonably sized ε, strictly adhering to the proportionality relationship may mean that the corresponding Δt is too large for the splitting approximation to be a good one. In practice, we have often found that it is helpful to take Δt to be slightly smaller (and consequently ε slightly larger) than what is expected in theory.

Simulation setup and parameters.

We first consider a tri-stable system (1) in , with drift term v taken to be the negative gradient of the potential (14) with wells {z₀, z₁, z₂} located at We illustrate this potential landscape in Fig 3A in the first two dimensions of . Simulated particles are initially isotropically distributed around the origin following the law (15) at t = 0, where denotes the standard normal distribution in with covariance I. Particles then evolve following drift-diffusion dynamics with σ² = 0.5. Whenever a particle falls in the vicinity of any of the potential wells {z₀, z₁, z₂}, it is removed with exponential rate 5. That is, in each time step dt, a particle located in a sink region is removed with probability 5 dt. We defined the sink region for each potential well z_i to be a ball of radius r = 0.25 centred at z_i.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 3. Potential-driven simulation.

(a) Illustration of the potential Ψ in the first two dimensions of the space . (b) Examples of simulated particle trajectories following the drift-diffusion process. (c) Snapshot particles shown in the first two dimensions of , with the value of R indicated. Source and sink regions correspond to R > 0 and R < 0 respectively. (d) Evolution of the dynamics recovered by StationaryOT.

https://doi.org/10.1371/journal.pcbi.1009466.g003

Exact sampling of snapshots from the steady state distribution ρ_eq of Eq (2) would require the solution of a high-dimensional PDE and is therefore computationally difficult. Instead, we obtain an approximate snapshot of the system at its steady state by simulating N = 250 trajectories from start to finish using the Euler-Maruyama method (16) For our simulations we employed a time step τ = 1 × 10⁻³, and from each trajectory we sampled a single particle state chosen at a random time chosen uniformly on to form the snapshot data . This scheme was also the one used for obtaining snapshots in [18].

Particles x_i located in the sink regions were labelled as ‘sink’ sites and assigned flux rates R_i so that the average sink flux rate was −5 and total flux rate for each well {z₀, z₁, z₂} matched the ground truth in proportion. Particles located in a ball of radius 0.25 of the origin were labelled as ‘source’ sites, corresponding to locations x_i with R_i > 0. Since we deal with finite samples, we assigned R_i uniformly on source sites such that the equilibrium condition ∑_i R_i = 0 was satisfied.

We display some example trajectories in Fig 3B, and illustrate the snapshot data in Fig 3C, where the values of R_i at source and sink sites are shown by colour.

Inferring dynamics using StationaryOT.

To apply StationaryOT, we chose a time step Δt = 25τ = 2.5 × 10⁻², noting that this is small compared to the average particle lifespan of 0.934 in this simulation. We solved the StationaryOT problem using entropy-regularised optimal transport using a range of regularisation parameter values ε in 10^−2.5 − 10¹. As we discuss in more detail later, we found that ε = 0.026 best matched the ground truth in terms of average fate probability correlation across the three lineages. For this choice of ε we computed a forward transition matrix P from the optimal transport coupling γ_Δt by row-normalising: The matrix P therefore describes a time-Δt evolution of probability densities on the discrete set . For an initial distribution π₀ supported on , we can compute the evolution {π₀P^k, k = 0, 1, 2, …} over steps of length Δt, which we take to be an estimate of the dynamics of the underlying drift-diffusion process. In Fig 3D we show the inferred process for k = 1, 5, 10, 20 where we have taken π₀ to be uniform on the source sites.

From the transition probabilities P_ij we may compute fate probabilities for each of the three lineages defined by the potential wells {z₀, z₁, z₂}. (These are absorption probabilities of the Markov chain P—see S1 Appendix for further details). We summarise these fate probabilities in Fig 4A–4D, and find that the correspondence between inferred and ground truth fate probabilities measured in terms of the Pearson correlation is high (r ≈ 0.99). As another measure of the accuracy of the estimated dynamics, we compute the mean-first passage time (MFPT) of each sampled point x_j. This is the expected time at which a Markov chain initialised at a randomly chosen source location x_i reaches x_j: where MFPT(x_j|x_i) denotes the conditional MFPT for a particle starting at x_i to hit state x_j. Comparing the MFPT estimates to the ground truth MFPT in Fig 4E, we find that the correspondence is high (r > 0.9).

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 4. Accuracy of inferred dynamics for potential-driven system.

(a) Colours representing estimated fate probabilities towards each of the wells {z₀, z₁, z₂} are displayed on the snapshot coordinates. (b-d) Correlation with ground truth fate probabilities. (e) Comparison of estimated MFPT (in terms of Markov chain steps) to ground-truth MFPT (in continuous time units). (f) Comparison of recovered velocities to ground truth velocity.

https://doi.org/10.1371/journal.pcbi.1009466.g004

Reconstructing the drift field v.

Since the transition probabilities encode the displacement law of the underlying process over a time interval Δt, we can also recover an estimate of the velocity field v by computing the expected time-Δt displacement of each cell: In Fig 4F we show the estimated velocity field alongside the ground truth v, and we measure the error by computing the mean cosine error between vector fields: We observe that the estimated field resembles the ground truth quite well near the potential wells where particles are subjected to a relatively strong drift, but struggles near the origin where the true velocity field has a small magnitude. Overall however, the cosine error is close to zero, indicating that our recovered velocity field matches the ground truth field well.

Effect of the choice of regularisation parameter ε and flux rates R.

We next turn to investigating the effects of the choice of the regularisation parameter ε on the quality of the recovered dynamics. To quantitatively measure this, we choose to compute the average correlation r between estimated and ground truth fate probabilities across the three lineages. We applied StationaryOT using both entropy and quadratic regularisations, and let ε vary on a log-scale from 10^−2.5 − 10¹ and 10⁻¹ − 10² respectively.

As shown in Fig 5A, in the case of entropy-regularised optimal transport we observe from the fate probability estimates that there is clearly a single optimal value of this parameter at ε = 0.026. This is larger than the theoretically optimal value of σ²Δt = 0.0125, in keeping with our observations discussed earlier (see Choice of ε and Δt.). However, StationaryOT with the theoretically optimal value fares only slightly worse and is located close to the maximum. When ε is chosen too small or too large, performance degrades. On the other hand, we find that performance when using a quadratic regularisation is much less sensitive to the choice of ε, with the correlation over ε showing a much flatter profile. We emphasise that ε is shown on a logarithmic axis in order to remove differences in the scales of ε for different regularisations.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 5. Effect of parameter choices on inference for potential-driven system.

(a) Correlation for varying regularisation parameter ε for entropic and quadratic regularisations. For entropic regularisation, the theoretically optimal value of ε is indicated in red. (b) Summarised correlations for systematic perturbation of flux rates towards each of the wells {z₀, z₁, z₂}. Note that the simplex represents the perturbation applied to the true flux rates rather than the flux itself, so the centre (1/3, 1/3, 1/3) of each simplex corresponds to using the true flux rates.

https://doi.org/10.1371/journal.pcbi.1009466.g005

Since flux rates R are also parameters that need to be specified, we examine the sensitivity to varying the flux rate in Fig 5B. We systematically perturb the proportion of particles that exit at each of the wells {z₀, z₁, z₂} by scaling the ground truth flux rates by values in the simplex. We observe that performance is optimal near (1/3, 1/3, 1/3) corresponding to no perturbation to the ground truth flux rates, and degrades moderately as bias is introduced to each well. We show results for entropic and quadratic regularisation where ε is chosen to be the optimal values of 0.026 and 0.43 respectively, and note that both choices of regularisation behave very similarly. While it is certain that perfect knowledge of the flux rates yields optimal performance, examining the level sets at r > 0.95, 0.975 leads us to conclude that StationaryOT still provides informative fate probabilities for a wide range of perturbations. Interestingly, the correlation remains reasonable even at the corners of the simplex, when all the flux is localised at a single well. This is due to particles in the vicinity of the other wells that diffuse into those wells randomly (as opposed to due to an inferred drift).

Laws on paths.

The SDE described by Eq (1) naturally induces a probability measure on the space of continuous functions valued in , from which one can sample cell trajectories. We discuss this point of view at length in related work on the non-stationary case [23]. From this perspective, we may treat the recovered process as inducing a law on discrete-time paths valued in , and we expect that a good estimate of the dynamics should correspond to a law on paths that is closer to the ground truth law. To illustrate this, in Fig 6A we display 100 sample paths over T = 25 timesteps, i.e. t ∈ {0, Δt, 2Δt, …, (T − 1)Δt}. The ground truth paths are obtained by sampling solutions to the Eq (1) using the Euler-Maruyama method with the initial condition Eq (15). To sample paths from the output of StationaryOT, we sampled first X₀ to be a random source cell and then let X_kΔt, 1 ≤ k ≤ 24 evolve following the Markov chain defined by the transition matrix P output by StationaryOT.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 6. Inferred dynamics in the space of paths for potential-driven system.

(a) Collections of 100 sample paths from the ground truth process Eq (1) as well as StationaryOT outputs for both entropic and quadratic OT with optimal and sub-optimal ε. The vertical axis corresponds to a projection 〈x, u〉 of the 10-dimensional state space onto a convenient 1-dimensional subspace defined by u = (cos(π/12), sin(π/12), 0, …, 0). (b) W₂ error on paths for StationaryOT reconstructions, shown for 5 repeated samplings of 250 paths.

https://doi.org/10.1371/journal.pcbi.1009466.g006

We compare the ground truth to the StationaryOT output for both entropic and quadratic OT for optimal and sub-optimal (taken as 10× the optimal value) choices of the regularisation parameter ε. Visually, it is clear that StationaryOT using both entropic and quadratic OT produces very similar output resembling the ground truth when ε is chosen to be optimal. On the other hand, when ε is chosen to be too large we observe a visible worsening of performance, with more paths jumping between branches. As we also observed in terms of fate probabilities, the performance of StationaryOT with quadratic regularisation appears to degrade more gracefully than with the entropic regularisation.

To provide a quantitative assessment of performance, the natural metric to use is the 2-Wasserstein (W₂) distance on the space of laws on paths, as we also argue in [23]. We refer the reader to S1 Appendix for details on the definition of the W₂ distance. Since we work in the setting of discrete time steps, the squared Euclidean (L₂) distance between a pair of paths f, g is taken to be (17) Using the W₂ metric for laws on paths, we computed the error of each reconstruction relative to the ground truth. Importantly, we note that since we are dealing with finite samples, the expected W₂ distance between independent collections of sample paths from the same distribution will be nonzero. Thus, as done in [23] we compute a baseline error as the W₂ distance between independent samplings of 250 paths from the ground truth. In Fig 6B we show the average W₂ error over 5 resamplings of 250 paths, from which we note that StationaryOT with entropic or quadratic OT yields results that are close to the baseline in W₂ error when ε was chosen to be optimal. On the other hand, picking ε to be too large leads to a higher error for both methods, but with entropic OT performing significantly worse than quadratic OT.

Simulation setup and parameters.

To illustrate this, we consider a process with a drift field given by the sum of a potential-driven term and a non-conservative vector field, i.e. (18) Again, we work in and we take (19) (20) We pick h = 0.5, controlling how rapidly the field f decays and the location of the potential well in Ψ. In the first two dimensions of , particles can be thought of as diffusing on a radially symmetric potential field with a ring of wells located about the origin, and subject to a superimposed anticlockwise vector field that decays away from the origin. We show a surface plot of Ψ(x) and a vector field plot of f(x) in Fig 7A.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 7. Non-conservative simulation.

(a) Illustration of potential-driven (Ψ) and non-conservative (f) components of the overall drift v. (b) Examples of simulated particle trajectories following the drift-diffusion process. (c) Snapshot samples shown in the first two dimensions of , with source (R > 0) and sink (R < 0) regions indicated. (d) Comparison of fate probabilities towards the sinks in the first quadrant. (e). Correlation of estimated fate probabilities to ground truth fates with (λ = 0.25) and without incorporation of velocity data (λ = 0). (f) Summary of fate probability correlation as a function of λ ∈ [0, 1].

https://doi.org/10.1371/journal.pcbi.1009466.g007

We initialise particles following the initial distribution that are then subject to the drift-diffusion process with diffusivity σ² = 0.1. The minimum of the circular potential well is located along a cylinder of radius 0.721 about the origin in the first two dimensions of , and we treat all points outside this cylinder as a sink region, in which particles are removed at exponential rate 5. We sample 500 particles from this process and designate cells found within a ball of radius 0.1 about the origin to be source cells, and cells located in the sink region to be sink cells. Sink cells were assigned a flux rate R_i = −5, and source cells were assigned a uniform flux rate so that ∑_i R_i = 0, as in the previous example. We illustrate in Fig 7B some example trajectories from this simulation, and in Fig 7C we display the sampled snapshot along with the flux rates.

StationaryOT with and without velocity data.

For each sampled cell x_i, we obtain velocity estimates by evaluating the drift vector field v(x_i) at its location. We then formed two cost matrices: C_euc, the matrix of squared Euclidean distances, and C_velo the matrix of cosine similarities as defined in Eq (13). Both matrices were normalised to have unit mean. Note here that this normalisation is purely an empirical choice, and no corresponding normalisation of the cost was performed in the potential-driven case because of the theoretical motivation in the potential-driven case.

We constructed the optimal transport cost matrix to be a convex combination of the Euclidean and velocity cost matrices: and we took λ = 0.25, 0, respectively corresponding to StationaryOT with and without velocity information. Both entropic and quadratic OT were used to solve for couplings, with ε = 0.05 and ε = 0.5 respectively. Since the setting of this simulation is rotationally invariant in the first two dimensions, we choose to summarise our results in terms of the absorption probabilities for cells entering the region i.e. the set of sink cells in the first quadrant in (x₁, x₂). As shown in Fig 7D, the ground truth fate probabilities clearly capture the rotational component of drift, with the set of cells fated towards forming a curled shape. We observe that StationaryOT with velocity data produces results qualitatively capturing this effect, whilst neglecting velocity information leads to a symmetric fate profile that reflects only the potential-driven component as expected. To quantitatively compare fates, we computed as previously the Pearson correlation between the estimated fate probabilities and the ground truth. We show this in Fig 7E, from which we observe that StationaryOT with velocity data produces a markedly improved fate correlation (r = 0.953) compared to StationaryOT without velocity data (r = 0.631). Finally, in Fig 7F we show the fate correlation r as a function of the parameter λ ∈ [0, 1] that controls the composition of the cost matrix C. The correlation improves rapidly as λ is increased from 0 and attains a maximum before it declines slowly as λ is further increased towards 1. We conclude that even relatively small choices of λ can greatly improve the accuracy of the inferred fate probabilities.

Laws on paths.

As in the case of the potential-driven system, we may examine sample paths from the ground truth process as well as the estimates output by StationaryOT. We sample trajectories with the initial condition

We illustrate these in Fig 8 in the first two dimensions of . Again, we observe that incorporation of velocity estimates yields results that clearly reflect the rotational trajectories in the ground truth. On the other hand, without using velocity information, we observe sample paths consistent with only the potential-driven component. Additionally, for either choice of regularisation we observe that StationaryOT overestimates the rotational drift as cells settle into the potential well. This effect can be attributed to the fact that the cosine similarity cost of Eq (13) depends only on the orientation of the rotational field, and thus is unaware of its decay as cells drift towards the well. In this situation, we can only expect to capture the rotational field qualitatively rather than quantitatively. We suggest that possible remedies for this effect may include weighting entries of C_velo by velocity magnitudes or using an alternative velocity cost that is based on squared Euclidean distances.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 8. Inferred dynamics in the space of paths for non-conservative system.

Collections of 100 sample paths drawn from the ground truth process Eq (1), as well as StationaryOT output with and without velocity estimates for both entropic and quadratic OT. We indicate the initial condition π₀ as dots.

https://doi.org/10.1371/journal.pcbi.1009466.g008

Sensitivity to noise.

Finally, are interested in investigating the behaviour of StationaryOT when the provided velocity estimates are subject to additive noise, that is where α is a scale constant chosen such that v₀/α has order 1, i.e. the noise term is on the same order as the signal. We pick . We applied StationaryOT using both entropic and quadratic OT for values of η ∈ [0, 2] and choices of regularisation ε chosen in the range 10⁻² − 10⁰ (logarithmic) for entropic OT and 0.5 − 10 (linear) for quadratic OT.

For additional comparison, for each noise level η we also computed a transition matrix based solely on cosine similarities of velocity estimates to k-nearest neighbour (k-NN) graph edges using the scVelo package [20] in which the transition law for each cell x_i is (21) In the above, κ is a scale parameter controlling the level of directedness in the resulting transition law, with larger κ corresponding to increased directedness in the transition law. We used κ in the range 2.5–25 and all other parameters were taken to be defaults.

In each case, performance was summarised as we did previously in terms of the fate correlation for the set . We show results summarised over 10 independent repeats in Fig 9A and we observe that, as expected, performance degrades for all methods as the level of noise increases. However, StationaryOT with either entropic or quadratic regularisation consistently produces more accurate fate estimates compared to the scVelo method. We argue that this effect reflects the fact that StationaryOT is a global method and solves for transition laws that best agree with the inputs across the dataset. On the other hand, the scheme described by Eq (21) is local in that the transition law for a single cell can be determined by only considering a single velocity vector and a few neighbouring locations.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 9. Effect of parameter choices on inference for non-conservative system.

(a) Correlation of estimated fate probabilities to ground truth as a function of noise η. (b) Sensitivity of entropic and quadratic regularisations to the choice of ε.

https://doi.org/10.1371/journal.pcbi.1009466.g009

Additionally, for moderate levels of noise (η ∈ [0, 1]), we observe that the quadratic regularisation outperforms the entropic regularisation. Roughly speaking, we suspect that this results from the fact that the transition laws recovered by quadratic OT are concentrated on a sparse subset of cells, limiting the effect that an error in any single velocity estimate can have on the overall Markov chain. On the other hand, since entropic regularisation yields dense transition laws, errors can be propagated across the full support of .

In Fig 9B we examine the sensitivity of StationaryOT performance on the choice of the regularisation ε. As we observed in the potential-driven setting, the entropic regularisation appears to depend strongly on the choice of ε whereas a quadratic regularisation behaves similarly across the values of ε used.

Overview.

We now apply StationaryOT to the scRNA-seq atlas dataset generated by Shahan et al. [42], which comprises of gene expression data from 1.1 × 10⁵ cells from the first 0.5 cm of the Arabidopsis thaliana root tip. Stem cells occur close to the tip of the root and differentiate into ten distinct lineages (see Fig 10), with cells becoming increasingly differentiated as they increase in distance from the stem cells. Additionally, the terminal 0.5 cm of the root captures all tissue developmental zones, including the root cap, meristem, elongation zone, and part of the maturation zone. While new cells are constantly produced in the meristem, the bottom 0.5 cm is expected to be in equilibrium as cell division and elongation push existing differentiated cells out of the 0.5 cm section of interest, preserving a constant profile of cell types as illustrated in Fig 10A. Lineages and developmental zones are shown anatomically on the bottom 0.5 cm of the root in Fig 10B as well as on a UMAP embedding of the dataset in Fig 11.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 10. The Arabidopsis root tip system.

(a) While individual cells divide (green), elongate (blue), and are displaced from the bottom 0.5 cm (red) as the root grows, cell populations remain in equilibrium. (b) The structure of the Arabidopsis thaliana root tip by developmental zone (left) and lineage (right) (Illustrations modified from the Plant Illustration repository [44]).

https://doi.org/10.1371/journal.pcbi.1009466.g010

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 11. Arabidopsis atlas cell annotations and StationaryOT output.

Developmental zone (left) and lineage annotations (centre) shown on a UMAP embedding. Putative fate probabilities from StationaryOT with entropic regularisation are visualised on the right, where each cell is coloured by putative fate and its saturation based on the magnitude of that probability. For over 80% of cells the putative fate matched the annotation, with the magnitude of the probability increasing later in development.

https://doi.org/10.1371/journal.pcbi.1009466.g011

Application of StationaryOT.

For each cell x_i, daily growth rates g_i were estimated from imaging data of the growing meristem over a week-long period [43]. Using these growth rates and the proportion of cells expected to be actively dividing, we estimated that roughly 5% of the cells in each lineage would be replaced in a 6-hour period (Δt = 0.25) and selected the 5% of most differentiated cells from each lineage as sinks, as defined by pseudotime. For these sink cells we set g_i = 0, i.e. they are completely removed over the time interval. The remaining cells were assigned a flux rate R_i chosen to agree with the biological growth estimates, i.e. for each cell x_i we take R_i such that exp(R_i) = g_i.

We applied StationaryOT using both entropic and quadratic regularisations with parameters and respectively, where the scale factor is taken as the mean value of the squared Euclidean cost matrix C. We found our results to be robust to changes of a factor of two in the number of sinks, the time step size, and ε. Due to computational limitations of the standard implementation of the method, we applied StationaryOT to a subset of 10,000 cells sampled from the full dataset, though we also demonstrate two methods to scale the analysis to the full atlas. For completeness, we display in UMAP coordinates fate probabilities for each lineage in S3 Fig.

In Arabidopsis root development, cell lineage is fixed early in development [43]. Thus, for each cell x_i we may regard the lineage j corresponding to the largest fate probability, i.e. argmax_j p_ij as its putative fate. We checked whether these putative fates matched the manually curated atlas annotation, and used the magnitude of the corresponding fate probability, 0 ≤ p_ij ≤ 1, as a measure of the confidence of prediction. StationaryOT with quadratic and entropic regularisation performed similarly in terms of the percentage of cells where the putative fate matched the atlas annotation, matching 81% and 80% of cells respectively (see Fig 12). Both regularisations also performed similarly in terms of the magnitudes of the putative fates, with the entropic regularisation achieving an average of 69%, increasing from an average of 40% for cells in the meristem to an average of 84% for cells in the maturation zone and quadratic achieving an average of 65%, ranging from 38% in the meristem to 80% in the maturation zone (see S2 Fig). This trend can also be seen clearly by visualising the maximum fate probability of each cell on the UMAP (see Fig 13). As in many cases external estimates of growth are not available, we created alternate growth rates based on cell cycle genes. Despite differences between the estimates, StationaryOT performed similarly matching 79% and 78% of cells for entropic and quadratic regularisation respectively (see S1 Appendix).

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 12. Comparison of StationaryOT performance to other methods.

Proportion of cells where the maximum probability matches the annotation by developmental zone and lineage.

https://doi.org/10.1371/journal.pcbi.1009466.g012

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 13. Comparison of fate probabilities found by StationaryOT to other methods.

Fate probabilities for StationaryOT with entropic and quadratic regularisations compared to the annotation, as well as PBA and CellRank output. The colour indicates the maximum fate probability (putative fate) of each cell and the colour saturation shows the magnitude of the fate probability.

https://doi.org/10.1371/journal.pcbi.1009466.g013

Both choices of regularisation performed well on nine of the ten lineages, struggling only with mature procambium cells, as shown in Fig 12. We believe this is due to an inconsistency between the pseudotime and developmental zone annotations, where cells in the elongation zone received higher pseudotimes than cells in the maturation zone, resulting in them being incorrectly set as terminal states (see S1 Fig). Both regularisations were robust to changes in parameters, where the percentage of cells whose putative fate matched the annotation changed by no more than 2% when changing a parameter by a factor of two. StationaryOT with quadratic regularisation was particularly robust, with performance degrading by no more than 2% when multiple parameters were changed by up to a factor of five (see S1 Appendix).

Comparison to PBA.

Since population balance analysis (PBA) [18] addresses the same problem as StationaryOT, it is natural to evaluate its performance on the Arabidopsis root dataset. We show results for 1% sinks, Δt = 0.25 (6-hour time step), diffusivity D = 2.5, and k = 10 for the k-NN graph. These parameters were found to yield the best results over a parameter sweep (see S1 Appendix). PBA was on par with the StationaryOT methods, with 81% of putative fates matching the annotation, compared to 80% and 81% for the StationaryOT analyses (see Fig 12). Average fate probabilities were also similar, with PBA achieving an average of 64% compared to 65% and 69% for the StationaryOT methods (see S2 Fig). Like StationaryOT, the average fate probabilities increased as the tissue matured (see Fig 13). Given that PBA and StationaryOT are methodologically distinct, the fact that they perform similarly is a strong indication that the results reflect the underlying biology, rather than artefacts from the respective models.

In general however, we found PBA to be more sensitive to parameter values than StationaryOT. Assigning 5% of cells in each lineage as sinks for a 6-hour time step (Δt = 0.25) is biologically motivated in order to balance the number of cells created due to growth with the number of sinks. Using this sink selection scheme, PBA was found to perform poorly for the columella lineage, incorrectly assigning all columella cells a putative fate of lateral root cap. This lowered the overall percentage of putative fates that matched the annotation to 72% (see S1 Appendix). Only through an extensive parameter sweep did we find the combination of parameters that resulted in columella cells receiving the correct putative fate. We found PBA to be generally more sensitive to other parameter changes, matching 7% fewer cells to the annotation when a single parameter was changed by a factor of two compared to only 2% for the StationaryOT methods (see S1 Appendix). This may be of concern, since in many applications there may not be sufficient prior biological knowledge to distinguish between good and bad parameter choices.

Comparison to CellRank.

CellRank is a trajectory inference method that uses both transcriptomic similarity and RNA velocity data to estimate transition laws for cells [22]. The method consists of three key steps: computing cell state transition probabilities, inferring macrostates from the resulting Markov chain, and computing fate probabilities to these macrostates. For ease of comparison between other methods discussed here, our summary will focus mostly on the computation of transition probabilities.

CellRank computes a transition matrix from a k-NN graph using a combination of transcriptomic similarity and RNA velocity data. First, a k-NN graph is computed using cell transcriptomic similarities and then symmetrised. Edge weights are assigned based on similarity estimates between neighbouring cell states. The resulting graph is then converted into a matrix containing similarity estimates between neighbours. For each cell, transition probabilities are calculated from RNA velocity data by considering the correlation of the RNA velocity vector with displacement vectors corresponding to edges in the k-NN graph. These correlations are used to create a categorical distribution on the neighbours of the cell, giving transition probabilities. To better account for noise in the velocity data, the final transition probabilities are taken to be a linear combination of velocity-based probabilities and similarity-based probabilities.

We applied CellRank to the same 10,000 cell subset of the Arabidopsis dataset used for the StationaryOT and PBA analyses. Using the output transition matrix, we computed fate probabilities and assigned putative fates as previously described. In terms of putative fates, we found that CellRank matched 73% of cells to the atlas annotation, compared to 80% and 81% achieved by the StationaryOT analyses. The main differences occurred in the lateral root cap and xylem tissues (see Fig 12). CellRank also had less confidence in fate prediction, having an average fate probability of 45% compared to greater than 60% for all other methods (see S2 Fig) and probabilities remained low through the elongation zone (see Fig 13). Finally, we note that CellRank uses a mixture of a directed transition matrix derived from RNA velocity and an undirected transition matrix computed from expression similarity.

Repeated subsampling.

We first randomly partition the dataset of interest into k subsets of size approximately 10⁴, or such that the computation time for StationaryOT is acceptable. Fates are computed for each subset, and this procedure is repeated j times with repeated random partitioning. We then average the computed fates on a cell-by-cell basis, to produce aggregated fate probabilities.

We applied this approach to the full 1.1 × 10⁵ cell atlas, partitioning it into 10 subsets and applying StationaryOT separately to each subset (see Fig 14). This was repeated 10 times to account for sampling error. Between the fates found directly for each subset and the consensus fates in the full atlas, 97% of cells shared the same putative fate and the maximum fate values had a correlation of 0.96. Accounting for all fate values, the correlation rose to 0.99.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 14. Output of StationaryOT on full atlas dataset.

Atlas annotation on the full dataset (1.1 × 10⁵ cells) shown in UMAP coordinates compared to fate probabilities computed on the full dataset respectively using the subsampling approach (using entropic regularisation for each subproblem) and memory-efficient GPU implementations of StationaryOT with entropic and quadratic regularisations.

https://doi.org/10.1371/journal.pcbi.1009466.g014

Memory-efficient GPU implementation with KeOps.

For both entropic and quadratic regularisations, algorithms for solution of the optimal transport minimisation problem can be implemented using the KeOps library [45] so as to avoid storing all N × N matrices in memory. Along with GPU acceleration, this allows StationaryOT to be applied directly to datasets with many more cells than can be handled by the standard implementation due to memory constraints.

In the case of entropy-regularised optimal transport, the Sinkhorn algorithm can be implemented in a straightforward manner so that the transport plan (a matrix of size N × N) is parameterised in closed form by two dual variables u and v (vectors each of length N) [31]. From here, one may construct a linear system to solve for fate probabilities of the form Ax = b where A is again parameterised by the dual variables (u, v) and thus not explicitly stored in memory. For quadratically-regularised optimal transport, a similar representation of the problem in terms of dual variables holds. We provide an implementation of the semi-smooth Newton method proposed in [41, Algorithm 2] that utilises the KeOps library. As mentioned earlier, quadratically regularised optimal transport has the property that the transport plans (and hence the transition laws of cells) will be sparse. In addition to being more interpretable, for large numbers of cells this means additional computational advantages, especially in terms of directly storing the entries of the sparse transition matrix, and computing the fate probabilities.

We applied StationaryOT to the full 1.1 × 10⁵ cell Arabidopsis root dataset using the KeOps implementation with GPU acceleration, using both entropic and quadratic regularisations. We found that solution of the optimal transport problem took roughly 15 and 20 minutes using entropic and quadratic regularisations respectively. We used the same parameter choices as used for the subsampled dataset. For an entropic regularisation, the resulting transportation plan was dense (effectively all entries were nonzero). In contrast, the quadratic regularisation yielded a very sparse transportation plan as expected (0.17% nonzero entries). Computation of the fate probabilities for the entropic regularisation was significantly more time-consuming than for the quadratic regularisation, taking approximately 10 minutes and 2 minutes respectively. We display the fate probabilities for entropic and quadratic regularisations in Fig 14. The difference in the runtimes reflects the fact that, compared to dense systems, sparse systems of linear equations can be solved much more efficiently using iterative methods. We compared the fates found for a 10,000 cell subset to the fates for StationaryOT with both entropic and quadratic regularisation on the full atlas and found them to perform similarly. For both methods, 90% of cells shared the same putative fate as the 10,000 subset and both had a 0.97 correlation for fate magnitudes accounting for all fates. With the entropic regularisation, the putative fate values were slightly higher correlated with the 10,000 cell subset, achieving a correlation of 0.92 compared to 0.87 for quadratic. All computations for the full atlas dataset were done on a Google Colaboratory instance with a 16GB NVIDIA Tesla V100 GPU.