ContextFlow: Context-Aware Flow Matching For Trajectory Inference From Spatial Omics Data

TL;DR

ContextFlow tackles the challenge of inferring continuous spatiotemporal tissue trajectories from snapshot spatial omics data by marrying flow matching with context-aware optimal transport. It introduces a transitional plausibility framework that encodes local tissue structure and ligand–receptor signaling into OT couplings, via two integration schemes: PACM and PAER. The method demonstrates consistent improvements over state-of-the-art baselines across three datasets, yielding trajectories that are both statistically accurate (via $\mathcal{W}_2$, MMD, and Energy metrics) and biologically coherent (reduced implausible transitions and coherent cell-type progression). The entropic variant, ContextFlow-H, provides robust generalization and eliminates the need for extensive hyperparameter tuning, offering a scalable, interpretable approach for spatiotemporal dynamics in spatial omics. $ContextFlow$ thus provides a principled, generalizable framework for integrating biological priors into trajectory inference from spatial omics data, with potential extensions to spatial flow modeling and perturbation analysis.

Abstract

Inferring trajectories from longitudinal spatially-resolved omics data is fundamental to understanding the dynamics of structural and functional tissue changes in development, regeneration and repair, disease progression, and response to treatment. We propose ContextFlow, a novel context-aware flow matching framework that incorporates prior knowledge to guide the inference of structural tissue dynamics from spatially resolved omics data. Specifically, ContextFlow integrates local tissue organization and ligand-receptor communication patterns into a transition plausibility matrix that regularizes the optimal transport objective. By embedding these contextual constraints, ContextFlow generates trajectories that are not only statistically consistent but also biologically meaningful, making it a generalizable framework for modeling spatiotemporal dynamics from longitudinal, spatially resolved omics data. Evaluated on three datasets, ContextFlow consistently outperforms state-of-the-art flow matching methods across multiple quantitative and qualitative metrics of inference accuracy and biological coherence. Our code is available at: \href{https://github.com/santanurathod/ContextFlow}{ContextFlow}

Figures (9)

• Figure 1: ContextFlow integrates local tissue organization and ligand-receptor communications to learn biologically meaningful trajectories from spatial omics data. Prior knowledge acts as a soft filter that discourages implausible transitions while preserving flexibility in trajectory inference.
• Figure 2: KL-Divergence between predicted and ground-truth cell type distributions.
• Figure 3: Comparison of biologically implausible cell type couplings between Stage 0 and Stage 1 of the Brain Regeneration Dataset 1_dataset_steroseq, under the Entropic-OT and ContextFlow Regularized-OT formulations. Biological implausibility is defined here as transitions involving excitatory--inhibitory lineage switches. Our formulation produces substantially fewer biologically implausible couplings (24) compared to MOTFM (54).
• Figure 4: Temporal progression of spatial distribution of different cell types for Brain Regeneration.
• Figure 5: Temporal progression of spatial distribution of different cell types for Mouse Organogenesis.

Theorems & Definitions (4)

• proof
• proof
• proof
• proof