Modeling Cell Type Developmental Trajectory using Multinomial Unbalanced Optimal Transport

All living organisms are subject to the effects of genetics. The genetic information stored within cells dictates the way of life, functioning like a sequence of commands that shape every organism into its eventual form. Since the discovery of DNA, the study of molecular genetics has revealed the dynamic nature of cellular development. A key aspect of this dynamic process is the variation in gene expression among different cell types and at different stages of development. Building on this understanding, biologists have studied the developmental processes of stem cells and differentiated cells (Keller, 2005), as well as stem cell reprogramming in the pursuit of new methods for stem cell therapy (Mothe et al., 2012). Although much research has been conducted on the molecular mechanisms underlying stem cell differentiation (Dixon et al., 2015), there is a growing interest for mathematical models that are more efficient and adaptable to heterogeneous datasets, in order to better understand these complex developmental processes. Waddington (2014, 2015) modeled the changes in gene expression as a landscape that influences the developmental trajectory of cells. In this analogy, various potential developmental pathways available to a cell are similar to different railway tracks, with each track representing a unique developmental trajectory. The observed gene expression patterns and the developmental costs are similar to the layout and travel costs of railway tracks, each guiding the cell’s developmental journey. At the population level, hundreds and thousands of cells undergo changes in gene expression based on their individual trajectories. In experimental settings, observations are typically made at multiple time points during cellular development. Single-cell RNA sequencing (scRNA-seq) enables the measurement of gene expression levels of individual cells at different time points (Tanay and Regev, 2017). However, it is often challenging to model the gene expression profile of a single cell and recover its developmental trajectory by analyzing scRNA-seq data from different time points, since for each cell we can only observe its expression profiles at one particular time instead of the whole trajectory. Several existing methods (Kester and Van Oudenaarden, 2018; Farrell et al., 2018; Wagner et al., 2018; Cao et al., 2019) aim to model such trajectories. Nevertheless, these methods do not directly address two main challenges. Firstly, they often focus on inferring trajectories for individual cells, potentially missing the broader developmental dynamics of the entire cell population. This can result in scenarios where individual optimization is not feasible for the whole population. Interactions and resource allocation among cells must be considered. For example, while a cell may differentiate into its expected optimal type under ideal conditions with ample resource, limited resources due to large number of cells can prevent all cells from reaching their expected state. Secondly, different cell types may proliferate at varying growth rates, and this heterogeneity in cellular growth rates affects the relative proportions of different cell types over time.

To tackle the two challenges mentioned above, Schiebinger et al. (2019), introduced Waddington-OT, a novel method that applies the concept of optimal transport to single-cell RNA sequencing (scRNA-seq) data analysis. Optimal transport, initially proposed by Monge (1781), was designed to model the transition process between different distributions. The Waddington-OT method models the developmental trajectory of cells in two key steps. Firstly, it employs optimal transport to infer temporal couplings between adjacent time points, effectively tracing the transition of the cell population from one time point to the next. Secondly, it introduces “unbalancedness” into the optimal transport framework. Traditional optimal transport models assume the conservation of total “mass” during transitions between states. However, in many biological processes, such as cellular growth, this conservation assumption is not valid, as the total mass can either increase or decrease by absorbing external nutrients or by death of cells. Unbalanced optimal transport addresses this by accommodating changes in mass, thus allowing the model to incorporate variations in cellular growth rates. This adaptation makes the Waddington-OT method a more realistic and flexible tool for modeling cell population dynamics over time.

The ability to model trajectories of cell types, which are groups of cells with similar expression profiles and functions, is particularly important in fields such as developmental biology, cell replacement therapy, and drug discovery (Keller, 2005). As opposed to modelling individual cells, analyzing different cell types reduces the model complexity compared to modelling individual cells while still preserving biologically meaningful heterogeneity. For example, in cell replacement therapy, trajectory modeling can help identify progenitor cell types with the highest potential for differentiation into the required cell types for therapeutic purposes. In this paper, we introduce Trajectory Inference with Multinomial Optimal Transport (TIMO), a novel approach that shifts the focus from individual cells to modeling trajectories of cell types. TIMO incorporates a unique statistic based on optimal transport, leveraging the benefits of a multinomial framework.

