NVIDIA Thrustで分子動力学(1)単原子分子集団のMD -

MDは時々使うけど、自分で実装したことがなかったので、一回くらいは、やっておくべき気がした。

実装
https://gist.github.com/vertexoperator/1d2b4c6b71138d59484665cb524143ac

最初なので、相互作用はLennard-Jonesポテンシャルのみで、周期境界条件、NVE(粒子数、体積、総エネルギー)一定という簡単なやつ。粒子の軌跡を追跡するだけでは詰まらないので、任意時点での温度と圧力を計算できるようにしてある。

普通に、全粒子間のLJ相互作用を全部計算すると、粒子数Nに対して、O(N^2)の計算量が必要になるので、分子動力学では、適当なカットオフ距離を決めて、カットオフ距離内の粒子との相互作用のみ計算することが多い(バッキンガムポテンシャルを使う場合も同様の方法が使えると思うけど、クーロン相互作用を計算する場合は、LJ相互作用より減衰しにくいので、ももう少し工夫がいる)。この場合、計算全領域をカットオフ距離幅で分割することで、計算量をO(N)に減らせる。

実装に難しいところはないけど、生のCUDAカーネルを直接書かないという目標は、一箇所だけ、いい方法が思いつかなくて達成できなかった(sortされた整数の配列がある時に、各整数の開始indexを計算する処理で、scan_by_key→unique_by_key→scan→gatherとかで出来るけど、これだけのことに、4つも関数を呼びたくない)。力の計算も、なんか、thrustの正しい使い方じゃない気もするけど、一応正しく、それなりの性能で動いてる。

時間積分は、velocity Verletだけど、今の所、知られてる時間積分では、どの方法であっても、長時間計算すると、エネルギーが発散するという問題は避けられないらしい。エネルギーを保存する時間積分法が欲しいところではある。

カットオフ距離2nm、タイムステップ1fsで、うちのマシンで適当に性能を測定すると、
倍精度、N=100万:1ステップ20msくらい
倍精度、N=1000万:1ステップ140msくらい
単精度、N=100万:1ステップ10msくらい
単精度、N=1000万:1ステップ70msくらい
という感じ。計算は、常温・約8気圧で、カットオフ距離は2nmと、やや大きめにしたものの、単精度と倍精度で、性能比が、ほぼ2：1なことから、残念な感じが察せられる。

力の計算の演算数を数えると、一粒子あたり
・3個の除算（最適化すれば、3個の乗算)と3個のfloor
・8〜11個の加減乗算が近傍グリッドにいる粒子の個数回
・(25個の加減乗算+1個の除算)がカットオフ距離内にいる粒子の個数回
となる。

100万粒子の計算では、系の一辺の長さを、約167nmに設定しているので、粒子が一様に分布していれば、近傍グリッドにいる粒子の個数は、平均46個。カットオフ距離内にいる粒子は(自分自身を除いて)6個強になる。使ったGPUの倍精度演算性能は420Gflopsと公表されていて、100万粒子に対して、420Gflops出れば、これくらいの演算をやるのに、2msあれば十分というくらい。

近傍グリッドは27個あるので、平均すると、1グリッドあたり、1〜2個の粒子しか入ってない。より計算需要があるだろう超高圧の気体や液体になれば、粒子の密度は、数千倍にもなって、一つのグリッドにいる粒子数が増大するので、多分、性能は向上する。希薄気体であっても、全ペアの計算をするよりは、ずっとマシ(100万粒子に対して、全ペアの計算をすれば、10Tflopsの演算性能があっても、1ステップ1秒以上は必要になるはず)なので、こんなんでもいいだろう。

一番はまったのは、初期位置を、単純に一様分布で生成すると、かなり近い位置に粒子が落ちることが結構あって、そうすると、運動エネルギーが普通でも、総エネルギーが異常に大きくなって困る。こういう希薄な気体でも、こういう問題が起きるのは、意外だった。

もう一つ、本来は初期化でやるべき気がすることとして、系の運動量及び角運動量を0にすることがあるが、やっていない。(初期位置を決めれば)運動量と角運動量が0という条件は、速度に関する6個の線形な拘束条件になるので、この拘束条件が6x3N行列Cで書けるとすると、 $I-C^{T}(CC^{T})^{-1}C$ を掛ければ、運動量と角運動量が0の部分空間に落とせる(射影すると、速度分布は、Maxwell-Boltzmann分布からは、ずれると思われる)。LJポテンシャルは方向に依存しない二体相互作用なので、運動量と角運動量は(原理的には)保存するはずである。

