STORIES: learning cell fate landscapes from spatial transcriptomics

In dynamic biological processes such as development, spatial transcriptomics is revolutionizing the study of the mechanisms underlying spatial organization within tissues. Inferring cell fate trajectories from spatial transcriptomics profiled at several ti me points has thus emerged as a critical goal, requiring novel computational methods. Wasserstein gradient flow learning is a promising framework for analyzing sequencing data across time, built around a neural network representing the differentiation pote ntial. However, existing gradient flow learning methods cannot analyze spatially resolved transcriptomic data. Here, we propose STORIES, a method that employs an extension of Optimal Transport to learn a spatially informed potential. We benchmark our approach using three large Stereo-seq spatiotemporal atlases and demonstrate superior spatial coherence compared to existing approaches. Finally, we provide an in-depth analysis of axolotl neural regeneration and mouse gliogenesis, recovering gene trends for known markers as Nptx1 in neuron regeneration and Aldh1l1 in gliogenesis and additional putative drivers.

Spatial transcriptomics technologies are revolutionizing the study of how cells organize within tissues1. Techniques based on high -throughput sequencing have enabled the unbiased discovery of gene expression patterns within their spatial context. For instance, recent studies have revealed previously unknown spatial organization at the tumor -microenvironment interface in melanoma and Alzheimer's disease amyloid plaque microenvironment 2,3. The most widely used spatially-resolved sequencing techniques (e.g. 10X Visium) measure spots larger than the typical cell size. However, recent technological developments based on barcoded arrays like Stereo -seq and HDST have reached single -cell resolution, effectively bridging functional and structural characterizations of the cell4,5. Recent works have leveraged Stereo-seq to produce large spatiotemporal atlases of various biological processes by profiling a system with spatial transcriptomics at several points in time4,6,7. These datasets are ideal for studying cellular dynamics within the tissue during processes such as development and the .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint onset of complex diseases, where cells undergo coordinated transcriptomic changes and spatial reorganization. Inferring the dynamics of biological processes from single -cell sequencing data requires tailored computational approaches known as trajectory inference methods8. Monocle initiated the field of trajectory inference by ordering cells along a pseudotime axis based on their transcriptomic similarities and analyzing gene expression trends along pseudotime 9. While pseudotime represents the progression along a differentiation process, pseudotime -based

do not provide a model for the underlying transcriptomic changes, and thus cannot predict a cell’s future transcriptomic state 10. RNA velocity has thus been proposed to predict changes in gene expression based on splicing dynamics11. However, velocity-based methods rely on simple kinetic models that can misinterpret cell dynamics, for instance in the case of transient boosts in transcription12. Multiple methods based on Optimal Transport (OT) have been developed for cases when several time points are available along differentiation. Waddington OT infers trajectories by computing probabilistic cell -cell transitions between adjacent time points 13. However, it delivers neither a notion of pseudotime nor a notion of velocity. Another class of OT -based

proposes a continuous model of population dynamics by training neural networks representing a generalized notion of velocity14,15. However, these methods do not order cells along a pseudotime axis. A promising OT -based framework for trajectory inference consists of learning a potential function governing a causal model of differentiation16–18. Framing cellular differentiation as the minimization of a potential function is rooted in systems biology and formalizes Waddington's idea of epigenetic landscape19,20. Furthermore, the potential function is a natural alternative to pseudotime, and its gradient yields a rigorous notion of velocity. OT-based approaches for trajectory inference using spatial transcriptomics through time have recently been developed 21–23. For instance, Moscot computes cell -cell transitions between adjacent time points using an extension of OT called Fused Gromov-Wasserstein (FGW), and stVCR learns a spatial velocity along with a gene expression velocity 22–24. However, these

do not provide an intuitive ordering of cells along differentiation, which limits their ability to discover the mechanisms driving dynamic biological processes. In addition, Moscot cannot predict the evolution of cells at future time points. Here, we propose STORIES, a novel trajectory inference method capable of learning a causal model of cellular differentiation from spatial transcriptomics through time using FGW. Unlike Moscot and stVCR, STORIES learns a potential function that defines each cell's stage of differentiation. In addition, unlike Moscot, STORIES allows one to predict the evolution of cells at future time points. Indeed, STORIES uses FGW as a machine learning loss to learn a continuous model of differentiation, while Moscot uses FGW to connect adjacent time points. We benchmarked our approach on three large -scale spatiotemporal Stereo -seq atlases, covering mouse development, zebrafish development, and axolotl regeneration 6,6,7. Furthermore, we used STORIES for the in -depth analysis of cellular trajectories in axolotl neural regeneration and mouse gliogenesis. We recover gene trends for known markers, such as Nptx1 in Nptx+ excitatory neuron regeneration and Aldh1l1 in gliogenesis. In addition, STORIES uncovers other possible driver genes and transcriptional regulators of cellular differentiation in these contexts, which may be of interest for further biological investigation. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint Finally, we provide STORIES as an open -source and user -friendly Python package (github.com/cantinilab/stories). It is based on the Scverse ecosystem, making it easy to interface STORIES with existing tools for single-cell analysis such as Scanpy and CellRank25– 27. In addition, STORIES benefits from the JAX ecosystem for deep learning and OT computation, enabling the fast handling of large datasets28,29.

STORIES: a new single-cell trajectory inference method for spatial transcriptomics profiled through time We developed SpatioTemporal Omics eneRgIES (STORIES), a novel tool for single -cell trajectory inference using omics data profiled through spatial and temporal dimensions (github.com/cantinilab/stories). STORIES allows studying dynamic biological processes in their spatial context by identifying cell fates, gene trends, and candidate transcriptional regulators (see Fig. 1). STORIES is based on the Optimal Transport (OT), a mathematical framework that enables the geometrically meaningful comparison of distributions, using various flavors of the Wasserstein distance30,31. OT also provides a valuable model for population dynamics: the so- called Wasserstein gradient flows were popularized by Jordan, Kinderlehrer, and Otto for their connection with the Fokker-Planck equation and were recently used for trajectory inference in single-cell transcriptomics16–18,32. However, existing methods for trajectory inference based on Wasserstein gradient flows are not equipped to deal with spatially resolved omics data. STORIES introduces key methodological innovations that allow one to address the specific challenges of including spatial information. As an input, STORIES takes slices of spatial transcriptomics profiled at several time points. For instance, Fig. 1A displays sections of axolotl brains profiled at different stages during regeneration. STORIES then learns the parameters of a neural network , which assigns a differentiation potential to each cell according to its gene expression profile (see Fig. 1B). The function formalizes the Waddington epigenetic landscape, where undifferentiated cells have a high potential and, as they differentiate, move towards low -potential transcriptomic states, which correspond to mature cell types20. The transition to these low-potential attractor states defines a causal model of cellular dynamics capable of predicting future gene expression patterns and suggesting potential driver genes and transcriptional regulators (see Fig. 1C). STORIES's potential-based approach provides two interpretable and biologically meaningful outputs: (i) the potential naturally orders cells along a differentiation process (ii) the vector gives the direction of the evolution of gene expression. On the contrary, pseudotime-based methods9,33,34 focus on the first aspect, and velocity -based11,35 methods focus on the second. Crucially, STORIES also innovates compared to state -of-the-art potential-based methods16–18 by enabling the use of spatial coordinates. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint Briefly, STORIES trains the neural network by predicting a distribution of gene expression profiles for each time point where . These distributions are then compared to the ground-truth distributions , and the parameters are updated to improve the predictions. Unlike existing potential -based methods, STORIES allows one to take into account the spatial coordinates of cells when comparing the distributions of gene expression. Formally, let denote the empirical distribution of cells at time , characterized by their gene expression profile , spatial coordinates and weight where . Similarly, let us denote the predictions of STORIES at time . Unlike the gene expression profiles , the spatial coordinates are not directly comparable because the slices are not necessarily aligned between time points. In other words, the spatial coordinates are defined up to an isometry (e.g., a rotation or translation). Existing potential-based methods train the neural network using a linear OT objective, which is sensitive to isometries. Our approach instead uses a recently developed quadratic extension of OT called Fused Gromov-Wasserstein (FGW), which renders the model invariant to spatial isometries24. The FGW distance, defined below and explained more thoroughly in the Methods section, allows one to compare the distributions and directly on gene expression profiles, and up to an isometry on spatial coordinates. (1) FGW seeks a matrix mapping cells from to such that minimizes the sum of three terms: (i) the linear term compares the gene expression coordinates (ii) the quadratic term compares pairwise distances and , which are not affected by translating or rotating the tissue (iii) an entropic regularization term . The parameter denotes the relative weight of spatial information. Our proposed objective function evaluates the predictions across all time points using a debiased version of the FGW distance denoted (see Methods): (2) For , Equation 2 corresponds to a model relying purely on linear OT and which does not leverage spatial information, as proposed in the state-of-the-art16–18. In the following, we refer to this as the linear method. Existing methods 16–18 propose different strategies to make the predictions . They vary in terms of teacher -forcing, number of steps between and , and whether steps are implicit or explicit (see "Discretization" in Methods). The linear