The use of a multinomial model offers several practical advantages over continuous models. Firstly, in Multinomial Optimal Transport, as detailed in Section 3, the cost functions are specifically designed to assign zero cost to cells that remain within the same cell type. In contrast, in Schiebinger et al. (2019), it is impossible for two cells of the same type to have identical gene expression profiles, leading to a non-zero cost between these two cells. This distinction makes the multinomial approach more effective in accurately detecting biological change points and in studying cellular sub-populations, i.e., cell types. Additionally, by focusing on the trajectories of cell types rather than individual cells, we significantly reduce the computational complexity from $\mathcal{O}(n^{2})$ to $\mathcal{O}(d^{2})$, where $n$ is the number of cells, typically in the thousands, and $d$ is the number of cell types, usually only around 10 (Pham et al., 2020).

To showcase the advantages of the TIMO framework, we demonstrate that TIMO can effectively detect biological change points and quantify cellular differentiation events during developmental trajectories. To empirically validate our approach, we applied TIMO to a comprehensive scRNA-seq dataset containing over 250,000 cells containing multiple cell types originated from mouse embryo fibroblast (MEF). MEF cells are commonly used in research related to modern medicine. Our model yielded the following key findings:

1. 1.

   We detected two biological change points corresponding to the emergence of mesenchymal epithelial transition (MET) cells and stromal cells, both of which have been confirmed by wet lab results in Schiebinger et al. (2019) as underlying points of cellular differentiation.

2. 2.

   We quantified cell type developmental trajectory and identified that MEF cells are the most potent cell type capable of differentiating into MET cells, and that both MEF and MET cells differentiate into stromal cells. These findings highly match the biological results from the original experiment (Schiebinger et al., 2019).

3. 3.

   TIMO successfully distinguishes the effects of growth from the effects of cellular differentiation. Notably, we detected no change during the known period of induced pluripotent stem (IPS) cell proliferation (i.e., cellular growth/division), unlike other modern statistical change point detection methods, which incorrectly identified growth periods as change points.

To the best of our knowledge, TIMO is the first method that can detect biological change points while accounting for the effects of cellular growth.

The rest of this paper is structured as follows. In Section 2, we provide definition on cell type developmental trajectories. In Section 3, we provide mathematical details on trajectory inference using TIMO. In Section 3.1, we introduce the test statistic and its application in biological change point detection. In Section 3.2, we present procedures on trajectory inference given a finite sample. Numerical results on simulated and real data sets are presented in Section 4 and Section 5. We end with a discussion in Section 6.

The study of cellular differentiation dynamics, encompassing landscapes, fates, and trajectories, poses formidable challenges for scientists. Deciphering these developmental pathways is a key area of study in scRNA-seq. To capture the gene expression profile of a cell, we employ a $G$-dimensional vector in gene expression space, where $G$ represents the total number of genes, where $G$ is often around $20,000$. Additionally, we assume that the observed developmental trajectory is limited to a finite number of time points, specifically $t=0,1,2,\ldots,T$, where $T$ is a finite positive integer.

In this section, we develop a new method, Trajectory Inference with Multinomial Unbalanced Optimal Transport (TIMO), to model the trajectories of different cell types under a discrete unbalanced optimal transport framework.

We first consider an ideal case, where the cells do not divide into multiple cells or different cell types divide at the same rate across time. If there is no differentiation among all cell types from time $t$ to $t+1$, one has $Q_{t}=Q_{t+1}$. We then have $\pi_{jk}^{t,t+1}=0$ (or equivalently $h_{kj}^{t+1|t}=0$) for all $j\neq k$, and $\pi_{jj}^{t,t+1}=q_{t,j}=q_{t+1,j}$ (or equivalently $h_{jj}^{t+1|t}=1$) for all $j$. In other words, the transition probability matrix $\pi^{t+1|t}$ is an identity matrix, indicating that any cell type will not differentiate to other types from time $t$ to $t+1$.

