MPSCF:基于GPU4PYSCF为金属蛋白体系生成力场参数

Chris——szk

一.前言
由于最近有空闲的GPU资源，所以整了一个框架来算金属蛋白/酶的力场参数，按照之前gpu4pyscf本身的框架(版本1.7.0)去计算，在结构优化部分其实不会太爆显存，但是在进行hessian计算的时候，往往会爆显存，最初在24GB单卡上测试68原子体系大小的hessian的时候就爆显存了，所以这次针对这方面进行了调整，同时发现gpu4pyscf本身在计算hessian的时候似乎是默认在单卡上进行，没有考虑在多卡情况下计算hessian的时候多卡之间的显存分配问题，所以一直爆显存。就目前来说的话，基本的单离子体系，同类型多离子体系，和简单配体和同类型离子体系都是能进行计算的，不过在含有配体的体系下还有一点不完善的地方，还在持续跟进，然后像有些复杂的体系，比如金属离子桥接体系，或者HOMO-LUMO gap较小的体系，需要提高一定的level-shift实现收敛，但是个人感觉这类体系对于初始结构的要求还是挺高的，所以也整合了一下tblite的pyberny，外加写了一下ANCopt去通过GFN2-XTB进行预优化。对于一般的简单金属酶体系可以不用开启这个开关，自己测试下来的话单一金属离子模型还是比较好算的，复杂体系的话在下面的json文件中自定义一下是否开启GFN2-XTB进行预优化或者采用哪一个优化器即可。
二.运行
因为整体的逻辑还是基于MCPB本身，所以后续还会考虑整合modseminario方法进行参数生成等。目前的运行效率的话在结构优化部分的速度还能接受，当然肯定是比不了MLIP的速度hh，主要耗时的还是计算hessian的时候CPHF部分，MPSCF框架针对hessian计算过程中的几个阶段进行了优化，详情可以参考MPSCF链接的相应notice/日志：MPSCF
自己一个人可能对于一些问题考虑的不是特别全面，如果各位感兴趣使用的话，可以给我提一些建议或者issue，我会尽力修改，下面是简单的使用方法：
2.1.自定义模版文件：

1. {
   
2.     "basis":                "def2-svp",
   
3.     "auxbasis":             "def2-svp-jkfit",
   
4.     "xc":                   "tpssh",
   
5. 
   
6.     "device":               "auto",
   
7.     "cuda_visible_devices": "0,1,2,3",
   
8. 
   
9.     "max_cycle":            300,
   
10.     "conv_tol":             1e-8,
    
11.     "grid_level":           5,
    
12.     "damp":                 0,
    
13.     "level_shift":          0.4,
    
14. 
    
15.     "maxsteps":             1,
    
16. 
    
17.     "hessian":              true,
    
18.     "hessian_device":       "gpu",
    
19.     "hessian_patch":        true,
    
20.     "cphf_grid_level":      null,
    
21. 
    
22.     "output_dir":           "outputs",
    
23. 
    
24.     "disp":                 "d3bj",
    
25. 
    
26.     "verbose":              4,
    
27. 
    
28.     "write_fchk":           true,
    
29. 
    
30.     "resp_a1":              5e-4,
    
31.     "resp_a2":              1e-3,
    
32.     "resp_b":               0.1,
    
33.     "hfree":                true,
    
34. 
    
35.     "chgmod":               1,
    
36. 
    
37.     "scalef":               1.0,
    
38.     "bond_avg":             true,
    
39.     "angle_avg":            true,
    
40. 
    
41.     "paraset":              "12_6",
    
42. 
    
43.     "ligand_metal_bond":    {"L01": true},
    
44.     "ligand_charge":         {"L01": -1},
    
45. 
    
46.     "optimizer":            "ancopt",
    
47.     "optimizer_maxsteps":   1000,
    
48.     "conv_tol_grad":        1e-3
    
49. }
    

复制代码

如果体系中没有配体分子的话，后面的ligand_metal_bond和ligand_charge部分就不需要自定义了，或者直接注释掉就行，optimizer的话对应的是GFN2-XTB预优化所采用的优化器，主要有两个参数可选，一个是“berny”，一个是上述模版中的“ancopt”，其实还可以加上ase的FIRE，后续考虑情况可能会加。其他的像level-shift之类的就不多赘述了，如果是简单的单离子体系直接设置为“0”就可以了，不影响正常收敛，damp几乎用不到设置（就目前自己遇到的案例来说），下面是关于gpu.json文件的模版：

1. {
   
2.   "cphf_subspace_frac": 0.80,
   
3.   "cphf_est_cycles": 15,
   
4. 
   
5.   "_make_h1": "make_h1 / _get_jk_ip device-0 peak control.",
   
6.   "make_h1_pinned_frac": 0.40,
   
7.   "make_h1_avail_div": 4,
   
8.   "make_h1_stream_tile_frac": 0.08,
   
9. 
   
10.   "_partial_hess": "Multi-GPU partial_hess_elec block sizing and load balancing.",
    
11.   "split_min_frac": 0.55,
    
12.   "hk_ao_aux_budget_frac": 0.30,
    
13.   "hk_ip1_ip1_budget_frac": 0.35,
    
14.   "aux_aux_budget_frac": 0.20,
    
15. 
    
16.   "_df": "DF JK-loop block sizing.",
    
17.   "df_blksize_factor": 0.05
    
18. }
    

复制代码

主要可供修改的就是df_blksize_factor,如果在您的显卡上出现了OOM，那么把它适当调低即可，但是不建议这么做，因为还是会增加一定的耗时，目前0.05对于大多数24GB显卡够用了，前提是显卡数量大于等于3，所以其实感觉就性价比来说不算特别高吧，但是最起码能用，个人觉得还是像个可以使用的玩具。
2.2.运行方式以及解释：

1. python main.py input.pdb output_H.pdb \
   
2.     --metal-charge ZN:2 \
   
3.     --step1 --step2 --step3 --step4 --step5 \
   
4.     --json scf_params.json --gpu-json gpu.json

复制代码

总共是5步，可以各自独立运行，没有强制性的先后顺序要求：

steps	job	describe
--step1	build MCPB MODEL	input/geometry
--step2	QM calculation	hessian
--step3	RESP	RESP2
--step4	FRCMOD GENERATION	Seminario
--step5	tleap	leap

在amber官网上的1okl金属蛋白，测试是通过的没有问题，同时欢迎各位老师指教）
下面是框架的链接：
https://github.com/KeithTab/MPSCF


忘记提了一下测试的情况，原先的96原子体系的hessian计算预计要125GB左右，用这个框架做的话差不多峰值显存需求约为40GB左右，感觉有两张或者三张显卡还是能算一算的