presented here ( ) incorporates the best-performing choices for our experiments (see "Discretization" in Methods and Supplementary Figure 1). The linear method thus offers an unbiased way to assess the performance of STORIES compared to the state-of-the-art, as explored in the next section. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint STORIES is implemented as an open-source Python package seamlessly integrated into the classical Python single -cell analysis pipeline ( github.com/cantinilab/stories). Users can thus take advantage of scverse tools like Scanpy, Squidpy, and CellRank for preprocessing and downstream analysis 26,27,36. In addition, STORIES provides a user -friendly visualization of driver genes and enriched transcription factors, thus helping biological interpretability. In the following sections, we extensively benchmark STORIES against the state -of-the-art16– 18 using large-scale spatiotemporal atlases. To this end, we compare the linear method to our proposed model across different values of the parameter . STORIES’s predictions are more biologically relevant than the state-of-the-art We assessed the effectiveness of STORIES in predicting cell states over time across three Stereo-seq spatiotemporal atlases: a mouse development atlas, a zebrafish development atlas, and an axolotl brain regeneration atlas 4,6,7. Details on data processing are provided in the Methods section. From each atlas, we created three sets: a training set, an early test set, and a late test set (see Fig. 2A). The test sets are composed of two time points, and the goal is to use the first time point to predict the second time point’s gene expression. The late test set is particularly challenging because its second slice comes from an entirely new time point, which may contain cell states not seen during training. For example, in the zebrafish atlas, fast muscle cells only appear at 24 hours post-fertilization (i.e. hpf), whereas the training set includes slices only up to 18 hpf. Two scores, defined in the Methods section, have been used for the benchmark. First, the gene expression prediction score evaluates how closely the prediction approaches the real gene expression at the next time point. Second, the spatial coherence score e valuates whether the predictions are consistent with space. For instance, predicting lung cells from liver cells would yield a bad spatial coherence score. We compared STORIES to the linear method described in the previous section, which incorporates the best-performing aspects of state-of-the-art methods16–18 for our experiments (see "Discretization" in Methods and Supplementary Fig. 1). Since the linear method is formally equivalent to STORIES with , it offers an unbiased way to assess the usefulness of spatial information compared to the state-of-the-art. To better evaluate the impact of spatial information, we reported results with a varying weight , representing the importance of the spatial term compared to the gene expression term (see Figure 2B). As expected, smaller values of perform better in terms of gene expression prediction score, while larger values perform better in terms of spatial coherence score. A good compromise between gene expression and space is obtained for , which we set as a default value. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint For this choice of , STORIES outperforms the linear method in terms of spatial coherence across all datasets for both early and late test sets (see Fig. 2B). In addition, STORIES outperforms the linear method in terms of gene expression prediction in the mouse atlas's early test set and the zebrafish atlas's early test set. This demonstrates that STORIES's predictions are biologically relevant on two levels: they are accurate in gene expression and coherent in space. The biological relevance of STORIES’s predictions is confirmed in Fig. 2C, which compares the training losses of STORIES and the linear method. Since both methods involve matching predictions with a reference population of cells, we compared their matchings for specific cell types (see Methods). First, in the axolotl atlas, STORIES correctly matches predictions from immature neurons (IMN) with Nptx+ excitatory neurons in the lateral pallium (NptxEX), and predictions from regeneration intermediate progenitor cells (rIPC2) with excitatory neurons in the dorsal pallium (dpEX)7. The linear method, on the contrary, incorrectly matches rIPC2 predictions with microglial cells (MCG) from a different anatomical region (see Fig. 2C). Second, in the zebrafish atlas, STORIES accurately matches predictions from the optic vesicle with cells located around the eye, and predictions from the polster with cells located within the head. In contrast, the linear method incorrectly matches optic vesicle predictions with a broad group of cells across different anatomical regions, and polster predictions with cells from the tail area (see Fig. 2C). Third, in the mouse atlas, STORIES correctly matches predictions from liver and lung cells with their respective organs. The linear method, instead, incorrectly matches lung cell predictions with a broad group of cells across organs (see Fig. 2C), suggesting poor biological coherence of the linear method. STORIES's superior performance in achieving biologically coherent and accurate gene expression predictions demonstrates the significant benefit of considering spatial information when learning a gradient flow model on spatial transcriptomics data. STORIES identifies trajectories and potential drivers of neuron regeneration in axolotls To further assess the potential of STORIES for trajectory inference in spatial transcriptomics through time, we first focused on axolotl brain regeneration. We trained STORIES as described in Methods on the subset of cells described in the original publication as involved in neuron regeneration: Wnt+ and reactive ependymoglial cells (wntEGC and reaEGC), regeneration intermediate progenitor cells (rIPC1 and rIP C2), immature neurons (IMN), Nptx+ lateral pallium excitatory neurons (nptxEX), dorsal pallium excitatory neurons (dpEX), and medial pallium excitatory neurons (mpEX)7. As shown in Fig. 3A, STORIES learns an energy landscape consistent with the original publication. Indeed, the potential assigns a high potential to progenitor states (wntEGC and reaEGC), a medium potential to intermediary states (rIPC1, rIPC2, and IMN), and a low potential to mature states (nptxEX, dpEX, and mpEX). .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint We computed cell-cell transitions by applying CellRank on the gradient of the trained potential (see Methods), as visualized in Fig. 3B. These transitions highlight that STORIES not only detects the correct stage of differentiation, but also recovers the t hree trajectories described in the original publication: wntEGC -mpEX, reaEGC -rIPC2-dpEX, and reaEGC -rIPC1-IMN- nptxEX7. Importantly, the authors identified these three trajectories by applying Moncole separately on three spatial regions and specifying EGCs as the starting point in each case. In contrast, STORIES achieves the same results without the need to isolate specif ic spatial regions and specify the starting point of the trajectory. Indeed, STORIES leverages spatial information to process all regions simultaneously and leverages temporal information to infer progenitor states from the data. We then narrowed further into the reaEGC -rIPC1-IMN-nptxEX trajectory, which the original publication studies in most detail7. First, we sought to confirm expected gene trends along this trajectory. The original study suggests Vim, which encodes a critical cytoskeletal protein, as a marker of reaEGC cells and Nptx1, which is involved in synaptic plasticity, as a marker of NptxEX. Accordingly, STORIES recovered a clear decreasing trend for Vim expression along differentiation and a clear increasing trend for Nptx1 expression (see Fig. 3C). Next, we performed unsupervised discovery of gene trends by fitting a spline regression model along the previously mentioned trajectory (see Methods). Fig. 3D reports the best candidate driver genes across differentiation stages. Interestingly, the early s tages of differentiation coincide with high expression of Hes5, which is known to maintain stemness in the context of neural differentiation 37, and Cdc25b, a cell -cycle regulator key to neuron production 38,39. Conversely, late stages of differentiation coincide with high expression of the microtubule - associated protein gene Map1a, crucial to neural development and regeneration 40, and L1cam, shown to promote axonal regeneration41. STORIES also outputs additional genes that represent possible drivers of neuron regeneration and would require further biological investigation. For instance, STORIES uncovered a trend for late expression of the scarcely studied Nsg2, which is thought to be involved in synaptic function and, like Nptx1, interacts with AMPA receptors42. Finally, our analysis revealed possible transcriptional regulators of the differentiation process (see Fig. 3F) by testing transcription factor (TF) enrichment using the curated literature-based TRRUST database (see Methods). The most significantly enriched TF, CTNNB1, encodes β- catenin, which has been described as an essential regulator in neuron regeneration in mouse models and in limb regeneration in axolotl43–45. Other top TFs include SP1 and MYC, described in the context of neuron regeneration and computationally retrieved in axolotl limb regeneration46–49. Additionally, we identify SOX6, MYCN, and REST, which are not widely studied in the context of regeneration but are known regulators in development 50–52. Interestingly, a recent study predicted REST as a regulator of neuron regeneration and validated this role in a mouse model53. STORIES thus learned a Waddington landscape that captures cell fate trajectories, continuous changes of gene expression, and the regulatory landscape underlying neuron regeneration in axolotls. Our model allowed for the unbiased discovery of potential drivers and mechanisms, possibly relevant for further biological investigations. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint STORIES identifies trajectories and potential drivers of gliogenesis in the developing mouse dorsal midbrain We then sought to highlight STORIES's potential in trajectory inference by studying mouse dorsal midbrain development. We trained STORIES as described in Methods on the subset of cells described in the original article as exhibiting a branching trajectory: radial glial cells (RGC) differentiating into either neuroblasts (NeuB) or glioblasts (GlioB)4. As visualized in Fig. 4A, STORIES learns an energy landscape consistent with the original publication. Indeed, the potential assigns a high potential to RGC and a low potential to the more differentiated NeuB and GlioB. We computed cell-cell transitions by applying CellRank on the gradient of the trained potential (see Methods), as visualized in Fig. 4B. These transitions highlight that STORIES not only detects the correct stage of differentiation but also recovers the expected branching from RGC to glial and neural cell fates4. Importantly, the original publication identified this branching using Monocle 3, which required manually setting RGC as the trajectories' starting point. On the contrary, STORIES achieves the same results without manual input by leveraging temporal information to infer the starting point from the data. Glial cells outnumber neurons in the brain, but their development has been studied less extensively54. Moreover, understanding gliogenesis is of critical therapeutic importance because of its parallels with glioma, the most common and deadliest form of brain cancer 55. Thus, we focused further on the RGC -GlioB trajectory. We first sought to confirm expected gene trends along this trajectory. The original study identifies Mki67, a proliferation marker, as highly expressed in RGC, and Aldh1l1, an astrocyte marker, as highly expressed in GlioB 4. Accordingly, STORIES recovered a decreasing trend for Mki67 expression along differentiation and an increasing trend for Aldh1l1 expression (see Fig. 4C). Next, we performed unsupervised discovery of gene trends by fitting a spline regression model along the previously mentioned trajectory (see Methods). Fig. 4D reports the best candidate driver genes across differentiation stages. The early stages of differ entiation coincide with a high expression of cell cycle genes Gmnn, Rrm2, and Hmgb256. Additionally, we observed a high expression of the alpha -tubulin gene Tuba1b in the early stages of differentiation, as previously described in the developing brain 57. Conversely, the late stages of differentiation coincide with the high expression of the glutamine synthetase gene Glul, a key astrocyte marker58,59. STORIES also outputs additional genes that represent possible drivers of gliogenesis and would require further biological investigation. For instance, Glis3 displays an increasing trend along gliogenesis (see Fig. 4E) but is little studied in this context. However, Glis3 was recently suggested as a therapeutic target to suppress proliferation in glioma60. Finally, our analysis revealed candidate transcriptional regulators of the differentiation process (see Fig. 4F) by testing transcription factor (TF) enrichment using the curated literature-based TRRUST database (see Methods). Among the most enriched TFs, SOX4 and NOTCH2 have been studied in gliogenesis 61,62. Additionally, STORIES recovers MYC, MYCN, and MAX, which have been studied in the context of glioma63. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint This second experiment in the context of gliogenesis confirms STORIES's ability to learn a Waddington landscape that captures cell fate trajectories, continuous gene expression changes, and regulatory processes. In doing so, we highlight known and potentia l mechanisms that motivate further biological investigations.

