今天介紹一篇文章，使用神經網路將分子動力學第一原理計算推進到新的尺度和精度。分子動力學模擬一般被用來揭示化學反應的微觀物理過程，用於光合作用，新藥尋找，奈米材料性質**等等。
劉慈欣告訴我們，經典力學下的三體運動已經非常複雜，而量子力學統治下的多體相互作用，計算複雜度更是高到天際。分子動力學模擬要根據每乙個時刻千萬粒子的位置計算勢能面（勢能的負梯度等於粒子受力），再根據受力和運動方程，更新所有粒子的座標與速度。此過程即便使用超級計算機，也只能計算很小的系統，遠離物理現實。
文中對比了各種大規模分子動力學模擬的原子數 （atoms), cpu計算核心數，gpu 計算核心數，以及單個原子單步模擬花費時間 tts (秒/步/原子，或 s/step/atom）。
對比來看，2023年之前模擬的系統最大包含 100 萬個矽原子（速度= [公式] s/step/atom)，模擬速度最快達到 [公式] s/step/atom (9000個水原子系統）。
今天介紹的這篇文章，模擬了4.03億個水分子h2o，1.13億個銅原子cu。速度也比之前最快的模擬提高了5000倍，為 [公式] s/step/atom， (下表的最後兩行）。
作者們詳細介紹了如何對開源** deepmd-kit 二次開發，各種優化，加速計算。這些優化包括超級計算機集群可擴充套件性優化， 面向 gpu 的資料結構優化，tensorflow 運算元優化與單雙精度混合計算加速優化。這篇部落格也集中於介紹這些優化方法（相比其他文章，deep potential 的網路架構極其簡單）。
從文章的配圖 2 也能看出，子圖（a）和（b）是優化之前的 deep potential 模型，而子圖 （c) - (g) 都在介紹如何面向gpu，更改資料結構，計算單雙精度以及 tensorflow 運算元，優化和加速計算。
子圖（b): 定義了一些類似於資料增廣的對稱性運算元，即在平移，旋轉，重排等對稱性變換下使得系統勢能不變。這提供了乙個從資料中學習物理對稱性的好方法。
子圖（c，d): 面向 gpu 的第乙個優化。gpu 計算的底層很有特點，以並行程式語言opencl和cuda 的混合術語舉例，一塊 gpu 一般有幾千個微處理器（processing element), 它們被分成幾十個小組，每個小組稱作乙個computing unit (cu), cu 裡面的32或64個執行緒運算過程同步，被稱為乙個 warp。計算的時候，如果同乙個warp裡不同執行緒走了不同的 if else 通道，那麼此 warp 內的所有執行緒都會既走 if 通道，又走 else 通道。warp 就像儀仗隊的一行士兵，佇列內所有士兵步調一致，不搞特殊。
第乙個優化是減少分支，讓同乙個warp內的所有執行緒進行同樣的操作。做法是將乙個粒子的鄰居們（也是粒子）按種類重排，這對應 子圖 c 和 d，重排後同一種類的鄰居粒子被同一warp內的執行緒操作，步調完全一致，摒除分支。
子圖 (e): 相比 array of structures (aos), gpu 計算更喜歡 structure of arrays (soa)資料結構，這樣gpu從視訊記憶體（global memory） 讀資料時速度更快。子圖 (e) 中對那些資料 reshape 來 reshape 去就是為了將 aos 轉化為 soa，讀到每個執行緒的 private memory 之後再轉化回來。
文中介紹了他們使用的另乙個技巧，資料壓縮。將粒子種類 [公式] ，粒子距離 [公式] 和粒子 index = j 壓縮後儲存在64位整數中。
基本上，這篇文章進行了非常精細的 gpu 並行加速優化，將所有計算耗時或者傳輸資料耗時的運算都移植到了 gpu 上進行。從 cpu 記憶體往 gpu 視訊記憶體移動資料非常緩慢，直接在 gpu視訊記憶體上開闢資料空間，在gpu上計算，減少記憶體-視訊記憶體之間的資料流動，是重要的gpu平行計算優化步驟。
當然，從 gpu 的 global memory（視訊記憶體）往執行緒的 private memory （暫存器）傳輸資料也有延遲，資料的 soa 儲存格式能夠使用快取加速讀取，對多執行緒共享的資料使用每個 computing unit 共享的記憶體 (shared memory)也能減小讀取延遲。目前未見文章在這一塊的優化。
子圖 (f): gpu上單精度(float)計算速度快於雙精度(double)計算速度。子圖 (f) 通過單雙精度混合計算，既保證精度，又加快運算速度。在一般的計算程式中，可多次嘗試，觀察哪些引數（資料）從雙精度改為單精度浮點數後不影響計算結果。此研究發現對原子的位置和能量使用雙精度，對神經網路的引數使用單精度，可以實現混合精度優化。混合精度相比雙精度加速1.5倍 （此處存疑，個人經驗gpu上單精度計算要遠快於雙精度），視訊記憶體節約 50%。
子圖 (g): 這一部分涉及對 tensorflow 計算圖和運算元的優化。此優化只針對當前物理問題。在勢能面的計算中，需要對大量粒子的貢獻求和，涉及很高很瘦的矩陣乘積 matmul 與求和 sum。在此過程中，sum 操作的計算需求劇增，tensorflow 未對此做特殊優化。本文使用了 cublas 提供的 general matrix multiply(gemm) 來加速 x*w + b，concat 和 sum 的計算。另外， [公式] 的微分等於 [公式] , 作者們定製了兩個 kernel 來計算 tanh 和其梯度 tanhgrad。感覺回到了自己手寫一大堆 opencl kernels 解流體力學方程的日子。
子圖之外，文章還對不同計算節點的 mpi 通訊進行了優化。mpi 節點間通訊主要有兩個目的：（1）將系統分塊後，每個節點計算乙個區塊內粒子雲的運動，實現粗粒度並行。每個區塊周邊要有一定的輔助區域作為邊界條件，輔助區域的粒子需要使用mpi通訊從周邊格仔抓取資訊來更新。 （2）輸出時需要對所有區塊的能量求和，mpi_allreduce 約化操作會因不同格仔的阻塞同步帶來延遲，文章使用非同步 reduce，並將求和頻率更改為 20 步一次。
Gromacs分子動力學模擬流程概述
gromacs分子動力學模擬主要可以分為以下幾個步驟，不同的體系步驟可能略有不同。在開始之前，先簡單了解一下預平衡 分子動力學模擬的最終目的是對體系進行抽樣，然後計算體系的能量，各種化學鍵，成分分析等等。打個比方說，我們有乙個蛋白質，我們想將它放入一種溶液中 可能是水，也可能不是 然後看看這個體系的...
分子動力學模擬 VMD多分子處理及指令碼（2）
file new molecule browse name.pdb load 將工作目錄指定到正確的路徑 cd 目錄執行模擬 mol new name.pdb和匯入晶體結構資訊一樣，但是檔案字尾為psf 在main視窗選中psf 然後右鍵，選擇load data to the molecule 出現...
LAMMPS分子動力學技術與應用
北京軟研國際資訊科技研究院 互動派 北京 教育科技 二零二零年九月一日 二零二零年九月一日 四 培訓課表 lammps分子動力學模擬技術與應用 專題線上培訓 課 程 內 容 第一天 上午 lammps基礎 1 分子動力學模擬入門理論 掌握lammps的in檔案中各命令的意義 1.1系綜理論 1.2主...