We now utilize the optimal transport tool to find the joint probability matrix $\pi^{t,t+1}$. Define $\Gamma_{d,d}^{\mu,\nu}\coloneqq\{\pi\geq 0;\pi\in\mathbb{R}^{d}\times\mathbb{R}^{d},\sum_{j=1}^{d}\sum_{k=1}^{d}\pi_{jk}=1,\pi 1_{d}=\mu,\pi^{{\mathrm{\scriptscriptstyle T}}}1_{d}=\nu\}$, where $\mu$ and $\nu$ are probability measures on $\{1,2,\ldots,d\}$ and $1_{d}$ is a $d$-dimensional column vector with each element $1$.

A preliminary approach involves determining the joint distribution $\pi^{t,t+1}$ by solving the following optimal transport problem,

$$\pi^{t,t+1}\coloneqq\operatorname*{arg\,min}_{\pi\in\Gamma_{d,d}^{Q_{t},Q_{t+1}}}\sum_{j=1}^{d}\sum_{k=1}^{d}m_{jk}\pi_{jk},$$

(1)

where $m_{jk}$ is the transport cost from cell type $j$ to cell type $k$. Here, the costs $m_{jk}=0$ if $j=k$, and $m_{jk}>0$ otherwise. The $\pi^{t,t+1}$ is called the transport map and $W(Q_{t},Q_{t+1})=\min_{\pi\in\Gamma_{d,d}^{Q_{t},Q_{t+1}}}\sum_{j=1}^{d}\sum_{k=1}^{d}m_{jk}\pi_{jk}$ is the transport cost from $Q_{t}$ to $Q_{t+1}$. We require $\pi^{t,t+1}\in\Gamma_{d,d}^{Q_{t},Q_{t+1}}$ because for any joint distributions of $X_{t}$ and $X_{t+1}$, their marginal distributions must be $Q_{t}$ and $Q_{t+1}$.

From a biological perspective, the costs represented by $m_{jk}$ can be interpreted as the energy or resources required for cellular differentiation from type $j$ to type $k$. In practice, numerous studies have revealed that differentiating cells undergo a substantial energy acquisition process (Fard et al., 2016; Guo et al., 2017; Banerji et al., 2013). In our context, these distinctions are reflected in the varying values of $m_{jk}$, which are determined by difference in gene expressions among various cell types. We will discuss how to choose $m_{jk}$’s at the end of this subsection.

Recent work by Shi et al. (2022) has demonstrated that cellular differentiation is intricately linked to energy availability. Building upon these biological insights, the differentiation process of a single-cell population can be conceptualized as an efficient allocation of available energy. This allocation aligns with the optimization problem described in (1).

However, in the real world, the optimal transport problem (1) can not be directly applied to model the cell type developmental trajectory. The main reason is that different cell types may grow at different rates (Yang et al., 2014), which makes the marginal distribution constraints invalid. To illustrate the idea, define $g_{t,j}$ the growth rate of cell type $j$ at time $t$. We have $q_{t+1,j}=\frac{q_{t,j}g_{t,j}}{\sum_{l=1}^{d}q_{t,l}g_{t,l}}$ and

$$Q_{t+1}=g_{t}(Q_{t})=\left(\frac{g_{t,1}q_{t,1}}{\sum_{k=1}g_{t,k}q_{t,k}},\frac{g_{t,2}q_{t,2}}{\sum_{k=1}g_{t,k}q_{t,k}},\ldots,\frac{g_{t,d}q_{t,d}}{\sum_{k=1}g_{t,d}q_{t,d}}\right).$$

(2)

Based on Zhang et al. (2021), the developmental process between $t$ and $t+1$ can be considered as a two-phase procedure. The first phase is the growth phase, where $Q_{t}$ grows to $g_{t}(Q_{t})$. The second phase is the transport phase, where $g_{t}(Q_{t})$ changes to $Q_{t+1}$ following the optimal transport principle. If the growth function $g_{t}(\cdot)$ is known, by modifying the constraint in (1), one can obtain the transport map $\pi^{t,t+1}$ by solving the following optimal transport problem

$$\pi^{t,t+1}=\operatorname*{arg\,min}_{\pi\in\Gamma_{d,d}^{g_{t}(Q_{t}),Q_{t+1}}}\sum_{j=1}^{d}\sum_{k=1}^{d}m_{jk}\pi_{jk}.$$

One key challenge is that the growth rates $g_{t,j}$’s cannot be directly identified based on the observed marginal distributions $Q_{t}$ and $Q_{t+1}$, making the above approach infeasible.