Recent technological advances in spatial transcriptomics have enabled the tracking of gene expression at single-cell resolution in the spatial context of the tissue. Large datasets of spatial transcriptomics profiled through time give a unique opportunity to understand dynamic biological processes such as development and disease onset. However, their analysis requires trajectory inference tools tailored to the specific challenges of spatial data. In this article, we proposed SpatioTemporal Omics eneRgIES (STORIES), a novel computational framework for trajectory inference from spatial transcriptomics profiled at several time points. STORIES enables a rich and spatially informed analysis of differentiation trajectories. To evaluate STORIES’s performance, we benchmarked it against the state -of- the-art in three large Stereo-seq datasets and highlighted the advantage of considering spatial information in trajectory inference from time -course single-cell data. We further showcased STORIES’s abilities in two concrete settings: axolotl neuron regeneration and mouse gliogenesis. STORIES offers a model of population dynamics tailored for single -cell resolution spatial transcriptomics technologies like Stereo -seq or Visium HD. Given the fast -paced developments in spatial transcriptomics1, the number of spatiotemporal atlases at single -cell resolution can be expected to increase steadily. At the same time, STORIES could be applied to low-resolution data (e.g. 10x Visium), which have a spot size larger than the typical cell, using deconvolution techniques64,65. In addition, STORIES could be adapted to imaging-based technologies like MERFISH which offer high resolution but can only detect a limited panel of genes66. STORIES provides an interpretable model of differentiation relying on a potential energy. Previous work shows that such potential landscapes arise naturally from simple gene regulatory networks (GRNs)19. However, potential energies cannot model complex GRNs, cell- cell communication, or oscillations within a cell state 67. Extensions to more complex energy functionals could thus lead to further insights into biological processes such as development, immune response, and the onset of complex disease. For instance, recent work in gradient flow learning explored interaction energies68, which could represent cell-cell communication in biology. The major novelty of STORIES is its ability to learn a spatially -informed potential. This methodological development is critical because dynamic processes such as development involve coordinated expression changes and tissue reorganization 69. However, the learned potential operates only on gene expression, so it does not allow the prediction of future positions of cells. Including a spatial component in the energy function may provide a more comprehensive view of biological processes by predi cting cell migration. Relating this to .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint existing models for morphogenesis, such as Alan Turing's reaction-diffusion model70–72, is an exciting avenue for further research.

