[tools] Protenix

安装

打开一个有gpu的环境，20-40GB RAM

source安装

git clone https://github.com/bytedance/Protenix.git
cd Protenix
pip install .

docker安装

从私人的hub上下在docker环境

# downloads from a private docker hub ai4s-cn-beijing
docker pull ai4s-cn-beijing.cr.volces.com/infra/protenix:v0.0.1

tag with username

# tag with my username
docker tag ai4s-cn-beijing.cr.volces.com/infra/protenix:v0.0.1 sky1ove/protenix

upload to docker hub

docker login

docker push sky1ove/protenix

下次直接docker pull sky1ove/protenix就可以。

打开docker 环境的terminal

docker run --gpus all -it -v $(pwd):/workspace -v /dev/shm:/dev/shm sky1ove/protenix /bin/bash

下载ccd 的database

python scripts/gen_ccd_cache.py -c release_data/ccd_cache/ -n [num_cpu]

目前pip install protenix的方法会在下载ccd这一步报错，已经提交了issue。但docker没问题。

因为是在docker里新建了release data文件夹，如果出现权限问题，在外面的terminal设置

sudo chown -R $USER:$USER 

Example test

测试是否安装成功

wget https://files.rcsb.org/download/7pzb.pdb
# run with pdb/cif file, and convert it to json file for inference.
protenix tojson --input examples/7pzb.pdb --out_dir ./output

# for CIF (same process)
# download cif file
wget https://files.rcsb.org/download/7pzb.cif
# run with pdb/cif file, and convert it to json file for inference.
protenix tojson --input examples/7pzb.cif --out_dir ./output

算msa

# run msa search with json file, it will write precomputed msa dir info to a new json file.
protenix msa --input examples/example_without_msa.json --out_dir ./output

# run msa search with fasta file which only contains protein.
protenix msa --input examples/prot.fasta --out_dir ./output

预测

rotenix predict --input examples/example.json --out_dir  ./output --seeds 101

# run with multiple json files, the default seed is 101.
protenix predict --input ./jsons_dir/ --out_dir  ./output

# if the json do not contain precomputed msa dir,
# add --use_msa_server to search msa and then predict.
# if mutiple seeds are provided, split them by comma.
protenix predict --input examples/example_without_msa.json --out_dir ./output --seeds 101,102 --use_msa_server

torchrun

并行

export LAYERNORM_TYPE=fast_layernorm
export USE_DEEPSPEED_EVO_ATTENTION=true

N_sample=5
N_step=200
N_cycle=10
seed=101

input_json_path="./examples/example.json"
dump_dir="./output"

python3 runner/inference.py \
--seeds ${seed} \
--dump_dir ${dump_dir} \
--input_json_path ${input_json_path} \
--model.N_cycle ${N_cycle} \
--sample_diffusion.N_sample ${N_sample} \
--sample_diffusion.N_step ${N_step}

# The following is a demo to use DDP for inference
# torchrun \
#     --nproc_per_node $NPROC \
#     --master_addr $WORKER_0_HOST \
#     --master_port $WORKER_0_PORT \
#     --node_rank=$ID \
#     --nnodes=$WORKER_NUM \
#     runner/inference.py \
#     --seeds ${seed} \
#     --dump_dir ${dump_dir} \
#     --input_json_path ${input_json_path} \
#     --model.N_cycle ${N_cycle} \
#     --sample_diffusion.N_sample ${N_sample} \
#     --sample_diffusion.N_step ${N_step}

单个文件，有pre-computed msa

protenix predict --input examples/example.json --out_dir  ./output --seeds 101

单个文件，没有msa

protenix predict --input examples/example_without_msa.json --out_dir ./output --seeds 101,102 --use_msa_server

注意如果需要多个seeds，用逗号隔开

多个文件，在一个目录下

protenix predict --input ./jsons_dir/ --out_dir  ./output

Input格式

得到MSA

可以用server先算出msa得到a3m文件，一个folder里应该有两个文件，一个pairing.a3m一个unpairing.a3m，然后这个文件的directory作为json文件里precomputed_msa_dir的输入

单个json

.json文件里是[{job1},{job2}]的格式，对于每个job里的{name:str, sequences:[]}

sequences，是[{item1},{item2}]

item1 如果是蛋白，就是proteinChain:{count:1,sequence:seq, msa: {}}

下面是一个.json的例子

[
    {
        "name": "job_name",
        "sequences": [
            {
                "proteinChain": {
                    "count": 1,
                    "sequence": "ABCDEFG",
                    "msa": {
                        "precomputed_msa_dir": "./msa",
                        "pairing_db": "uniref100",
 
                    }
                }
            },
            {
                "ligand": {
                    "count": 1,
                    "ligand": "CCD_Q2T" # or SMILES
                }
            }
        ]
    }
]

一个.json文件里放多个job

那就是[{},{},{}]的格式。

跑json

protenix predict --input input.json --out_dir  ./output --seeds 101

default 种子是101