README.md

<!--
 * @Author: zhuww
 * @email: zhuww@sugon.com
 * @Date: 2023-04-06 18:04:07
 * @LastEditTime: 2023-12-26 15:54:01
-->
# AF2
## 论文
- [https://www.nature.com/articles/s41586-021-03819-2](https://www.nature.com/articles/s41586-021-03819-2)

## 模型结构
模型核心是一个基于Transformer架构的神经网络，包括两个主要组件：Sequence to Sequence Model和Structure Model，这两个组件通过迭代训练进行优化，以提高其预测准确性。

![img](./docs/alphafold2.png)

## 算法原理
AlphaFold2通过从蛋白质序列和结构数据中提取信息，使用神经网络模型来预测蛋白质三维结构。

![img](./docs/alphafold2_1.png)


## 环境配置

### Docker
    
    docker pull image.sourcefind.cn:5000/dcu/admin/base/jax:0.4.23-ubuntu20.04-dtk24.04.1-py3.10

    docker run --shm-size 50g --network=host --name=alphafold2 --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v 项目地址(绝对路径):/home/ -v /opt/hyhal:/opt/hyhal:ro -it <your IMAGE ID> bash

    # 1. 一般依赖项安装
    pip install -r requirements_dcu.txt

    pip install dm-haiku==0.0.11 flax==0.7.1 jmp==0.0.2 tabulate==0.8.9 --no-deps jax

    pip install orbax==0.1.6 orbax-checkpoint==0.1.6 optax==0.2.2

    python setup.py install

    # 2、hh-suite 

    git clone https://github.com/soedinglab/hh-suite.git
    mkdir -p hh-suite/build && cd hh-suite/build
    cmake -DCMAKE_INSTALL_PREFIX=. ..
    make -j 4 && make install
    export PATH="$(pwd)/bin:$(pwd)/scripts:$PATH"

    wget https://github.com/TimoLassmann/kalign/archive/refs/tags/v3.4.0.zip
    unzip 3.4.0.zip && cd kalign-3.4.0
    mkdir build 
    cd build
    cmake .. 
    make 
    make test 
    make install

    # 3. openmm + pdbfixer

    sudo apt install doxygen

    wget https://github.com/openmm/openmm/archive/refs/tags/8.0.0.zip

    unzip 8.0.0.zip && cd openmm-8.0.0 && mkdir build && cd build

    cmake .. && make && sudo make install && sudo make PythonInstall

    wget https://github.com/openmm/pdbfixer/archive/refs/tags/1.9.zip

    unzip 1.9.zip && cd pdbfixer-1.9 && python setup.py install 

<!-- ## 环境配置
提供[光源](https://www.sourcefind.cn/#/image/dcu/custom)拉取推理的docker镜像：
```
docker pull image.sourcefind.cn:5000/dcu/admin/base/custom:alphafold2-2.3.2-dtk23.10-py38
# <Image ID>用上面拉取docker镜像的ID替换
# <Host Path>主机端路径
# <Container Path>容器映射路径
docker run -it --name alphafold --privileged --shm-size=32G  --device=/dev/kfd --device=/dev/dri/ --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --ulimit memlock=-1:-1 --ipc=host --network host --group-add video -v <Host Path>:<Container Path> <Image ID> /bin/bash
```

镜像版本依赖：
* DTK驱动：dtk23.10
* Jax: 0.3.25
* TensorFlow2: 2.11.0
* python: python3.8 -->

## 数据集
推荐使用AlphaFold2中的开源数据集，包括BFD、MGnify、PDB70、Uniclust、Uniref90等,数据集大小约2.62TB。数据集格式如下：
```
$DOWNLOAD_DIR/                             
    bfd/  
        bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt_hhm.ffindex
        bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt_hhm.ffdata 
        bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt_cs219.ffindex                           
        ...
    mgnify/                                
        mgy_clusters_2022_05.fa
    params/                                
        params_model_1.npz
        params_model_2.npz
        params_model_3.npz
        ...
    pdb70/                                
        pdb_filter.dat
        pdb70_hhm.ffindex
        pdb70_hhm.ffdata
        ...
    pdb_mmcif/                            
        mmcif_files/
            100d.cif
            101d.cif
            101m.cif
            ...
        obsolete.dat
    pdb_seqres/                            
        pdb_seqres.txt
    small_bfd/                           
        bfd-first_non_consensus_sequences.fasta
    uniref30/                            
        UniRef30_2021_03_hhm.ffindex
        UniRef30_2021_03_hhm.ffdata
        UniRef30_2021_03_cs219.ffindex
        ...
    uniprot/                               
        uniprot.fasta
    uniref90/                             
        uniref90.fasta
```

此处提供了一个脚本download_all_data.sh用于下载使用的数据集和模型文件：

    ./scripts/download_all_data.sh 数据集下载目录

数据集快速下载中心：[SCNet AIDatasets](http://113.200.138.88:18080/aidatasets) ，项目中数据集可从快速下载通道下载：[alphafold](http://113.200.138.88:18080/aidatasets/project-dependency/alphafold) 。


## 推理
分别提供了基于Jax的单体和多体的推理脚本.
```bash
    # 进入工程目录
    cd alphafold2_jax
```

### 单体
```bash
    ./run_monomer.sh
```
单体推理参数说明：download_dir为数据集下载目录，monomer.fasta为推理的单体序列；`--output_dir`为输出目录；`model_names`为推理的模型名称，`--model_preset=monomer`为单体模型配置；`--run_relax=true`为进行relax操作；`--use_gpu_relax=true`为使用gpu进行relax操作（速度更快，但可能不太稳定），`--use_gpu_relax=false`为使用CPU进行relax操作（速度慢，但稳定）；若添加--use_precomputed_msas=true则可以加载已有的MSAs，否则默认运行MSA工具。

### 多体
```bash
    ./run_multimer.sh
```
多体推理参数说明：multimer.fasta为推理的多体序列，`--model_preset=multimer`为多体模型配置；`--num_multimer_predictions_per_model`为每个模型预测数量，其他参数同单体推理参数说明一致。

## result
`--output_dir`目录结构如下：
```
<target_name>/
    features.pkl
    ranked_{0,1,2,3,4}.pdb
    ranking_debug.json
    relaxed_model_{1,2,3,4,5}.pdb
    result_model_{1,2,3,4,5}.pkl
    timings.json
    unrelaxed_model_{1,2,3,4,5}.pdb
    msas/
        bfd_uniclust_hits.a3m
        mgnify_hits.sto
        uniref90_hits.sto
        ...
```

[查看蛋白质3D结构](https://www.pdbus.org/3d-view)
<div style="display: flex; justify-content: center; align-items: center;">
  <img src="./docs/result_pdb.png" alt="Image">
  <div style="position: absolute; top: 50%; left: 50%; transform: translate(-50%, -50%); background: rgba(0, 0, 0, 0.5); color: #fff; padding: 10px;">
    红色为真实结构，蓝色为预测结构。
  </div>
</div>

### 精度
测试数据：[casp15](https://www.predictioncenter.org/casp15/targetlist.cgi)、[uniprot](https://www.uniprot.org/)，
使用的加速卡:1张 Z100L-32G

1、plddts/iptm+ptm

单体见<target_name>/ranking_debug.json中的`plddts`，多体见<target_name>/ranking_debug.json中的`iptm+ptm`


2、其它精度值计算：[https://zhanggroup.org/TM-score/](https://zhanggroup.org/TM-score/)

准确性数据：
| 数据类型 | 序列类型 | 序列 | 长度 | GDT-TS | GDT-HA | plddts/iptm+ptm | TM score | MaxSub | RMSD |
| :------: | :------: | :------: |:------: |:------: | :------: | :------: | :------: |:------: |:------: |
| fp32 | 单体 | T1029 | 125 | 0.434 | 0.256 | 93.984 | 0.471 | 0.297 | 7.202 |
| fp32 | 单体 | T1024 | 408 | 0.664 | 0.470 | 87.076 | 0.829 | 0.518 | 3.516 |
| fp32 | 多体 | H1106 | 236 | 0.203 | 0.144 | 0.860 | 0.181 | 0.151 | 20.457 |

## 应用场景

### 算法类别
蛋白质预测

### 热点应用行业
医疗,科研,教育

## 预训练权重
预训练权重快速下载中心：[SCNet AIModels](http://113.200.138.88:18080/aimodels) ，项目中的预训练权重可从快速下载通道下载：[alphafold](http://113.200.138.88:18080/aimodels/findsource-dependency/alphafold-params) 。

## 源码仓库及问题反馈
* [https://developer.hpccube.com/codes/modelzoo/alphafold2_jax](https://developer.hpccube.com/codes/modelzoo/alphafold2_jax)

## 参考
* [https://github.com/deepmind/alphafold](https://github.com/deepmind/alphafold)