Learning stochastic dynamics and cell-fate landscapes from single-cell snapshots via optimal transport

The temporal dynamics and stochasticity of gene expression are critical to cell fate decisions, yet integrating snapshot omics data across multiple time points remains a major challenge. Here, we introduce DiffusionOT, a dynamic machine learning framework that infers cellular trajectories from multi-time point single-cell transcriptomics by incorporating stochastic effects. DiffusionOT transforms stochastic differential equations into ordinary differential equations, using optimal transport and neural networks to solve a high-dimensional landscape model. Through an unsupervised learning of the stochastic force in the data, DiffusionOT allows robust inference of the underlying stochastic dynamics of cell-state transitions. The framework includes a stochastic trajectory analysis module for lineage tracing and a gene perturbation module for in silico knockout and overexpression experiments. Benchmarks on simulated and four real-world datasets, including a spatial Stereo-seq dataset, demonstrate DiffusionOT's accuracy and efficiency in inferring state-transition velocities, cellular trajectories, population growth, gene regulatory networks, and cell-fate landscape.