To overcome this challenge, we utilize the Kullback-Leibler (KL) divergence penalties to relax the marginal distribution constraints, effectively mimicing the effects of cellular growth on probability distributions. In particular, we can infer the cellular development trajectory by solving the following discrete unbalanced optimal transport problem

$$\pi^{t,t+1}=\operatorname*{arg\,min}_{\pi\in\Gamma_{d,d}^{+,Q_{t+1}}}\{\sum_{j=1}^{d}\sum_{k=1}^{d}m_{jk}\pi_{jk}+\lambda KL\left(\pi 1_{d}||Q_{t}\right)\},$$

(3)

where $\Gamma_{d,d}^{+,Q_{t+1}}\coloneqq\{\pi>0;\pi\in\mathbb{R}^{d}\times\mathbb{R}^{d},\sum_{i=1}^{d}\sum_{j=1}^{d}\pi_{ij}=1,\pi^{\mathrm{\scriptscriptstyle T}}1_{d}=Q_{t+1}\}$. Here $KL$ denotes the KL divergence, which is defined as $KL(P||Q)=\sum_{k=1}^{d}p_{k}\log\left(\frac{p_{k}}{q_{k}}\right)$ for any two multinomial probability measures with $d$ categories $P=(p_{1},p_{2},\ldots,p_{d})^{\mathrm{\scriptscriptstyle T}}$ and $Q=(q_{1},q_{2},\ldots,q_{d})^{\mathrm{\scriptscriptstyle T}}$. In LABEL:sec:comp_algo of the supplementary material, we provide a detailed algorithm for solving the problem in (3).

In problem (3), the equality constraint only holds for the column sum constraints $\pi^{\mathrm{\scriptscriptstyle T}}1_{d}=Q_{t+1}$, whereas the KL divergence penalization relaxes the row sum constraints $\pi 1_{d}=Q_{t}$ to account for different growth rates of cell types. When $\lambda=0$, the model does not enforce the row sum constraints at all. The optimal solution is selected from $\Gamma_{d,d}^{+,Q_{t+1}}$. As $\lambda$ increases, the row sum of the optimal solution $\pi^{t,t+1}$ more closely resembles $Q_{t}$, that is, the discrepancy in growth rates among cell types approaches $0$. When $\lambda=\infty$, it is equivalent to imposing the row sum constraints $\pi 1_{d}=Q_{t}$, which corresponds to the case that every cell grows or dies at the same rate.

In this section, we conduct a simulation study to evaluate the performance of TIMO. We consider a population of cells spanning a time horizon from $t=0,1,\ldots,T$, with $T=50$. In this population, there are a total of $d=10$ distinct cell types, each expressing $G=50$ different genes.

We generate the cell type $X_{t}$ from a multinomial distribution $Q_{t}=(q_{t,1},q_{t,2},\ldots,q_{t,10})^{{\mathrm{\scriptscriptstyle T}}}$. Initially, $Q_{0}$ is set as $(0.1,0.1,\ldots,0.1)^{\mathrm{\scriptscriptstyle T}}$. We then define $g_{t}(Q_{t})$ for $t=0,1,\ldots,T-1$ based on Equation 2. Here, the relative growth rate of the $j$th cell type at time $t$ is set as $g_{t,j}=\exp(\nu\sin(\frac{t+j-1}{d})\pi)$, where we consider two settings $\nu=0.1$ and $\nu=0.25$.

Given $Q_{0}$ and $g_{t,j}$, we generate $Q_{1},Q_{2},\ldots,Q_{50}$ as follows:

$$Q_{t+1}=\begin{cases}\frac{c_{1}^{{\mathrm{\scriptscriptstyle T}}}g_{t}(Q_{t})}{\|c_{1}^{{\mathrm{\scriptscriptstyle T}}}g_{t}(Q_{t})\|_{1}}&t\in\{10,30\}\\ \frac{c_{2}^{{\mathrm{\scriptscriptstyle T}}}g_{t}(Q_{t})}{\|c_{2}^{{\mathrm{\scriptscriptstyle T}}}g_{t}(Q_{t})\|_{1}}&t\in\{20,40\}\\ g_{t}(Q_{t})&t\notin\{10,20,30,40\},\end{cases}$$

