Featurization of Single Cell Trajectories Through Kernel Mean Embedding of Optimal Transport Maps

Abstract

Longitudinal single-cell data has spurred the development of computational trajectory models with the power to make time-resolved, testable predictions about cell fates. As ”real-time” trajectory inference methods proliferate, there is a growing need for tools that integrate their inherently high-dimensional outputs. In this work, we propose a novel strategy to facilitate downstream analysis of single-cell optimal-transport trajectory models, by constructing feature vectors that contain information about a cell’s state across the entirety of its trajectory. This approach leverages kernel mean embedding of distributions to create trajectory features with applications in several domains, including cell clustering and comparison of perturbation response trajectories. We demonstrate how k-means clustering on trajectory features produces interpretable clusters that respect the underlying cell trajectories. Furthermore, we develop a divergence metric between single-cell trajectories based on the maximum mean discrepancy (MMD). We use this trajectory divergence to show that modeling perturbation trajectories may help uncover experimentally interesting perturbations at higher significance levels than by comparing perturbation responses at only a single time point.