Data collection Zebrafish For the training and validation sets, we used the five first time points showcased in Figure 2A of the original publication6: 3.3hpf slice 1, 5.25hpf slice 10, 10hpf slice 11, 12hpf slice 8, 18hpf slice 8. For the test set, we used: ● 10hpf slice 17 and 12hpf slice 5 to evaluate prediction within the time range seen during training. These two slices were studied together in Figure S4A of the original publication6. ● 18hpf slice 11 and 24hpf slice 4 to evaluate prediction outside the time range seen during training. These two slices were studied together in Figure S4E of the original publication6. Altogether, this represented 17,920 cells after preprocessing. Mouse For the training and validation sets, we used the 7 first time points showcased in Figure 3A of the original publication4: E9.5 E1S1, E10.5 E1S1, E11.5 E1S1, E12.5 E1S1, E13.5 E1S1, E14.5 E1S1, E15.5 E1S1. For the test set, we used: ● E13.5 E1S2 and E14.5 E1S2 to evaluate prediction within the time range seen during training. These two slices were studied in Figure S2C of the original publication4. ● E15.5 E1S2 and E16.5 E1S1 to evaluate prediction outside of the time range seen during training. The first slice was studied in Figure S2C and the second in Figure 3A of the original publication4. Altogether, this represented 794,063 cells after preprocessing. Dorsal midbrain We retained one slice per time point. We used the 3 slices showcased in Figure S7A of the original publication4: E12.5 (E1S3), E14.5 (E1S3), E16.5 (E1S3). As in the original publication, we subset the analysis to the RGC, NeuB, and GlioB cell types. Altogether, this represented 4,581 cells after preprocessing. Note that the "Mouse" dataset and the "Dorsal midbrain" dataset originate from the same experiments but have different resolutions (bin 50 vs image-based segmentation), so the same neural network weights cannot be used in both cases. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint Axolotl Benchmark. For the training and validation sets, we used the five first time points showcased in Figure 3B of the original publication7: 2DPI (rep1), 5DPI (rep1), 10DPI (rep1), 15DPI (rep4), 20DPI (rep2). We manually removed spatial outliers in 10DPI (rep1), 15DPI (rep4), and 20DPI (rep2). For the test set, we used: ● 10DPI (rep2) and 15DPI (rep1) to evaluate prediction within the time range seen during training. These two slices were showcased in Figure S9 of the original publication7. ● 20DPI (rep3) and 30DPI (rep2) to evaluate prediction outside of the time range seen during training. These two slices were showcased in Figure S9 of the original publication7. As in the original publication, we restricted the analysis to the dorsal part of the injured hemisphere. Altogether, this represented 22,083 cells. In-depth analysis. For the analysis in Section 3 of Results, we used slices 2DPI (rep1), 5DPI (rep1), 10DPI (rep1), 15DPI (rep4), 20DPI (rep2), and 30DPI (rep2). As before, we manually removed spatial outliers and restricted the analysis to the dorsal part of the injured hemisphere. As in the original publication, we subset the data to the following cell types: nptxEX, reaEGC, wntEGC, dpEX, mpEX, IMN, rIPC1, rIPC2. This represented 5,904 cells. Preprocessing For all datasets, we performed the following preprocessing steps. Cell and gene quality control. Using Scanpy's sc.pp.filter_cells, we removed cells with less than 200 expressed genes. Then, we removed the top 0.1% of cells with the most expressed genes. Finally, we removed genes expressed in less than 3 cells using Scanpy's sc.pp.filter_genes. Normalization and highly variable gene selection. We applied Scanpy's "Pearson residuals normalization" and selected 10,000 highly variable genes. Scanpy computed highly variable genes for each batch and merged them to avoid selecting batch-specific genes. Dimensionality reduction. Using Scanpy's sc.tl.pca, we applied Principal Component Analysis (PCA) to reduce the data to 50 dimensions. Section 3 and Section 4 of Results used a subset of relevant cell types. This was done after PCA but before batch correction. Batch correction. We applied Harmony on the PCA components to correct the batch effects, using Scanpy's sc.external.pp.harmony_integrate, a wrapper around harmonypy73. Visualization. Using Scanpy's sc.tl.umap, we applied UMAP to project the batch-corrected data into two dimensions. In Section 3 and Section 4 of Results we applied Isomap instead of UMAP. Indeed, we found that visually, Isomap respected cell type transitions better than UMAP. We used scikit-learn's sklearn.manifold.Isomap. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint Wasserstein gradient flow learning with a quadratic objective Notations Let us consider a time point and a discrete distribution of cells. We denote where the vector represents the gene expression of the -th cell, and the weights are such that . In the following, for a function , we define the pushforward measure as . Wasserstein gradient flow Similarly to previous works16–18, we model the evolution of as a Wasserstein gradient flow for a potential energy. For one single cell , the Euclidean gradient flow is a that verifies , the continuous counterpart of gradient descent. Wasserstein gradient flows extend this to the space of measures74,75. We say that is a Wasserstein gradient flow for the energy if its density verifies the continuity equation As in previous works16–18, we do not know the potential a priori and aim instead to learn a neural network from snapshots for . See "neural network architecture" for details about . Discretization Our approach boils down to learning such that for given parameters and an initial population , the predicted populations are close to the observed snapshots . To make these predictions, existing potential-based methods16–18 differ in three main aspects: number of steps, teacher-forcing, and discretization scheme. As detailed below, we select the best-performing choices for these three aspects, as measured by validation loss in the Zebrafish atlas. The linear method ( ) thus differs from existing works by combining their best-performing aspects. Number of steps. Hashimoto et al. and Yeo et al.16,17 make intermediary predictions between and , i.e. . Bunne et al. perform a single step instead, i.e. 18. In our experiments, multiple steps did not improve results (see Supplementary Figure 1A) so we chose the computationally less expensive single-step method. Teacher-forcing. Hashimoto et al. and Yeo et al.16,17 predict from . Bunne et al. introduce teacher-forcing, i.e. predicting from 18. In our experiments, teacher-forcing improved results (see Supplementary Figure 1B), so we used it throughout this work. Discretization scheme. To predict from an earlier population , we used the forward Euler discretization scheme as Hashimoto et al. and Yeo et al.16,17. In our discrete setting, this corresponds to for each cell . Bunne et al. propose using a backward Euler scheme to improve stability for large 18. However we .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint did not find this to improve results in our experiments (see Supplementary Figure 1C), so we used the computationally less expensive forward method. Pairwise information Our setting differs from previous works16–18 in that for each cell we have access to spatial coordinates . Spatial coordinates are defined up to an isometry. For instance, a slice of an embryo may be rotated without changing the problem. Consequently, one cannot simply concatenate and to leverage the spatial coordinates. The next paragraph details how we used the coordinates to inform the problem, while still defining only on gene expression. Learning the potential For each , we compared the prediction to the reference snapshot . Linear model. Let us consider two discrete probability distributions and . To compare and , previous works16–18 use the Sinkhorn divergence76, defined as where is the entropy-regularized Optimal Transport77 (entropy-regularized OT), defined as Here, and is the Shannon entropy. is a debiased version of , such that . Quadratic model. The Sinkhorn divergence between and only compares distributions of gene expression. Instead, we propose a debiased Fused Gromov-Wasserstein (debiased FGW) loss to enforce the spatial coherence of the predictions. Let us consider and . Gromov-Wasserstein (GW) is a quadratic extension of OT well suited to compare measures defined up to an isometry78. A debiased version of GW has been used to learn a Generative Adversarial Network (GAN)79. Fused Gromov-Wassertein (FGW) combines a linear and a quadratic OT term80. In our setting, it is natural to use the linear term for gene expression and the quadratic term for spatial coordinates: Analogously to Bunne et al., we introduce a debiased FGW to ensure the loss vanishes for an exact match79. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint A weight corresponds to the debiased entropy-regularized GW81. As , we recover the Sinkhorn divergence . Final loss function. For , the full objective reads where with . In other words, gene expression is predicted with teacher-forcing and a single forward step (see "Discretization"), and the spatial coordinates remain unchanged. This loss is optimized using minibatches, as done by Yeo et al. and Bunne et al.17,18 A theoretical study of minibatch Optimal Transport for debiased linear and quadratic OT is provided by Fatras et al.82 Choice of the quadratic weight To investigate the effect of the relative weight of the linear term in FGW, we reported results for . Choosing would not allow learning the potential because it would ignore gene expression. As expected, with a low weight ( ), STORIES behaves as the linear method (see Section 2 of Results). To compare our approach with the linear model proposed in previous works16–18, we also trained the model with a Sinkhorn divergence, i.e. . The benchmark in Section 2 of Results suggests good performances for values of of the order of , with performing best. The value of can be adjusted by the user depending on the dataset. In Section 3 and Section 4 of Results we set a value of . Computational Optimal Transport OTT solvers We use the OTT package to solve OT problems in a fast, GPU-enabled, and differentiable manner29. In particular, we rely on the Sinkhorn and GromovWasserstein solvers. We set the entropic regularization . Linear term: gene expression The linear OT terms are defined on gene expression space, for which we chose the first components of the Harmony-aligned Principal Component Analysis. We chose components as Bunne et al.18 Before training the neural network, we normalized the points as to make the linear and quadratic terms comparable. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint Quadratic term: Spatial coordinates Before training the neural network and for each slice separately, we centered the spatial coordinates and scaled them to unit variance to make linear and quadratic terms comparable. Neural network architecture We implemented a Multi-Layer Perceptron (MLP) with two hidden layers of dimension 128 and GeLU activations using Flax83. This is a similar architecture to previous works16–18. The linear output layer has no bias since it would not influence the values of . Likewise, a soft activation like GeLU is preferable to the classical ReLU because we manipulate . Indeed, the derivative of ReLU is simply a unit step function, which is discontinuous and not very expressive. Neural network training Data loading For each time point described in “Data Collection", 75% of the cells were used as training samples and 25% as validation. At each training or validation iteration, a batch containing 1,000 cells per time point was sampled uniformly without replacement. In development, the early time points contain fewer cells than the later time points. If less than 1,000 cells were available for a time point, we used all available cells. To reflect the train/validation split, one in four iterations performs a validation step. Optimizer We used Optax's implementation of the AdamW optimizer, with parameters b1=0.9, b2=0.999, eps=1e-8, and weight_decay=1e-428,84. We set the learning rate using Optax's cosine scheduler, with an initial value of 1e-2 and 10000 decay steps. To ensure convergence, when performing 10 steps, we set the learning rate to 1e-3. Similarly, when performing an implicit step, we set the learning rate to 1e-4. Early stopping We set the maximum number of iterations to 15,000 but stopped the training when the validation loss had not improved in 150 iterations. We kept the weights associated with the lowest validation loss. Seeds We ran every experiment with 10 random seeds: 17158, 20181, 12409, 5360, 21712, 21781, 24802, 13630, 9668, and 651. The random seed reproducibly determines the train/validation split and weight initialization. For the plots in Figure 2C, and the analysis in Section 3 and Section 4 of Results, the experiments correspond to the randomly chosen seed 20181. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint Computational growth rate Weight of marginals When comparing a prediction to the reference snapshot , Yeo et al. proposes setting the weights of the prediction proportionally to a computationally derived growth rate17. The motivation is that cells with a larger growth rate should be matched to more descendants. This idea was introduced in Waddington OT and recently reimplemented in MOSCOT13,22. In the next paragraphs, we follow MOSCOT's implementation. Proliferation and apoptosis Computing the growth rate relies on cell-wise proliferation and apoptosis scores. We used Scanpy's sc.tl.score_genes with lists of genes collected from the literature. The gene lists are described in more detail at the end of the section. We calculated gene scores on raw counts, after quality filtering of cells and genes. Calculating the growth rate For a given cell , let us call the proliferation score and the apoptosis score. We then define the birth rate and the death rate as We used MOSCOT's default parameters , , , , , . Finally, we defined the cell's growth rate as where is the time difference between populations. We obtained by normalizing the growth rate In this equation, plays the role of the softmax's inverse temperature. The histogram will thus be sharper for large values of . Most slices in our experiments were evenly sampled, so we set a fixed , which yielded sharp enough weight differences between cell types. In the benchmark Previous works tested gene sets to compute the growth rate in the case of mice and humans, but not of zebrafish and axolotl13,17,22. We thus kept uniform marginals for the benchmark in Section 2 of Results. In Section 3 and Section 4 of Results, we analyzed a restricted number of cell types, where we could validate the biological coherence of the computed growth rate and of the learned potential (see details in paragraphs below). Dorsal midbrain For Section 3 of Results, we used murine proliferation and apoptosis gene sets from MOSCOT. Proliferation genes come from https://doi.org/10.1038/nature20123 and apoptosis .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint genes from https://www.gsea-msigdb.org/gsea/msigdb/cards/HALLMARK_P53_PATHWAY. See Supplementary Text 1 and Supplementary Text 2 for the gene names. Axolotl neuron regeneration For Section 4 of Results, we used an NSC axolotl gene set described in the original publication https://www.science.org/doi/10.1126/science.abp9444 to represent proliferation and a human apoptosis gene set from MOSCOT, originally from gsea-msigdb's HALLMARK_APOPTOSIS. See Supplementary Text 3 and Supplementary Text 4 for the gene names. Evaluation We solve the FGW problem with and and report the terms and separately. Since quantifies the error in terms of gene expression, we call this quantity the gene expression prediction score. Similarly, since quantifies the error in terms of spatial coordinates, we call this quantity the spatial coherence score. Figure 3C compares the transport plans involved in the linear method and STORIES with . Formally, let us consider the indicator vector where if the -th cell corresponds to a given cell type, and otherwise. The transport plan between the prediction and the ground truth is applied to the indicator, yielding a vector representing the mass transported from towards each cell in the second time point. Trajectory inference Gene imputation In our analysis, the gene expression trends (Figure 3 and Figure 4) would be negatively affected by the sparsity of gene expression. Cellrank demonstrated good performances in identifying gene expression trends with MAGIC27,85. We thus applied MAGIC gene imputation after all other preprocessing steps. We computed the exponentiated Markov transition matrix on the Harmony-aligned PCA space instead of the original PCA. We did not use the imputed signal for tasks other than gene expression trends. Potential visualization The neural potential is a functional defined on the 20-dimensional space of Harmony- aligned principal components. To visualize the potential as a Waddington-like landscape defined on 2 dimensions, we proceed similarly to Qin et al.86. ● First, we compute the potential associated with each cell . ● Then, we use Scipy's RBF interpolation and the 2-dimensional Isomap coordinates of the cells to define a potential on a 2-dimensional grid87. ● The cells are projected on the surface using the interpolator. ● Finally, the maximum value is thresholded. We rendered the resulting surface and point cloud using Blender's Python API. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint Cell-cell transition matrix We used CellRank's VelocityKernel with a velocity for a trained . Based on this kernel, we computed a cell-cell transition matrix, and the trajectory plots Figure 3B, Figure 4B. Gene expression trends We fit the MAGIC-imputed gene expression as a function of the learned potential, using Scipy's Spline Regression87. To plot the gene expression cascades in Figure 3D, Figure 4D, we order genes by the value of potential for which the maximum regressed expression value is achieved. Genes are then split into equally sized groups illustrating different stages of differentiation (10 groups for Figure 3D and 2 groups for Figure 4D). Finally, regressed values for the genes with the best regression scores in each group are displayed (3 per group in Figure 3D and 15 per group in Figure 4D). Transcription factor enrichment We perform TF-target enrichment based on the TRRUST dataset, which contains “Activation", "Repression", and "Unknown" links based on a curated literature search88. For each TF in the database, we perform a Wilcoxon rank-sum test comparing the list of regression scores of its target genes, and the list of regression scores of the other genes. Figure 3B and Figure 4B display the TFs ranked by -value. Data availability ● We retrieved the mouse Stereo-seq atlas from Chen et al.4, available at https://db.cngb.org/stomics/mosta/. ● We retrieved the zebrafish Stereo-seq atlas from Liu et al.6, available at https://db.cngb.org/stomics/zesta/. ● We retrieved the axolotl Stereo-seq atlas from Wei et al.7, available at https://db.cngb.org/stomics/artista/. ● We retrieved mouse proliferation genes from Tirosh et al.89 (see Supplementary Text 1) and apoptosis genes from gsea-msigdb's HALLMARK_P53_PATHWAY (see Supplementary Text 2). ● We retrieved an NSC axolotl gene set from Wei et al.7 (see Supplementary Text 3) and a human apoptosis gene set from gsea-msigdb's HALLMARK_APOPTOSIS (see Supplementary Text 4). Code availability The Python package for STORIES is hosted at https://github.com/cantinilab/stories. It can be installed easily by running “pip install stories-jax”. Code to reproduce the experiments and figures is available at https://github.com/cantinilab/stories_reproducibility/. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint Figure captions Figure 1. Overview of STORIES. (A) STORIES takes as an input spatial transcriptomics through time. (B) STORIES learns the parameters of a neural network representing the differentiation potential of a cell based on its transcriptomic profile. The objective function is based on Fused Gromov-Wasserstein, which leverages both the transcriptomic profile and the spatial coordinates. (C) The gradient of the function delivers a velocity that can be used to perform trajectory inference. The potential itself is a natural alternative to pseudotime and allows the study of gene trends along differentiation. Finally, STORIES can highlight possible transcription factors regulating differentiation. Figure 2. Benchmark of STORIES on three large datasets. (A) Visual representation of the three datasets in our benchmark. From left to right, an axolotl brain regeneration dataset, zebrafish development, and mouse development. Slices in each dataset are split into a train set (blue), an early test set (orange), an d a late test set (green). (B) Gene expression prediction score (top) and spatial coherence score (bottom) in early test sets (orange) and late test sets (green) across the three datasets. Scores a re reported for ten seeds, across seven values of quadratic weight parameter , including the linear method ( , light orange/green). (C) Visual representation of the optimal transport matching involved in the loss of the linear