where $c_{1}=e^{\eta}(1_{5},-1_{5})^{{\mathrm{\scriptscriptstyle T}}}$ and $c_{2}=e^{\eta}(-1_{5},1_{5})^{{\mathrm{\scriptscriptstyle T}}}$ with $1_{d}$ denotes a $d$-dimensional vector with all entries equal $1$. Here, we generate $Q_{t}$ differently at time points $t=10,20,30,40$ to mimic the cell type differentiation at these time points. Given $Q_{0},Q_{1},\ldots,Q_{T}$, we generate the gene expression $Y_{t}\in\mathbb{R}^{50}$ based on the conditional distribution $Y_{t}|X_{t}$. In particular, $Y_{t}|X_{t}=j$ is generated from the multivariate normal distribution $\mathcal{N}_{G}(\mu_{j},\Sigma_{j})$, with $\mu_{j}=\{0.5(j-d)-1\}1_{G}$ and $\Sigma_{j}=I_{G}$.

To define the underlying cost, we first reduce $Y_{t}$ to two dimensions by the PHATE procedure. The underlying cost $m_{jk}=||\zeta_{j}-\zeta_{k}||_{2}^{2}$, where $\zeta_{j}\in\mathbb{R}^{2}$ and $\zeta_{k}\in\mathbb{R}^{2}$ are centroids of reduced two-dimensional expressions of the $j$th and $k$th cell types, respectively. Given the underlying costs $m_{jk}$’s and the marginal distributions $Q_{t}$’s, the underlying transport plan $\pi^{t,t+1}$ is defined as: $\pi^{t,t+1}\coloneqq\operatorname*{arg\,min}_{\pi\in\Gamma_{d,d}^{+,Q_{t+1}}}\{\sum_{j,k}^{d}\pi_{jk}m_{jk}+\lambda KL(\Pi 1_{d}||Q_{t})\},$ with default value $\lambda=1$.

To generate the data, we sample $n$ cells at each time point, where we consider $n=1000$ and $2000$. Denote $x_{t,i}$ the cell type of the $i$th sampled cell at time $t$ and $y_{t,i}$ its gene expression. Each $x_{t,i}$ is independent and identically distributed (i.i.d) from $Q_{t}$, and $y_{t,i}$ is i.i.d. from $\mathcal{N}_{G}(\mu_{x_{t,i}},\Sigma_{x_{t,i}})$. We can obtain the empirical distributions $\widehat{Q}_{t}=\{\widehat{q}_{t,j}\}_{1\leq j\leq d}^{\mathrm{\scriptscriptstyle T}}$, where $\widehat{q}_{t,j}=n^{-1}\sum_{i=1}^{n}1(x_{t,i}=j)$. To estimate the centroid $\zeta_{j}$, we first project $y_{t,i}$ to the two-dimensional space using the PHATE procedure, denoted by $\tilde{y}_{t,i}$. For each cell type, its centroid $\widehat{\zeta}_{j}$ is estimated as $\widehat{\zeta}_{j}=\{\sum_{t=0}^{T}\sum_{i=1}^{n}1(x_{t,i}=j)\}^{-1}\sum_{t=0}^{T}\sum_{i=1}^{n}\tilde{y}_{t,i}1(x_{t,i}=j)$. We can then obtain the cost estimate $\widehat{m}_{jk}=||\widehat{\zeta}_{j}-\widehat{\zeta}_{k}||_{2}^{2}$. Finally, we estimate the transport plan $\pi^{t,t+1}$ by solving the following optimization problem

$$\widehat{\pi}^{t,t+1}\coloneqq\operatorname*{arg\,min}_{\pi\in\Gamma_{d,d}^{+,\widehat{Q}_{t+1}}}\{\sum_{j,k}^{d}\pi_{jk}\widehat{m}_{jk}+\lambda KL(\Pi 1_{d}||\widehat{Q}_{t})\}.$$