NVE条件で計算すると、温度や圧力は別に保存量じゃないので一定値にはならないが、系が平衡状態にあれば、(相転移点の近傍にいない限りは)ある平均値の周りで変動するものと思われる。日常的な感覚に基づく推測でも、外部とエネルギーのやり取りがない孤立系で、突然、温度や圧力が不安定化することは、あんまりないと思われる。けれど、所与の温度や圧力が、どのようなNVE条件で成立するか事前に知ることは、通常、そんなに簡単じゃない。普通、計算したい物理量は、特定の温度や圧力条件下における値なので、これは、少し不便。

なので、分子動力学でも、温度や圧力を制御する方法が色々と考えられている。これらの方法の多くは、1980年台に提案されたものらしい。Nose-Hoover thermostat, Andersen barostat, Berendsen thermostatなどがよく使われるが、実装は、別に難しくない。

何らかのthermostatを使ったNVT計算を、統計力学に於けるカノニカル・アンサンブルの実現と見なすようなことは理論的根拠がないので、これらのbarostat/thermostatは、所与の温度、圧力に対応する密度とエネルギーを発見するために使われる(この値が間違ってないことは理論的に保証されないと思うので、NVE条件による"production run"で、温度や圧力はチェックすべきである)。

生体高分子の計算をやる場合は、NVE一定条件でやると、タンパク質のフォールディングによって温度が変わることがある(十分多くの水分子を用意すれば、系の温度変化は極めて小さくなるだろうが、計算量が増える)ようなので、production runを、NVE条件でやらずに、上に書いたようなthermostat/barostatを使っているのを見ることがあるかもしれない。

実用的な分子動力学に必要な一般化として、多原子分子の計算もある。(古典)分子動力学では、電子の自由度を無視しているので、多原子分子の扱いは、色々無理があるけど、電子の自由度を考慮するには、量子論的な計算が必要になって計算量が爆発的に増えるので、古典分子動力学は、色んな業界で使われている。

電子の自由度を無視した結果として、分子内の相互作用は、二体相互作用の形ではなく、多体相互作用の形でモデル化されることも、しばしばある。こういう相互作用を許すと、運動量や角運動量の保存則が壊れる気がする(そもそも、原子核だけでは、これらの量は保存しないので問題ないとも言える)けど、計算結果は、それなりに有用なものとして扱われる。分子内の相互作用の計算は、実装としては、そんなに難しいものでもない。

分子間の相互作用は、Lennard-Jonesポテンシャルに加えて、Coulomb力が考慮されることが多い。Coulomb力は、LJポテンシャルより、減衰しにくいので、LJポテンシャルのように単純に小さなカットオフを使って遠くの粒子からの影響を切り捨てるわけにはいかないし、それどころか、Coulombポテンシャルは周期境界条件を考慮して、周期対称性を持つ形にしなければならない。つまり、 $1/r^2$ ではなく、 $\sum_{\mathbf{k} \in L} 1/(\mathbf{r} + \mathbf{k})^2$ の形を取る( $L$ は周期格子)

Coulombポテンシャルを考慮した計算も、O(N^2)で良ければ、定義どおりに実装してみることは、特に難しくない。計算量を減らすアルゴリズムとしては、particle Mesh Ewald法というものが、よく使われて、O(N log N)になる。最近は、fast multipole method(FMM)を使った実装が増えてきていて、こっちは、O(N)になる。

cf)FMMについては、2020年に実装しましたというだけの論文が出てる程度には枯れてない話題
A GPU-Accelerated Fast Multipole Method for GROMACS: Performance and Accuracy
https://www.ncbi.nlm.nih.gov/pmc/articles/PMC7660746/

分子内と分子間の相互作用をモデル化する方法は多分色々提案されているが、特に、水分子については、専用のモデルがいくつも作られている。有名なものとしては、SPC,SPC/E, TIP4P,TIP5Pなどがある。Wikipediaで
水モデル
https://ja.wikipedia.org/wiki/%E6%B0%B4%E3%83%A2%E3%83%87%E3%83%AB
を読めば、大体、どういうものかはわかる。

次は、とりあえず、水の分子動力学計算を実装して、いくつかの物理量を計算して、水モデルによる違いでも見ることにする。