单一凡：dflow + 电子结构计算——过程与终点的思考
DeepModeling Hackathon 2.0 于8月31日圆满落下帷幕，最终5位选手在决赛中脱颖而出，充分展示了各自的风采。未来，选手们也将携手DeepModeling社区，用创意和代码为更多人创造价值。
今天，小编为大家带来AI辅助电子结构计算：ABACUS, DeePKS赛道一等奖单一凡的获奖作品展示，一起来看看吧~
01
项目介绍
本Hackathon项目参考DP-GEN2对ABACUS-DeePKS的内嵌工作流进行重写，取得了以下成果：
利用dflow工作流套件实现了DeePKS + ABACUS 的迭代训练
ABACUS的SCF计算在多机器间并行进行
迭代终点的自动判断
如今有很多像DeePMD这样很著名的机器学习势能训练工具，使得用于大规模和长时间的分子模拟的量子力学模型准确性大幅提高。
然而，机器学习势函数的训练首先需要生成大量的第一性原理计算数据，对于精度要求很高的场景（如Quantum Monte-Carlo、杂化密度泛函等），生成这些数据会耗费极大甚至难以负担的计算资源。
常见的解决方法就是“钞能力”
不去Push一下老板怎么发Science？但如果有那么小部分情况，恰巧财力没有那么雄厚，就可以考虑使用DeePKS这样一个非常有潜力的工具。
DeePKS是一种机器学习驱动的密度泛函模型，初始被搭载到PySCF软件上并应用于分子体系，取得了极大成功。这种神奇的泛函模型现在已经可以被应用于周期性体系，即DeePKS + ABACUS。
ABACUS是一个开源的国产密度泛函理论计算工具，由北大的陈默涵老师带队开发。我非常敬佩他十年磨一剑的坚定，这款软件也正在快速崛起。
而DeePKS构建了一座跨越低精度计算开销和高精度计算结果的桥梁，使得我们能更轻松地走上雅各布天梯。
搭建跨越雅各布天梯的桥梁，是通过迭代计算实现的
首先利用ABACUS的SCF计算去求解一个低精度的结果作为初始化，之后利用DeePKS的机器学习方法去拟合这样低精度结果与高精度结果之间的差距，生成一个初始模型。
有了初始模型后，就把它作为修正项，直接去修正ABACUS的输入。经过低精度的SCF计算后得到收敛结果，那么这样的结果再通过DeePKS去与高精度结果对比来缩小差距，生成第二个模型。
之后这第二个模型又会被输入ABACUS，继续计算并收集结果。
迭代过程可能是永无止境的，不过终能达到模型拟合的极限，或者使结果收敛到一个可以接受的误差范围。低精度的SCF计算由于修正模型的加持，在低开销的情况下，可以输出一个高精度的DFT结果。
化简为繁，实现弯道超车
这样的收敛过程比较形象地说就是像图里那个山谷一样，把三维空间内寻找最低点的操作，拆解成两个二维竖直平面内的收敛过程，虽然多走了一些路，但是开销更低。
通过这段时间的学习，一步一步把DeePKS和DP-GEN2的源码都啃完，设计出了这样一个比较复杂的dflow代码结构。
三层“套娃”结构
从最里面开始看，把abacus_scf和deepks_train步骤都拆成了两个op去进行执行，用steps进行打包。之后在abacus_scf和deepks_train之间加上数据转换的convert_scf，再在他们外层进行steps打包。
因为初始化的参数init_iter_config和后续迭代的参数iter_config是不一样的，那就需要再把初始化的步骤和后续迭代步骤的外层进行第三次steps打包。
最后就形成了这样一个三层“套娃”的结构。通过对iter_block这样一个steps，我们称为Super-OP，进行多次复用，就组成了初始化以及后续的多次迭代的过程。
而iter_block间传递的就是训练所得到的模型，其他数据均在内部进行传递。以上这些都可以在argo页面或者Bohrium平台页面上一键下载。
过程：自动分发并行计算
采用一套独特的命名规则，对每一个结构的SCF计算文件夹进行命名，按需打包好后即可送入dflow内置的Slices功能进行分发。之后再解析文件夹名，按原始的存储结构对回传的结果进行排布。最终实现了SCF多机器并行过程的无感知分发和收集。
run_scf_abacus = Step( 'run-scf-abacus', template=PythonOPTemplate( run_scf_abacus_op, slices = Slices( "{{item}}", input_artifact = ["task_path"], output_artifact = ["task_path"], ), python_packages = upload_python_package, **run_template_config, ), parameters={ "scf_abacus_config" : scf_abacus_steps.inputs.parameters['scf_abacus_config'], "n_iter" : scf_abacus_steps.inputs.parameters["n_iter"], }, artifacts={ 'task_path' : prep_scf_abacus.outputs.artifacts['task_paths'], "stru_file" : scf_abacus_steps.inputs.artifacts["stru_file"], "model" : scf_abacus_steps.inputs.artifacts["model"], }, with_param=argo_range(argo_len(prep_scf_abacus.outputs.parameters["task_names"])), key = step_keys['run-scf-abacus'], executor = run_executor, util_image="base_dflow_deepks", **run_config, )终点：自动判断迭代终点
我在初赛中对迭代过程加入了波动检测算法，在训练结束时侦测模型训练误差曲线，当曲线没有较大范围波动则认为模型修正达到极限，从而停止迭代过程。
如左图中所示，第一轮迭代(iter.00)中，训练误差曲线呈下降趋势，模型得到了较大程度的修正，说明此轮迭代的输入模型仍不能使SCF的收敛值接近目标值，因此需要执行第二轮迭代。
而右图为第二轮迭代(iter.01)，模型训练误差趋近于一直线，说明输入模型已经接近于其能达到的修正极限，则终止迭代过程。
化繁为简，来自效率的挑战
代码采用了模块化设计，使得后续的功能开发或者新功能迭代的过程中，可以复用之前设计出来的任意一个模块，减少重复修改的开销。
去掉所有细节，最后呈现的工作流其实是非常简单且直观的。从数据输入开始，再到转化，并行SCF计算，训练模型，最后进入下一轮迭代。
简单易用且高效，是软件设计中不变的追求。
受限于比赛时间和个人能力，本次Hackathon项目提交的代码仅是能满足运行的最低要求，仍然有诸多不足之处亟待优化。
这次比赛陪伴了我一整个夏天，我在这里学到了很多，收获了很多，也掉了很多头发，认识到了自己的不足。期待能为AI4S尽自己的一份力量！
最后感谢所有支持和帮助我的老师同学们！
02
导师评价
欧琪
本科毕业于中国科学技术大学，博士就读于美国宾夕法尼亚大学，导师Joe Subotnik。随后在普林斯顿(合作导师Emily Carter)和清华大学(合作导师帅志刚)分别做过博士后。博士期间研究背景为分子体系非绝热过程的方法论开发，参与了量化软件Q-Chem中TDDFT电子态之间旋轨耦合及非绝热耦合的实现。博士后期间参与开发了potential functional embedding theory，以及有机光电材料的光物理性质计算方法。于2021年10月正式加入AISI，任职量子化学研究员岗位，主要职务为ABACUS软件及机器学习辅助泛函DeePKS的开发及完善。
作为硕士一年级的学生，单一凡同学具有卓越的代码编写能力和软件工程基础。ABACUS-DeePKS本身的训练流程较为复杂，需要迭代进行ABACUS的SCF计算以及DeePKS-kit的模型训练步骤，直至模型收敛。此前，该流程需要较多的人工干预或依赖使用者自己编写的脚本完成，自动化程度较差。初赛中一凡快速掌握了ABACUS-DeePKS的迭代训练流程，并在dflow中实现了自动化判断模型是否收敛的功能，相比于此前肉眼比较两轮训练的能量差值，一凡提出的自动化判断方法足够鲁棒且更为高效。在决赛过程中，一凡以DP-GEN2的框架为基础，成功实现了ABACUS-DeePKS的全自动化dflow工作流，整个流程包括了三层嵌套，并引入了Super-OP的设计，能够实现多次复用，包括初始化和后续的多次迭代。目前工作流已经具备基本功能，跑通单一体系的训练任务。一凡也将参与到后续工作流完善的工作中，使得该工作流具有更高的普适性及灵活性，如实现混合体系训练以及自动评估模型训练是否理想等。该自动化工作流能够极大推动DeePKS项目本身的发展，特别是大规模的应用以及普适模型的训练等方面。
- End -
（如需转载图文请与公众号后台联系）
-------------------------------
推荐阅读