We first examine the estimation error for the transport plan $\pi^{t,t+1}$. We conduct $1000$ Monte Carlo runs and we report the average of $\frac{1}{4}\sum_{t\in\mathcal{C}}\|\widehat{\pi}^{t,t+1}-\pi^{t,t+1}\|_{F}^{2}$ at four change points $t\in\mathcal{C}=\{10,20,30,40\}$ and $\frac{1}{46}\sum_{t\notin\mathcal{C}}\|\widehat{\pi}^{t,t+1}-\pi^{t,t+1}\|_{F}^{2}$ at all $46$ non-change points $t\notin\mathcal{C}$ for different settings of $\eta$ and $\nu$ in Table 1. The results show the estimation errors become smaller as $n$ increases, suggesting consistency of the proposed estimator.

We next evaluate the proposed change point detection procedure in Section 3.1. In this setting, we have biological change points at time $t=10,20,30,40$. We also compare the performance with two other change point detection methods for multinomial data. In particular, we consider the method in Wang et al. (2018), which is designed for detection of multiple change points for multinomial time series data. We also compare with Matteson and James (2014), which is designed for change point detection for multivariate time series data. Here, we treat the probability weight vector at every time point as one vector observations in the time series and perform their method.

Suppose the underlying change points are $\mathcal{C}=\{t_{1},t_{2},\ldots,t_{L}\}$ and the estimated change points are $\widehat{\mathcal{C}}=\{\widehat{t}_{1},\widehat{t}_{2},\ldots,\widehat{t}_{L^{\prime}}\}$. Similar to Aminikhanghahi and Cook (2017) and Truong et al. (2020), we report the following quantities to evaluate the performance:

$$\text{precision}=\frac{|\mathcal{C}\cap\widehat{\mathcal{C}}|}{|\widehat{\mathcal{C}}|},\quad\text{recall}=\frac{|\mathcal{C}\cap\widehat{\mathcal{C}}|}{|\mathcal{C}|},\quad\text{F-score}=\frac{2\times\text{precision}\times\text{recall}}{\text{precision}+\text{recall}}.$$

Among these three quantities, F-score is considered the most important measure for change point detection. The results are summarized in Tables 2 and 3.

Our proposed method TIMO showcases outstanding performance across various scenarios. It is evident that TIMO signficantly leads in precision and F-score in every instance. While ECP excels in recall, TIMO consistently attains values near the value $1$.