(top) and STORIES for (bottom), for the three datasets. In each dataset, the left slice displays two cell types and the right slice displays the cells they are matched with at the following time point. Figure 3. Trajectory inference with STORIES in axolotl neuron regeneration. (A) 3-D representation of the potential landscape learned with STORIES. The x and y axes are Isomap coordinates, and the z-axis is an interpolation of the potential. Colors represent cell types involved in the regeneration process. (B) Visual representation of cell -cell transitions computed using CellRank from STORIES’s velocity vectors. (C) Smoothed gene expression for Vim and Nptx1 along the potential computed by STORIES . The blue line is a spline regression of expression from potential. (D) Normalized gene expression regressed using a spline model along the potential computed by STORIES. Genes are ordered by the potential for which they achieve maximum expression (E) Smoothed gene expression for Hes5 and Nsg2 along the potential comput ed by STORIES . The blue line is a spline regression of expression from potential. (F) Enrichment scores of Transcription Factors targeting candidate driver genes. Figure 4. Trajectory inference with STORIES in mouse gliogenesis. (A) 3-D representation of the potential landscape learned with STORIES. The x and y axes are Isomap coordinates, and the z-axis is an interpolation of the potential. Colors represent radial glial cells which differentiate into either neuroblasts or glioblasts. (B) Visual representation of cell- cell transitions computed using CellRank from STORIES’s velocity vectors. (C) Smoothed gene expression for Mki67 and Aldh1l1 along the potential computed by STORIES. The blue line is a spline regression of expression from potential. (D) Normalized gene expression regressed using a spline model along the potential computed by STORIES. Genes are ordered by the potential for which they achieve maximum expression (E) Smoothed gene expression for Tuba1b and Glis3 along the potential computed by STORIES . The blue line is a spline .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint regression of expression from potential. (F) Enrichment scores of transcription factors targeting candidate driver genes.

