gavinraym / trajectorynet

TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics

Pytorch Implementation of TrajectoryNet

This library runs code associated with the TrajectoryNet paper [1].

In brief, TrajectoryNet is a Continuous Normalizing Flow model which can perform dynamic optimal transport using energy regularization and / or a combination of velocity, density, and growth regularizations to better match cellular trajectories.

Our setting is similar to that of WaddingtonOT. In that we have access to a bunch of population measurements of cells over time and would like to model the dynamics of cells over that time period. TrajectoryNet is trained end-to-end and is continuous both in gene space and in time.

Installation

TrajectoryNet is available in pypi. Install by running the following

pip install TrajectoryNet

This code was tested with python 3.7 and 3.8.

Example

Basic Usage

Run with

python -m TrajectoryNet.main --dataset SCURVE

To run TrajectoryNet on the S Curve example in the paper. To use a custom dataset expose the coordinates and timepoint information according to the example jupyter notebooks in the /notebooks/ folder.

If you have an AnnData object then take a look at notebooks/Example_Anndata_to_TrajectoryNet.ipynb, which shows how to load one of the example scvelo anndata objects into TrajectoryNet. Alternatively you can use the custom (compressed) format for TrajectoryNet as described below.

For this format TrajectoryNet requires the following:

1. An embedding matrix titled [embedding_name] (Cells x Dimensions)
2. A sample labels array titled sample_labels (Cells)

To run TrajectoryNet with a custom dataset use:

python -m TrajectoryNet.main --dataset [PATH_TO_NPZ_FILE] --embedding_name [EMBEDDING_NAME]
python -m TrajectoryNet.eval --dataset [PATH_TO_NPZ_FILE] --embedding_name [EMBEDDING_NAME]

See notebooks/EB-Eval.ipynb for an example on how to use TrajectoryNet on a PCA embedding to get trajectories in the gene space.

References

[1] Tong, A., Huang, J., Wolf, G., van Dijk, D., and Krishnaswamy, S. TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics. In International Conference on Machine Learning, 2020. arxiv ICML

---

If you found this library useful, please consider citing:

@inproceedings{tong2020trajectorynet,
  title = {TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics},
  shorttitle = {TrajectoryNet},
  booktitle = {Proceedings of the 37th International Conference on Machine Learning},
  author = {Tong, Alexander and Huang, Jessie and Wolf, Guy and {van Dijk}, David and Krishnaswamy, Smita},
  year = {2020}
}