InfoAlign learns molecular representation from bottleneck information derived from molecular structures, cell morphology, and gene expressions. Readers are referred to the paper for more details: https://arxiv.org/abs/2406.12056v1.
This code was developed and tested with Python 3.11.7, PyTorch 2.1.0+cu118, and torch-geometric 2.5.2. All dependencies are specified in the requirements.txt
file.
We provide a checkpoint which can be downloaded from Hugging Face. Please place the model weights (pretrain.pt
) under the ckpt
folder along with its configurations in the YAML file.
For fine-tuning and inference, use the following code:
python main.py --model-path ckpt/pretrain.pt --dataset finetune-chembl2k
python main.py --model-path ckpt/pretrain.pt --dataset finetune-broad6k
python main.py --model-path ckpt/pretrain.pt --dataset finetune-biogenadme
python main.py --model-path ckpt/pretrain.pt --dataset finetune-moltoxcast
Note: Please visit Hugging Face for the cell morphology and gene expression features in the ChEMBL2k and Broad6K datasets.
To pretrain the model from scratch, first download the pretraining dataset from Hugging Face. Place all pretrain data files under the raw_data/pretrain/raw
folder. Then run the following code:
python main.py --model-path "ckpt/pretrain.pt" --lr 1e-4 --wdecay 1e-8 --batch-size 3072
The pretrained result will be saved in the ckpt
folder with the name pretrain.pt
.
For readers interested in data collection, here are the sources:
-
Cell Morphology Data
- JUMP dataset: The data are from "JUMP Cell Painting dataset: morphological impact of 136,000 chemical and genetic perturbations" and can be downloaded here. The dataset includes chemical and genetic perturbations for cell morphology features.
- Bray's dataset: "A dataset of images and morphological profiles of 30,000 small-molecule treatments using the Cell Painting assay". Download from GigaDB. Processed version available on Zenodo.
-
Gene Expression Data
- LINCS L1000 gene expression data from the paper "Drug-induced adverse events prediction with the LINCS L1000 data": Data.
-
Relationships
- Gene-gene, gene-compound relationships from Hetionet: Data.
If you find this repository useful, please cite our paper:
@article{liu2024learning,
title={Learning Molecular Representation in a Cell},
author={Liu, Gang and Seal, Srijit and Arevalo, John and Liang, Zhenwen and Carpenter, Anne E and Jiang, Meng and Singh, Shantanu},
journal={arXiv preprint arXiv:2406.12056},
year={2024}
}