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

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}

💡 Research Summary

ContextFlow introduces a novel framework for inferring temporal trajectories from longitudinal spatial omics data by embedding biologically relevant contextual information into a flow‑matching paradigm. Traditional flow‑matching methods learn a time‑varying velocity field that transports a source distribution to a target distribution, often relying on optimal transport (OT) couplings that ignore spatial organization and cell‑cell communication. Consequently, generated trajectories may be statistically optimal but biologically implausible.

To address this, ContextFlow incorporates two complementary spatial priors: (1) local tissue smoothness and (2) ligand‑receptor communication networks. The smoothness prior assumes that neighboring cells within a microenvironment experience similar mechanical and signaling cues, leading to coherent expression profiles. By aggregating the expression of each cell’s neighborhood, the method quantifies a transition plausibility score between consecutive time points. The ligand‑receptor prior leverages curated interaction databases (e.g., CellPhoneDB) to identify signaling pairs that are biologically feasible between cell types. Both priors are encoded into a transition plausibility matrix that modifies the OT coupling cost.

Two integration schemes are proposed. In the cost‑based scheme, the plausibility matrix is added directly to the transport cost before solving the entropic OT problem with the Sinkhorn algorithm. In the entropy‑based scheme, the priors are introduced as additional entropy regularization terms, effectively reshaping the coupling distribution. Both schemes remain compatible with minibatch OT, preserving an O(N²) computational complexity and enabling efficient GPU training.

The model is trained using conditional flow matching (CFM) with entropic OT couplings for each consecutive pair of time points. For a given interval