CI-FM: Cellular Interaction Foundation Model

Overview

This is the PyTorch implementation of the CI-FM model -- an AI model that can simulate the activities within a living tissue (AI virtual tissue). The current version of CI-FM has 138M parameters and is trained on around 23M cells of spatial genomics. The signature functions of CI-FM are:

Embedding of celllular microenvironments via embeddings = model.embed(adata) (the 1st Figure below panel D top);
Inference/simulation of cellular gene expressions within a certain microenvironment via expressions = model.predict_cells_at_locations(adata, target_locs) (the 1st Figure below panel D bottom, and the 2nd Figure below).

The detailed usage of the model can be found in the tutorial. Before running the tutorial, please set up an environment following the environment instruction.

More information about the model can be found in the preprint.

Environment

I use conda to manage the environment conda create -n $MYENV python=3.11, but it is not the only way to do that. After entering the environment, you can clone the model via:

# Make sure you have git-lfs installed (https://git-lfs.com), e.g. as:
git lfs install
# or if you are under conda environment, you can also install git-lfs via conda as:
conda install anaconda::git-lfs

# clone the model and enter the directory
git clone https://huggingface.co/ynyou/CIFM
cd ./CIFM/
git lfs install ; git lfs pull

The environment can be set up via:

# my python version is 3.11, which is manage when creating conda environment: conda create -n $MYENV python=3.11
pip install torch==2.1.0 --index-url https://download.pytorch.org/whl/cu118
pip install torch-scatter torch-sparse torch-cluster torch-geometric -f https://data.pyg.org/whl/torch-2.1.0+cu118.html
pip install scanpy e3nn lightning==2.1.0 transformers numpy==1.26.4

where the compatibility between torch and torch-geometric is not always guaranteed (lines 2-3), since these two guys are very picky on the platform and the version of other packages. You may need trial-and-error to find the right versions of torch and torch-geometric that works for you, e.g., in some machine I once installed via:

pip install torch==2.0.1 --index-url https://download.pytorch.org/whl/cu117
pip install torch-scatter torch-sparse torch-cluster torch-geometric -f https://data.pyg.org/whl/torch-2.0.1+cu117.html

Citation

If you use this code for you research, please cite our paper.

@article{you2025building,
  title={Building Foundation Models to Characterize Cellular Interactions via Geometric Self-Supervised Learning on Spatial Genomics},
  author={You, Yuning and Wang, Zitong and Fleisher, Kevin and Liu, Rex and Thomson, Matt},
  journal={bioRxiv},
  pages={2025--01},
  year={2025},
  publisher={Cold Spring Harbor Laboratory}
}