We analyze the scRNA-seq data (Schiebinger et al., 2019) which can be downloaded from NCBI Gene Expression Omnibus (https://www.ncbi.nlm.nih.gov/geo/query/acc.cgi?acc=GSE122662). The cells were originated from mouse embryonic fibroblasts (MEFs) in serum conditions, and this cellular populations contained multiple cell types due to differentiation. The experiment ran over 18 days during which a total of $259,155$ cells were collected. We pre-processed the scRNA-seq data using R package Seurat (Satija et al., 2015). Specifically, we first remove low-quality cells (e.g. cells with very few genes or cell doublets/multiplets with an aberrantly high gene count). Next, the gene expression levels of each cell were normalized by the total expression levels. The normalization is followed by multiplicative scaling and log transformation. Finally, the Seurat recipe selects only the genes that demonstrate high cell-to-cell variation, i.e., the potentially biologically significant genes. After pre-processing, there are a total of $182,908$ cells collected every $12$ hours from day $0$ to day $18$. The number of cells at each time point is provided in LABEL:fig:samplesize in the supplementary material. For each cell, we obtaine its gene expression level over $G=216$ genes, i.e., $Y_{t}\in\mathbb{R}^{216}$. For simplicity, we denote these observation time $\{0,1,2,\ldots,36\}$, where the time point $t$ denotes the $0.5t$-th day.

We leverage the cell sets provided by Schiebinger et al. (2019) to assign labels to each cell based on its corresponding cell type. Our dataset encompasses the following $d=7$ distinct cell types: Mouse Embryonic Fibroblasts (MEFs), Mesenchymal-Epithelial Transition (MET) Cells, Induced Pluripotent Stem (IPS) Cells, Stromal Cells, Epithelial Cells, Neural Cells and Trophoblasts.

We first perform dimensionality reduction using PHATE. Figure 5 displays the PHATE-embedded data, featuring labels for both time and cell types. To outline the observed patterns in Figure 5, we can observe the dominance of MEFs in the population during the earlier days of development. As developmental progresses, distinct branches corresponding to the final stages of development emerge. These branches are populated by various cell types, including epithelial cells, trophoblasts, and induced pluripotent stem cells (IPS cells). Notably, the middle region is predominantly occupied by mesenchymal-epithelial transition (MET) cells, signifying transitional states on the path to becoming epithelial cells.

We now demonstrate the performance of our proposed method, TIMO, on the MEF dataset. TIMO detects two biological change points on days 4.5 and 10 (see LABEL:fig:ot_cost in the supplementary material for a plot of the statistics $\{W_{0},W_{1},\ldots,W_{T-1}\}$). Transport maps for these two time points are presented in Figure 6 and Figure 7. From day 4.5 to 5, the transport map highlights the early stage of cells transitioning into the MET state. According to Schiebinger et al. (2019), the first significant divergence occurs around day 1.5 and intensifies over the subsequent days, notably around day 4. From day 10 to 10.5, the transport map illustrates a substantial portion of the population undergoing differentiation, with many cell types transitioning into stromal cells. This aligns with the peak population of stromal cells around days 10-11, as reported by Schiebinger et al. (2019). In Figures LABEL:fig:tempora0–LABEL:fig:tempora17_5 of the supplementary material, we also display all the transport maps corresponding to the remaining non-change points. We can see that the transport maps predominantly take the form of diagonal matrices, except at the two detected change points. Importantly, our model does not identify the increase in IPS cell proportion and decrease in stromal cell proportion as change points, further highlighting that our model effectively distinguishes the effects of growth from differentiation.

Next, we compare our results with the detailed biological findings from Schiebinger et al. (2019). From a biological standpoint, this entire timeline corresponds to the reprogramming of IPS cells. The developmental process of this population reveals two critical branches. In the initial phases of IPS cell reprogramming, the original MEF population bifurcates into two branches: MET and stromal. Additionally, Schiebinger et al. (2019) established that IPS cells, along with other cell types like epithelial cells, neural cells, and trophoblasts, also originate from the MET branch. However, only a very small percentage (less than 10%) of MET cells undergo differentiation into IPS cells or neural cells. The increase in the proportion of IPS cells is primarily driven by rapid cell growth and proliferation rather than differentiation. In contrast, the stromal branch emerges at a later stage, with its population peaking around days 10-11. However, beyond this peak, the proportion of stromal cells starts to decline due to reduced cellular growth and increased apoptosis (i.e., cellular death). TIMO successfully distinguishes the effects of cellular growth from cellular differentiation.

Finally, we compare the change points detected by TIMO with those detected by the other two methods. The ECP method detected change points on days 4.5, 8.5, 10, and 15.5, while the MN method detected change points on days 4.5, 8, 10, 11.5, 14.5, and 16.5. Both methods fail to distinguish non-change points on days 14-16, whereas TIMO succeeds by accounting for growth rates with unbalanced regularity in its minimization problem.

Our study introduces TIMO, an innovative method for modeling the development trajectory of cell types using unbalanced optimal transport. Our method can be used to identify key biological change points and elucidating the temporal evolution of various cell types. We applied TIMO to a real scRNA-seq dataset from mouse embryonic fibroblasts. The results accurately pinpointed significant changes in cell type developmental process, supported by experimental observations. The transition probabilities among different cell types are consistent with the experimental observations.

While our method has shown promising results in cell type trajectory modeling, it also has some limitations and provides opportunities for future research. First, we computed the transition matrix at the cell type level. While this approach is computationally more efficient than optimal transport analysis at the single-cell level, it overlooks the heterogeneity among cells within the same cell type. Second, we applied TIMO to datasets with known cell types for each individual cell. Future improvements could involve integrating clustering techniques to make our method applicable to a broader array of datasets without predefined labels. Additionally, our approach for detecting multiple change points currently relies on constructing statistics using only two adjacent time points. This has been sufficient for our real-world applications due to significant changes in the underlying biological process. However, in scenarios with more moderate changes, it might be beneficial to incorporate moving window techniques to enhance stability and accuracy. Finally, in the field of unbalanced optimal transport, determining the appropriate regularization parameter $\lambda$ through a data-adaptive method remains an unsolved challenge.