1. Palla, G., Fischer, D. S., Regev, A. & Theis, F. J. Spatial components of molecular tissue biology. Nat. Biotechnol. 40, 308–318 (2022). 2. Hunter, M. V., Moncada, R., Weiss, J. M., Yanai, I. & White, R. M. Spatially resolved transcriptomics reveals the architecture of the tumor-microenvironment interface. Nat. Commun. 12, 6278 (2021). 3. Chen, W.-T. et al. Spatial transcriptomics and in situ sequencing to study Alzheimer’s disease. Cell 182, 976–991 (2020). 4. Chen, A. et al. Spatiotemporal transcriptomic atlas of mouse organogenesis using DNA nanoball-patterned arrays. Cell 185, 1777–1792 (2022). 5. Vickovic, S. et al. High-definition spatial transcriptomics for in situ tissue profiling. Nat.

16, 987–990 (2019). 6. Liu, C. et al. Spatiotemporal mapping of gene expression landscapes and developmental trajectories during zebrafish embryogenesis. Dev. Cell 57, 1284–1298 (2022). 7. Wei, X. et al. Single-cell Stereo-seq reveals induced progenitor cells involved in axolotl brain regeneration. Science 377, eabp9444 (2022). 8. Saelens, W., Cannoodt, R., Todorov, H. & Saeys, Y. A comparison of single-cell trajectory inference methods. Nat. Biotechnol. 37, 547–554 (2019). 9. Trapnell, C. et al. The dynamics and regulators of cell fate decisions are revealed by pseudotemporal ordering of single cells. Nat. Biotechnol. 32, 381–386 (2014). 10. Tritschler, S. et al. Concepts and limitations for learning developmental trajectories from single cell genomics. Development 146, (2019). 11. La Manno, G. et al. RNA velocity of single cells. Nature 560, 494–498 (2018). 12. Bergen, V., Soldatov, R. A., Kharchenko, P. V. & Theis, F. J. RNA velocity—current challenges and future perspectives. Mol. Syst. Biol. 17, e10282 (2021). .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint 13. Schiebinger, G. et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell 176, 928–943 (2019). 14. Tong, A., Huang, J., Wolf, G., Van Dijk, D. & Krishnaswamy, S. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. in International conference on machine learning 9526–9536 (PMLR, 2020). 15. Huguet, G. et al. Manifold interpolating optimal-transport flows for trajectory inference. Adv. Neural Inf. Process. Syst. 35, 29705–29718 (2022). 16. Hashimoto, T., Gifford, D. & Jaakkola, T. Learning population-level diffusions with generative RNNs. in International Conference on Machine Learning 2417–2426 (PMLR, 2016). 17. Yeo, G. H. T., Saksena, S. D. & Gifford, D. K. Generative modeling of single-cell time series with PRESCIENT enables prediction of cell trajectories with interventions. Nat. Commun. 12, 3222 (2021). 18. Bunne, C., Papaxanthos, L., Krause, A. & Cuturi, M. Proximal optimal transport modeling of population dynamics. in International Conference on Artificial Intelligence and Statistics 6511–6528 (PMLR, 2022). 19. Huang, S., Guo, Y.-P., May, G. & Enver, T. Bifurcation dynamics in lineage- commitment in bipotent progenitor cells. Dev. Biol. 305, 695–713 (2007). 20. Allen, M. Compelled by the diagram: thinking through CH Waddington’s epigenetic landscape. Contemporaneity 4, 119 (2015). 21. Shen, X. et al. Inferring cell trajectories of spatial transcriptomics via optimal transport analysis. bioRxiv 2023–09 (2023). 22. Klein, D. et al. Mapping cells through time and space with moscot. bioRxiv 2023–05 (2023). 23. Peng, Q., Zhou, P. & Li, T. stVCR: Reconstructing spatio-temporal dynamics of cell development using optimal transport. bioRxiv 2024–06 (2024). 24. Vayer, T., Chapel, L., Flamary, R., Tavenard, R. & Courty, N. Optimal transport for structured data with application on graphs. ArXiv Prepr. ArXiv180509114 (2018). .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint 25. Virshup, I. et al. The scverse project provides a computational ecosystem for single- cell omics data analysis. Nat. Biotechnol. 41, 604–606 (2023). 26. Wolf, F. A., Angerer, P. & Theis, F. J. SCANPY: large-scale single-cell gene expression data analysis. Genome Biol. 19, 15 (2018). 27. Weiler, P., Lange, M., Klein, M., Pe’er, D. & Theis, F. Unified fate mapping in multiview single-cell data. (2023) doi:10.1101/2023.07.19.549685. 28. DeepMind et al. The DeepMind JAX Ecosystem. (2020). 29. Cuturi, M. et al. Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein. ArXiv Prepr. ArXiv220112324 (2022). 30. Peyré, G., Cuturi, M., & others. Computational optimal transport: With applications to data science. Found. Trends® Mach. Learn. 11, 355–607 (2019). 31. Villani, C. & others. Optimal Transport: Old and New. vol. 338 (Springer, 2009). 32. Jordan, R., Kinderlehrer, D. & Otto, F. The variational formulation of the Fokker– Planck equation. SIAM J. Math. Anal. 29, 1–17 (1998). 33. Street, K. et al. Slingshot: cell lineage and pseudotime inference for single-cell transcriptomics. BMC Genomics 19, 1–16 (2018). 34. Chen, H. et al. Single-cell trajectories reconstruction, exploration and mapping of omics data with STREAM. Nat. Commun. 10, 1903 (2019). 35. Bergen, V., Lange, M., Peidli, S., Wolf, F. A. & Theis, F. J. Generalizing RNA velocity to transient cell states through dynamical modeling. Nat. Biotechnol. 38, 1408–1414 (2020). 36. Palla, G. et al. Squidpy: a scalable framework for spatial single cell analysis. BioRxiv 2021–02 (2021). 37. Ohtsuka, T., Sakamoto, M., Guillemot, F. & Kageyama, R. Roles of the basic helix- loop-helix genes Hes1 and Hes5 in expansion of neural stem cells of the developing brain. J. Biol. Chem. 276, 30467–30474 (2001). 38. Peco, E. et al. The CDC25B phosphatase shortens the G2 phase of neural progenitors and promotes efficient neuron production. Development 139, 1095–1104 .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint (2012). 39. Bonnet, F. et al. Neurogenic decisions require a cell cycle independent function of the CDC25B phosphatase. Elife 7, e32937 (2018). 40. Nunez, J. & Fischer, I. Microtubule-associated proteins (MAPs) in the peripheral nervous system during development and regeneration. J. Mol. Neurosci. 8, 207–222 (1997). 41. Schäfer, M. K. & Frotscher, M. Role of L1CAM for axon sprouting and branching. Cell Tissue Res. 349, 39–48 (2012). 42. Chander, P., Kennedy, M. J., Winckler, B. & Weick, J. P. Neuron-specific gene 2 (NSG2) encodes an AMPA receptor interacting protein that modulates excitatory neurotransmission. eneuro 6, (2019). 43. Marchetti, B. Wnt/β-catenin signaling pathway governs a full program for dopaminergic neuron survival, neurorescue and regeneration in the MPTP mouse model of Parkinson’s disease. Int. J. Mol. Sci. 19, 3743 (2018). 44. Duraikannu, A., Martinez, J. A., Chandrasekhar, A. & Zochodne, D. W. Expression and manipulation of the APC-β-catenin pathway during peripheral neuron regeneration. Sci. Rep. 8, 13197 (2018). 45. Kawakami, Y. et al. Wnt/β-catenin signaling regulates vertebrate limb regeneration. Genes Dev. 20, 3232–3237 (2006). 46. Miras-Portugal, M. T. et al. Nucleotides in neuroregeneration and neuroprotection. Neuropharmacology 104, 243–254 (2016). 47. Zhang, Y., Zhao, Q., Chen, Q., Xu, L. & Yi, S. Transcriptional control of peripheral nerve regeneration. Mol. Neurobiol. 60, 329–341 (2023). 48. Yang, S.-G., Wang, X.-W., Qian, C. & Zhou, F.-Q. Reprogramming neurons for regeneration: The fountain of youth. Prog. Neurobiol. 214, 102284 (2022). 49. Jhamb, D. et al. Network based transcription factor analysis of regenerating axolotl limbs. BMC Bioinformatics 12, 1–12 (2011). 50. Tolchin, D. et al. De novo SOX6 variants cause a neurodevelopmental syndrome .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint associated with ADHD, craniosynostosis, and osteochondromas. Am. J. Hum. Genet. 106, 830–845 (2020). 51. Chen, J. & Guan, Z. Function of oncogene Mycn in adult neurogenesis and oligodendrogenesis. Mol. Neurobiol. 59, 77–92 (2022). 52. Ballas, N. & Mandel, G. The many faces of REST oversee epigenetic programming of neuronal genes. Curr. Opin. Neurobiol. 15, 500–506 (2005). 53. Cheng, Y. et al. Transcription factor network analysis identifies REST/NRSF as an intrinsic regulator of CNS regeneration in mice. Nat. Commun. 13, 4418 (2022). 54. He, F. & Sun, Y. E. Glial cells more than support cells? Int. J. Biochem. Cell Biol. 39, 661–665 (2007). 55. Laug, D., Glasgow, S. M. & Deneen, B. A glial blueprint for gliomagenesis. Nat. Rev. Neurosci. 19, 393–403 (2018). 56. Tirosh, I. et al. Dissecting the multicellular ecosystem of metastatic melanoma by single-cell RNA-seq. Science 352, 189–196 (2016). 57. Hausrat, T. J., Radwitz, J., Lombino, F. L., Breiden, P. & Kneussel, M. Alpha-and beta-tubulin isotypes are differentially expressed during brain development. Dev. Neurobiol. 81, 333–350 (2021). 58. Anlauf, E. & Derouiche, A. Glutamine synthetase as an astrocytic marker: its cell type and vesicle localization. Front. Endocrinol. 4, 144 (2013). 59. Jurga, A. M., Paleczna, M., Kadluczka, J. & Kuter, K. Z. Beyond the GFAP-astrocyte protein markers in the brain. Biomolecules 11, 1361 (2021). 60. Liu, Z., Liu, L., Qi, Y., Li, H. & Pan, S. GLIS family zinc finger 3 promoting cell malignant behaviors and NF-κB signaling in glioma. Brain Res. 1770, 147623 (2021). 61. Braccioli, L., Vervoort, S. J., Puma, G., Nijboer, C. H. & Coffer, P. J. SOX4 inhibits oligodendrocyte differentiation of embryonic neural stem cells in vitro by inducing Hes5 expression. Stem Cell Res. 33, 110–119 (2018). 62. Taylor, M. K., Yeager, K. & Morrison, S. J. Physiological Notch signaling promotes gliogenesis in the developing peripheral and central nervous systems. Development 134, .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint 2435–2447 (2007). 63. Hirvonen, H. et al. Differential expression of myc, max and RB1 genes in human gliomas and glioma cell lines. Br. J. Cancer 69, 16–25 (1994). 64. Andersson, A. et al. Single-cell and spatial transcriptomics enables probabilistic inference of cell type topography. Commun. Biol. 3, 565 (2020). 65. Kleshchevnikov, V. et al. Cell2location maps fine-grained cell types in spatial transcriptomics. Nat. Biotechnol. 40, 661–671 (2022). 66. Chen, K. H., Boettiger, A. N., Moffitt, J. R., Wang, S. & Zhuang, X. Spatially resolved, highly multiplexed RNA profiling in single cells. Science 348, aaa6090 (2015). 67. Weinreb, C., Wolock, S., Tusi, B. K., Socolovsky, M. & Klein, A. M. Fundamental limits on dynamic inference from single-cell snapshots. Proc. Natl. Acad. Sci. 115, E2467–E2476 (2018). 68. Terpin, A., Lanzetti, N. & Dörfler, F. Learning Diffusion at Lightspeed. ArXiv Prepr. ArXiv240612616 (2024). 69. Zhou, R., Yang, G., Zhang, Y. & Wang, Y. Spatial transcriptomics in development and disease. Mol. Biomed. 4, (2023). 70. Turing, A. M. The chemical basis of morphogenesis. Bull. Math. Biol. 52, 153–197 (1990). 71. Zhu, J., Zhang, Y.-T., Alber, M. S. & Newman, S. A. Bare bones pattern formation: a core regulatory network in varying geometries reproduces major features of vertebrate limb development and evolution. PLoS One 5, e10892 (2010). 72. Economou, A. D. et al. Periodic stripe formation by a Turing mechanism operating at growth zones in the mammalian palate. Nat. Genet. 44, 348–351 (2012). 73. Korsunsky, I. et al. Fast, sensitive and accurate integration of single-cell data with Harmony. Nat. Methods 16, 1289–1296 (2019). 74. Ambrosio, L., Gigli, N. & Savaré, G. Gradient Flows: In Metric Spaces and in the Space of Probability Measures. (Springer Science & Business Media, 2005). 75. Santambrogio, F. {Euclidean, metric, and Wasserstein} gradient flows: an overview. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint Bull. Math. Sci. 7, 87–154 (2017). 76. Genevay, A., Peyré, G. & Cuturi, M. Learning generative models with Sinkhorn divergences. in International Conference on Artificial Intelligence and Statistics 1608– 1617 (2018). 77. Cuturi, M. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. in Advances in Neural Information Processing Systems (eds. Burges, C. J., Bottou, L., Welling, M., Ghahramani, Z. & Weinberger, K. Q.) vol. 26 (Curran Associates, Inc., 2013). 78. Mémoli, F. Gromov–Wasserstein distances and the metric approach to object matching. Found. Comput. Math. 11, 417–487 (2011). 79. Bunne, C., Alvarez-Melis, D., Krause, A. & Jegelka, S. Learning generative models across incomparable spaces. ArXiv Prepr. ArXiv190505461 (2019). 80. Vayer, T., Chapel, L., Flamary, R., Tavenard, R. & Courty, N. Fused Gromov- Wasserstein distance for structured objects: theoretical foundations and mathematical properties. ArXiv Prepr. ArXiv181102834 (2018). 81. Peyré, G., Cuturi, M. & Solomon, J. Gromov-wasserstein averaging of kernel and distance matrices. in International Conference on Machine Learning 2664–2672 (2016). 82. Fatras, K. et al. Minibatch optimal transport distances; analysis and applications. ArXiv Prepr. ArXiv210101792 (2021). 83. Heek, J. et al. Flax: A neural network library and ecosystem for JAX. (2023). 84. Loshchilov, I. & Hutter, F. Decoupled weight decay regularization. ArXiv Prepr. ArXiv171105101 (2017). 85. Van Dijk, D. et al. Recovering gene interactions from single-cell data using data diffusion. Cell 174, 716–729 (2018). 86. Qin, X. et al. A single-cell perturbation landscape of colonic stem cell polarisation. bioRxiv 2023–02 (2023). 87. Virtanen, P. et al. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nat. Methods 17, 261–272 (2020). 88. Han, H. et al. TRRUST v2: an expanded reference database of human and mouse .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint transcriptional regulatory interactions. Nucleic Acids Res. 46, D380–D386 (2018). 89. Tirosh, I. et al. Single-cell RNA-seq supports a developmental hierarchy in human oligodendroglioma. Nature 539, 309–313 (2016).

The project leading to this manuscript has received funding from the European Union (ERC StG, MULTIview-CELL, 101115618). In addition, this work has been funded by the French government under management of Agence Nationale de la Recherche as part of the "Investissements d'avenir" program, reference ANR-19-P3IA-0001 (PRAIRIE 3IA Institute) and by the Inception program (Investissement d’Avenir grant ANR-16-CONV-0005). The work of G. Peyré was supported by the European Research Council (ERC project NORIA) and the French government under management of Agence Nationale de la Recherche as part of the ‘Investissements d’avenir’ program, reference ANR19-P3IA-0001 (PRAIRIE 3IA Institute). We acknowledge the help of the HPC Core Facility of the Institut Pasteur and Déborah Philipps for the administrative support. Competing interests The authors declare no competing interests. .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint .CC-BY-ND 4.0 International licensemade available under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is The copyright holder for this preprintthis version posted July 26, 2024. ; https://doi.org/10.1101/2024.07.26.605241doi: bioRxiv preprint