{"id":18887112,"url":"https://github.com/thunlp-mt/get","last_synced_at":"2025-04-14T22:33:45.994Z","repository":{"id":238782739,"uuid":"796511560","full_name":"THUNLP-MT/GET","owner":"THUNLP-MT","description":"This repo contains the codes for our paper \"Generalist Equivariant Transformer Towards 3D Molecular Interaction Learning\" (ICML 2024).","archived":false,"fork":false,"pushed_at":"2024-09-21T02:40:09.000Z","size":2232,"stargazers_count":30,"open_issues_count":2,"forks_count":3,"subscribers_count":4,"default_branch":"main","last_synced_at":"2025-03-28T10:54:14.773Z","etag":null,"topics":["affinity-prediction","graph-transformer","unified-molecules"],"latest_commit_sha":null,"homepage":"https://proceedings.mlr.press/v235/kong24b.html","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/THUNLP-MT.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-05-06T05:00:43.000Z","updated_at":"2025-03-20T17:30:06.000Z","dependencies_parsed_at":null,"dependency_job_id":"951566a4-60b4-45e8-b40d-c8c2ac58b503","html_url":"https://github.com/THUNLP-MT/GET","commit_stats":null,"previous_names":["thunlp-mt/get"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/THUNLP-MT%2FGET","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/THUNLP-MT%2FGET/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/THUNLP-MT%2FGET/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/THUNLP-MT%2FGET/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/THUNLP-MT","download_url":"https://codeload.github.com/THUNLP-MT/GET/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248973005,"owners_count":21191897,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["affinity-prediction","graph-transformer","unified-molecules"],"created_at":"2024-11-08T07:34:14.331Z","updated_at":"2025-04-14T22:33:40.978Z","avatar_url":"https://github.com/THUNLP-MT.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# GET: Generalist Equivariant Transformer Towards 3D Molecular Interaction Learning\n\n![cover](./assets/cover.png)\n\nThis repo contains the codes for our paper [Generalist Equivariant Transformer Towards 3D Molecular Interaction Learning (ICML 2024)](https://arxiv.org/abs/2306.01474).\n\n## Quick Links\n- [Cleaned Codes of GET](#cleaned-codes-of-get)\n- [Setup](#setup)\n    - [Environment](#environment)\n    - [Datasets](#datasets)\n- [Experiments](#experiments)\n    - [PDBbind Benchmark](#pdbbind-benchmark)\n    - [Protein-Protein Affinity (PPA)](#protein-protein-affinity-ppa)\n    - [Ligand Binding Affinity (LBA)](#ligand-binding-affinity-lba)\n    - [Ligand Efficacy Prediction (LEP)](#ligand-efficacy-prediction-lep)\n    - [Data Augmentation from Different Domains](#data-augmentation-from-different-domains)\n    - [Zero-Shot on Nucleic Acid and Ligand Affinity](#zero-shot-on-nucleic-acid-and-ligand-affinity)\n- [Contact](#contact)\n- [Reference](#reference)\n\n## Cleaned Codes of GET\n\nIf you are interested in using GET for other tasks, we have provided a cleaned version in `./get_clean` that can be easily integrated into other projects. The only requirements are:\n\n```bash\ntorch\ntorch_scatter\nscipy\n```\n\n**Example Code**\n\n```python\nfrom get_clean import GET, fully_connect_edges, knn_edges\nimport torch\n\n# use GPU 0\ndevice = torch.device('cuda:0')\n\n# Dummy model parameters\nd_hidden = 64   # hidden size\nd_radial = 16   # mapped size of key/value in attention (greatly influence space complexity)\nn_channel = 1   # number of channels for coordinates, usually 1 as one atom only has one coordinate\nd_edge = 16     # edge feature size\nn_rbf = 16      # RBF kernal size\nn_head= 4       # number of heads for multi-head attention\n\n# Dummy variables h, x and fully connected edges\n# 19 atoms, divided into 8 blocks\nblock_ids = torch.tensor([0,0,1,1,1,1,2,2,2,3,4,4,5,6,6,6,6,7,7], dtype=torch.long).to(device)\n# 8 blocks, divided into 2 graphs\nbatch_ids = torch.tensor([0,0,0,0,0,1,1,1], dtype=torch.long).to(device)\nn_atoms, n_blocks = block_ids.shape[0], batch_ids.shape[0]\nH = torch.randn(n_atoms, d_hidden, device=device)\nX = torch.randn(n_atoms, n_channel, 3, device=device)\n# fully connect edges\nsrc_dst = fully_connect_edges(batch_ids)\n# if you want to test knn_edges, you can try:\n# src_dst = knn_edges(block_ids, batch_ids, X, k_neighbors=5)\nedge_attr = torch.randn(len(src_dst[0]), d_edge).to(device)\n\n# Initialize GET\nmodel = GET(d_hidden, d_radial, n_channel, n_rbf, d_edge=d_edge, n_head=n_head)\nmodel.to(device)\nmodel.eval()\n\n# Run GET\nH, X = model(H, X, block_ids, batch_ids, src_dst, edge_attr)\n```\n\n## Setup\n\n### Environment\n\nWe have prepared the configuration for creating the environment with conda in *env.yml*:\n```bash\nconda env create -f env.yml\n```\n\n### Datasets\n\nWe assume all the datasets are downloaded to the folder *./datasets*.\n\n#### 1. Protein-Protein Affinity (PPA)\n\nFirst download and decompress the protein-protein complexes in [PDBbind](http://www.pdbbind.org.cn/download.php) (registration is required):\n\n```bash\nwget http://www.pdbbind.org.cn/download/PDBbind_v2020_PP.tar.gz -P ./datasets/PPA\ntar zxvf ./datasets/PPA/PDBbind_v2020_PP.tar.gz -C ./datasets/PPA\nrm ./datasets/PPA/PDBbind_v2020_PP.tar.gz\n```\n\nThen process the dataset with the provided script:\n\n```bash\npython scripts/data_process/process_PDBbind_PP.py \\\n    --index_file ./datasets/PPA/PP/index/INDEX_general_PP.2020 \\\n    --pdb_dir ./datasets/PPA/PP \\\n    --out_dir ./datasets/PPA/processed\n```\n\nThe processed data will be saved to *./datasets/PPA/processed/PDBbind.pkl*. This should result in 2502 valid complexes with affinity labels.\n\nWe still need to prepare the test set, i.e. Protein-Protein Affinity Benchmark Version 2. We have provided the index file in *./datasets/PPA/PPAB_V2.csv*, but the structure files need to be downloaded from the [official site](https://zlab.umassmed.edu/benchmark/). In case the official site is down, we also uploaded [a backup on Zenodo](https://zenodo.org/record/8318025/files/benchmark5.5.tgz?download=1).\n\n```bash\n# If the offical url is not working, please replace it with https://zenodo.org/record/8318025/files/benchmark5.5.tgz?download=1\nwget https://zlab.umassmed.edu/benchmark/benchmark5.5.tgz -O ./datasets/PPA/benchmark5.5.tgz\ntar zxvf ./datasets/PPA/benchmark5.5.tgz -C ./datasets/PPA\nrm ./datasets/PPA/benchmark5.5.tgz\n```\n\nThen process the test set with the provided script:\n\n```bash\npython scripts/data_process/process_PPAB.py \\\n    --index_file ./datasets/PPA/PPAB_V2.csv \\\n    --pdb_dir ./datasets/PPA/benchmark5.5 \\\n    --out_dir ./datasets/PPA/processed\n```\n\nThe processed dataset as well as different splits (Rigid/Medium/Flexible/All) will be saved to *./datasets/PPA/processed*. This should result in 176 valid complexes with affinity labels.\n\n#### 2. PDBBind Benchmark (established splits)\n\nFirst download and extract the raw files:\n\n```bash\nmkdir ./datasets/PDBBind\nwget \"https://zenodo.org/record/8102783/files/pdbbind_raw.tar.gz?download=1\" -O ./datasets/PDBBind/pdbbind_raw.tar.gz\ntar zxvf ./datasets/PDBBind/pdbbind_raw.tar.gz -C ./datasets/PDBBind\nrm ./datasets/PDBBind/pdbbind_raw.tar.gz\n```\n\nThen process the dataset with the provided script:\n```bash\npython scripts/data_process/process_PDBbind_benchmark.py \\\n    --benchmark_dir ./datasets/PDBBind/pdbbind \\\n    --out_dir ./datasets/PDBBind/processed\n```\n\nThis should result in 4709 valid complexes with affinity labels.\n\nWhat is different here is that if you want to use fragment-based representation of small molecules, you need to process the data here:\n\n```bash\npython scripts/data_process/process_PDBbind_benchmark.py \\\n    --benchmark_dir ./datasets/PDBBind/pdbbind \\\n    --fragment PS_300 \\\n    --out_dir ./datasets/PDBBind/processed_PS_300\n```\n\n#### 3. Ligand Binding Affinity (LBA)\n\nYou only need to download and decompress the LBA dataset:\n\n```bash\nmkdir ./datasets/LBA\nwget \"https://zenodo.org/record/4914718/files/LBA-split-by-sequence-identity-30.tar.gz?download=1\" -O ./datasets/LBA/LBA-split-by-sequence-identity-30.tar.gz\ntar zxvf ./datasets/LBA/LBA-split-by-sequence-identity-30.tar.gz -C ./datasets/LBA\nrm ./datasets/LBA/LBA-split-by-sequence-identity-30.tar.gz\n```\n\n#### 4. Zero-Shot Inference on Nucleic-Acid-Ligand Affinity (NLA)\n\nWe need to use protein-protein data, protein-nucleic-acid data, and protein-ligand data for training, then evaluate the zero-shot performance on nucleic-acid-ligand affinity. All the data are extracted from PDBBind database. We have got protein-protein data in PPA and protein-ligand data in LBA, now we further need to get other data.\nTo get protein-nucleic-acid data:\n\n```bash\nwget http://www.pdbbind.org.cn/download/PDBbind_v2020_PN.tar.gz -P ./datasets/PN\ntar zxvf ./datasets/PN/PDBbind_v2020_PN.tar.gz -C ./datasets\nrm ./datasets/PN/PDBbind_v2020_PN.tar.gz\n```\nThen process the data:\n\n```bash\npython scripts/data_process/process_PDBbind_PN.py \\\n    --index_file ./datasets/PN/index/INDEX_general_PN.2020 \\\n    --pdb_dir ./datasets/PN \\\n    --out_dir ./datasets/PN/processed\n```\n\nThis should result in 922 valid complexes with affinity labels.\n\nTo get nucleic-acid-ligand data:\n\n```bash\nwget http://www.pdbbind.org.cn/download/PDBbind_v2020_NL.tar.gz -P ./datasets/NL\ntar zxvf ./datasets/NL/PDBbind_v2020_NL.tar.gz -C ./datasets\nrm ./datasets/NL/PDBbind_v2020_NL.tar.gz\n```\n\nThen process the data:\n\n```bash\npython scripts/data_process/process_PDBbind_NL.py \\\n    --index_file ./datasets/NL/index/INDEX_general_NL.2020 \\\n    --pdb_dir ./datasets/NL \\\n    --out_dir ./datasets/NL/processed\n```\n\nThis should result in 134 valid complexes with affinity labels.\n\n#### 5 (Optional). Ligand Efficacy Prediction (LEP)\n\nYou only need to download and decompress the LEP dataset:\n\n```bash\nmkdir ./datasets/LEP\nwget \"https://zenodo.org/record/4914734/files/LEP-split-by-protein.tar.gz?download=1\" -O ./datasets/LEP/LEP-split-by-protein.tar.gz\ntar zxvf ./datasets/LEP/LEP-split-by-protein.tar.gz -C ./datasets/LEP\nrm ./datasets/LEP/LEP-split-by-protein.tar.gz\n```\n\n## Experiments\n\n:warning: Due to the non-deterministic behavior of `torch.scatter`, the reproduced results might not be exactly the same as those reported in the paper, but should be very close to them.\n\n### PDBbind Benchmark\n\nWe have provided the script for running the experiment with 3 random seeds:\n\n```bash\npython scripts/exps/exps_3.py \\\n    --config ./scripts/exps/configs/PDBBind/identity30_get.json \\\n    --gpus 0\n```\n\n### Protein-Protein Affinity (PPA)\n\nWe have provided the script for splitting, training and testing with 3 random seeds:\n\n```bash\npython scripts/exps/PPA_exps_3.py \\\n    --pdbbind ./datasets/PPA/processed/PDBbind.pkl \\\n    --ppab_dir ./datasets/PPA/processed \\\n    --config ./scripts/exps/configs/PPA/get.json \\\n    --gpus 0\n```\n\n### Ligand Binding Affinity (LBA)\n\nWe have provided the script for training and testing with 3 random seeds:\n\n```bash\npython scripts/exps/exps_3.py \\\n    --config ./scripts/exps/configs/LBA/get.json \\\n    --gpus 0\n```\n\nIf you want to use fragment-based representation of small molecules, please replace the config with `get_ps300.json`.\n\n### Ligand Efficacy Prediction (LEP)\n\nWe have provided the script for training and testing with 3 random seeds:\n\n```bash\npython scripts/exps/exps_3.py \\\n    --config ./scripts/exps/configs/LEP/get.json \\\n    --gpus 0\n```\n\n### Data Augmentation from Different Domains\n\nTo enhance the performance on PPA with additional data on LBA:\n\n```bash\npython scripts/exps/mix_exps_3.py \\\n    --config ./scripts/exps/configs/MIX/get_ppa.json \\\n    --gpus 0\n```\n\nTo enhance the performance on LBA with additional data on PPA:\n\n```bash\npython scripts/exps/mix_exps_3.py \\\n    --config ./scripts/exps/configs/MIX/get_lba.json \\\n    --gpus 0\n```\n\nTo enhance the performance on PDBbind with additional data on PPA:\n\n```bash\n# !!!Here we use exps_3.py instead of mix_exps_3.py\npython scripts/exps/exps_3.py \\\n    --config ./scripts/exps/configs/MIX/get_pdbbind.json \\\n    --gpus 0\n```\n\n\n### Zero-Shot on Nucleic Acid and Ligand Affinity\n\nThis experiment needs two 12G GPU (so a total of 2000 vertexes in a batch). If you use a single \u003e=24G GPU, please change the value of `max_n_vertex_per_gpu` to 2000.\n\n```bash\npython scripts/exps/NL_zeroshot.py \\\n    --config scripts/exps/configs/NL/get.json \\\n    --gpu 1 2\n```\n\n## Contact\n\nThank you for your interest in our work!\n\nPlease feel free to ask about any questions about the algorithms, codes, as well as problems encountered in running them so that we can make it clearer and better. You can either create an issue in the github repo or contact us at jackie_kxz@outlook.com.\n\n## Reference\n\n```bibtex\n@InProceedings{pmlr-v235-kong24b,\n  title = \t {Generalist Equivariant Transformer Towards 3{D} Molecular Interaction Learning},\n  author =       {Kong, Xiangzhe and Huang, Wenbing and Liu, Yang},\n  booktitle = \t {Proceedings of the 41st International Conference on Machine Learning},\n  pages = \t {25149--25175},\n  year = \t {2024},\n  editor = \t {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix},\n  volume = \t {235},\n  series = \t {Proceedings of Machine Learning Research},\n  month = \t {21--27 Jul},\n  publisher =    {PMLR},\n  pdf = \t {https://raw.githubusercontent.com/mlresearch/v235/main/assets/kong24b/kong24b.pdf},\n  url = \t {https://proceedings.mlr.press/v235/kong24b.html},\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fthunlp-mt%2Fget","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fthunlp-mt%2Fget","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fthunlp-mt%2Fget/lists"}