magattacaのブログ

日付以外誤報

OpenMMをステップバイステップで 〜Part 5 シミュレーションプログラムの中身を理解しよう〜

OpenMMの使い方を順番にたどる記事のPart 5です。前回までで、GUIOpenMM Setupを利用して作成したファイル(①タンパク質構造、②シミュレーション実行プログラム)を用いて、Google Colab上でシミュレーションが実行できることを確認できました。

今回は「シミュレーションプログラムの中身がどのようになっているか?」、 OpenMMのコードの書き方について眺めて見たいと思います。

また、コードの解釈に先立って「OpenMM プログラムの全体像(各クラス間の相関関係)」を把握したいと思います。

1. プログラムの全体像(相関図)

1-1. Diagram of classes in OpenMM 6.0

OpenMMのプログラムの全体像について参考になる図がありました。バージョン 6.0*1についてのものですが、クラス構成は以下のようになっているそうです。

OpenMMのクラス構成

上図はSimTKプロジェクトにあるOpenMMDocumentsで公開されているOpenMM Class Diagramです。*2

各クラスの図形の意味は左下の点線枠内にあります。

OpenMMはPythonスクリプトで実行できるようになっていますが、内部の仕組みはC++で書かれているものもあります。これらは「API Layer Class」として緑色ひし形でダイアグラムに示されています。SystemContextStateといったクラスです(OpenMM C++ API)。

一方で、Pythonで書かれているものは「App Layer Class」 として黄色円形で示されています。SimulationReportersなどです(OpenMM Python API)。

また、「App Layer Class」のうちファイルを読み込んで働くものは灰色円形となっています。PDBFileや他のMDソフトで作成した入力(AmberPrmtopFIle etc.)などです。

さらに、各クラスに紐づいたメソッドや生成されるデータは、それぞれ赤枠カード黄枠カードになっています。

これらのクラスの関係性を矢印でつないでいるのが上図のダイアグラムです。*3

1-2. ちょっと簡略化

ダイアグラムがかなり盛りだくさんなので少し簡略化してみます。

メインになりそうな部分を白抜き表示にしています。それぞれこういう役割のクラスのようです。

ドキュメントに基づくクラスの説明

説明文だけだとよくわかりません・・・。

2. OpenMMのクラスとMD計算上の役割?

2-1. イメージ

文章だとクラスの役割がわかりにくいので、MDシミュレーションの実行工程とムリやり対応づけてみましょう。

まず、MD計算の流れはこんな感じでした。

これをふまえて、クラスと対応づけるとこんな感じでしょうか???

定義した計算System)でシミュレーションSimulation)を行って、得られた系の時間発展(状態の"文脈", Context)から知りたい時点の情報State)を適宜取り出すイメージです。

2-2. くどく説明

もう少しくどく書きます。

① まず、シミュレーションを行う系の定義をSystemクラスで行います。「どんなタンパク質?水分子はいくつ?どんな箱に置くの?力場は?…」みたいな基本的な条件設定のイメージです。

② ついで、計算の実行に関する設定項目(「時間発展の積分計算はどうするか?」)や出力方法(「結果のデータやトラジェクトリはどうやって出力する?」)に関する設定を、先のSystem)と結びつけるためにSimulationクラスを利用します。

③ このSimulationを使ってエネルギー最小化minimizeEnergy())やトラジェクトリの生成(プロダクション、step())を行うことになります。

④ また、MD計算は時間変化を追うシミュレーションなので、系の状態のデータを格納するためにContextクラスが必要になります。例えばシミュレーション開始時点ではタンパク質の各原子の位置情報(pdb.positions)を設定したりします。

⑤ その上で、特定の時点での状態の情報を取り出すためには別途Stateクラスを利用することになります。シミュレーションの要である情報を格納したContext変数には直接アクセスできないStateを経由してアクセスする)仕組みになっているようです。

⑥ 最後に、シミュレーション全体の時間変化やデータはreporterを使って出力することになるようです。

3. サンプルコードを眺めよう!

OpenMMの概観がわかったところで、OpenMM HPにあるサンプルコードを例にPythonの書き方を眺めてみましょう!

3-1. パラメータの復習

コードをみる前に、過去記事でとりあげたOpenMM SetupGUIツール)で設定したパラメータのイメージを貼っておきます。

設定項目のイメージ

たくさんあって大変です。。。

3-2. サンプルコード

サンプルコードではどうなっているでしょうか?

from openmm.app import *
from openmm import *
from openmm.unit import *
from sys import stdout

pdb = PDBFile('input.pdb')
forcefield = ForceField('amber99sb.xml', 'tip3p.xml')
system = forcefield.createSystem(pdb.topology, nonbondedMethod=PME, nonbondedCutoff=1*nanometer, constraints=HBonds)
integrator = LangevinIntegrator(300*kelvin, 1/picosecond, 0.002*picoseconds)
simulation = Simulation(pdb.topology, system, integrator)
simulation.context.setPositions(pdb.positions)
simulation.minimizeEnergy()
simulation.reporters.append(PDBReporter('output.pdb', 1000))
simulation.reporters.append(StateDataReporter(stdout, 1000, step=True, potentialEnergy=True, temperature=True))
simulation.step(10000)

(設定項目の多さから考えれば)10数行でMD計算の実行ができるというシンプルさがOpenMMの良さですね!

3-3. コードの中身

3-3-1. ステップ分割

ではコードの中身を見てみましょう。import部分は省略して、残りは以下のように分割できそうです。

サンプルコードを工程で分割

青い点線で「⓪系の準備、①エネルギー最小化、②シミュレーション本番」の3つのステップに区切っています。ここでは「平衡化」に相当する操作がないことにご留意ください。

各コードの右側には「今まで見てきたOpenMMのどのクラスに相当するか?」を書いています。

SystemSimulation、そしてContextと順番に構築して「⓪ 系の準備」を行います。ついで「② エネルギー最小化(minimizeEnergy())」を経て「③ シミュレーション本番(step(10000))」を行っています。

データの出力設定(reporters)は「保存したいステップ(③ シミュレーション本番)の直前」に行うというのもポイントになりそうです。

なお、サンプルコードでは状態の確認などを行っていないのでStateクラスが出てこないことにご留意ください。

3-3-2. 系の準備-1

各ステップを細かく見ていきましょう。まずは系の準備です。

最初に、Systemオブジェクトをタンパク質やその他粒子のトポロジーに基づいてForceFieldから構築します。

まず、(水分子の配置などの前処理済みの)PDBファイルの読み込みPDBFile())と、力場のパラメータを定義したXMLファイルの読み込みForceField())を行います。

ここで、ForceFieldは水モデルの力場も合わせて読み込んでいることにご注意ください(amber99sb.xmltip3p.xml)*4

読み込んだらトポロジーpdb.topology)と、(相互作用)の計算方法を引数にとってcreateSystem()することでSystemオブジェクトが構築できます。

なお、力の計算の設定は非結合相互作用(粒子・メッシュ・エバルト法、nonbondedMethod=PME)やカットオフ距離拘束条件などです。

3-3-3. 系の準備-2

次に、Simulationオブジェクトを構築します。これはtopologysystemintegratorを引数にとることで構築できます。

イラストはいらすとやさんより

Systemの構築に使ったtopologyがもう一度出てくるのが少しわかりづらいです。

積分計算方法(Integrator)としてサンプルではLangevinIntegratorを利用しています。これはランジュバン動力学を利用するもので、温度制御ができます。

LangevinIntegratorの引数は(temperature, frictionCoeff, stepSize)で、サンプルでは熱浴の温度を表す引数に「300*kelvin」が設定されています。

なお、nanometerkelvinといった単位が利用できるのは最初に以下を実行しているからです。

from openmm.unit import *

単位系が使えるのは便利なので忘れずに実行したいですね。

3-3-4. 系の準備-3

系の準備の最後はsimulationからのcontextの構築です。

位置についてContextを作成(イラストはいらすとやさんより)

ここではPDBファイルの情報のうち「位置 pdb.position」を利用します。

以上で系の準備は終わりです。

3-3-5. エネルギー最小化

ついで、系のエネルギーを最小化します。

minimzeEnergy()simluationに行うだけです。シンプル!

3-3-6. シミュレーション本番

最後にシミュレーション本番です。トラジェクトリのプロダクションです。

Reporterを追加して計算のstepを踏んでいく(イラストはいらすとやさんより)

step()により、指定したタイムステップのシミュレーションが実行できます。

ただし、ここではトラジェクトリを出力したいので、時間発展させる前に種々のreporterオブジェクトをreporters.append()によりsimulationに追加しています。

出力内容としてPDBReporter()は指定したファイルパスに、一連のフレームを(一定間隔ごとに)PDBファイルとして書き出します。

また、StateDataReporter()はエネルギーや温度といった状態の情報を書き出します。

それぞれの細かな設定はマニュアルをご参照ください。

以上がサンプルコードの中身でした。

4. OpenMM Setupで作成したコードを眺めよう!

4-1. サンプルコードとの違いの概略

サンプルコードを見ることで、シミュレーション実行のためのコードの流れがつかめました。

もう少し長いコード例として、過去記事OpenMM Setupで作成したものを眺めてみましょう!

コード全体の流れは、予めオブジェクトの作成やメソッドのための色々な引数変数に代入する形で定義しておき、後ほどシミュレーション実行コードを書くという順番です。

また、サンプルコードのシミュレーションとの大きな違いとして以下の点が挙げられます。

  • ① NPTアンサンブルで圧力制御も行う
  • ② シミュレーション本番前に平衡化のステップも行う

細かな違いとしては以下の点が挙げられます。

  • ③ Reporterの種類の違い
  • ④ 計算機環境(platform)の設定を明示的に行っている

4-2. コード全体を眺めよう

違いをふまえた上でコード全体を眺めましょう!

コードの途中にコメントアウトする形で説明を追記しました。

# 必要なライブラリのインポート

from openmm import *
from openmm.app import *
from openmm.unit import *

# ファイルの読み込み(PDBと力場XML → Systemのベースになるファイル)

pdb = PDBFile('3poz-processed.pdb')
forcefield = ForceField('amber14-all.xml', 'amber14/tip3pfb.xml')

# Systemの構築に使うパラメータ変数の設定  
## 非結合相互作用の計算方法や拘束条件の定義

nonbondedMethod = PME
nonbondedCutoff = 1.0*nanometers
ewaldErrorTolerance = 0.0005
constraints = HBonds
rigidWater = True
constraintTolerance = 0.000001

# 積分計算(Integrator)の設定に必要なパラメータ変数の設定  
## ここではLangevinMiddleIntegrator
dt = 0.002*picoseconds
temperature = 300*kelvin
friction = 1.0/picosecond

# アンサンブルNPTとしているため圧力制御も行う(サンプルコードとの違い)
## 圧力制御のパラメータ変数を設定
pressure = 1.0*atmospheres
barostatInterval = 25

# シミュレーションを行うための設定
## シミュレーション時間の設定  
## 平衡化も行うため別途タイムステップを定義している(サンプルコードとの違い)
steps = 500000
equilibrationSteps = 50000

## 計算環境をPlatformで設定している(サンプルコードとの違い)
platform = Platform.getPlatformByName('CUDA')
platformProperties = {'Precision': 'single'}

## 予めreporterとして何を使うかも変数に定義している
dcdReporter = DCDReporter('trajectory.dcd', 10000)
dataReporter = StateDataReporter('log.txt', 1000, totalSteps=steps,
    step=True, speed=True, progress=True, potentialEnergy=True, temperature=True, separator='\t')
checkpointReporter = CheckpointReporter('checkpoint.chk', 10000)

# シミュレーションの準備  
## ForceFiledからSystemを構築する
print('Building system...')
topology = pdb.topology
positions = pdb.positions
system = forcefield.createSystem(topology, nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,
    constraints=constraints, rigidWater=rigidWater, ewaldErrorTolerance=ewaldErrorTolerance)

## 圧力制御の機能をaddForceする(サンプルコードとの違い)
system.addForce(MonteCarloBarostat(pressure, temperature, barostatInterval))

## 積分計算の設定
integrator = LangevinMiddleIntegrator(temperature, friction, dt)
integrator.setConstraintTolerance(constraintTolerance)

## simulationの構築
simulation = Simulation(topology, system, integrator, platform, platformProperties)
simulation.context.setPositions(positions)

# XML結果出力ための設定
with open("system.xml", mode="w") as file:
    file.write(XmlSerializer.serialize(system))
with open("integrator.xml", mode="w") as file:
    file.write(XmlSerializer.serialize(integrator))

# エネルギー最小化の実行
print('Performing energy minimization...')
simulation.minimizeEnergy()

# 平衡化のステップ(サンプルコードとの違い)
print('Equilibrating...')
simulation.context.setVelocitiesToTemperature(temperature)
simulation.step(equilibrationSteps)

# シミュレーション本番の実行の前にreporterオブジェクトを追加する

print('Simulating...')
simulation.reporters.append(dcdReporter)
simulation.reporters.append(dataReporter)
simulation.reporters.append(checkpointReporter)
simulation.currentStep = 0

# シミュレーション本番の実行
simulation.step(steps)

# シミュレーション最後の状態をXMLで書き出す
simulation.saveState("final_state.xml")

ずいぶん長くなっているように見えますが、実際の操作としてはサンプルコードとそれほど変わっていなさそうです。

4-3. コードの違いを少し詳しく

折角なのでサンプルコードと違う部分を少し詳しくみてみましょう。

4-3-1. NPTアンサンブル(圧力制御)

まず、アンサンブルの違いです。

NPTアンサンブル達成のためMonteCarloBarostatという圧力制御を系に加えています。これは、Systemを構築後にaddForce()の形で追加することで利用可能です。

圧力制御の方法(イラストはいらすとやさんより)

温度制御に関わるIntegratorSimulationを構築する際の引数であったのと異なり、その前の「系の構築段階で力として加える」という点は要注意です。

4-3-2. プラットフォーム

続いて計算機環境の設定です。Platformクラスを利用して設定しています。

計算機プラットフォームの設定(イラストはいらすとやさんより)

ここではGoogle ColabのGPUCUDA)を利用するためgetPlatformByName('CUDA')となっています。

Platformはそのプロパティ{'Precision': 'single'})と合わせてSimulationを構築する時に引数として与えられています。

Simulationが計算の実行や出力のさまざまなオブジェクトを取りまとめる役割だったことを考えると納得できそうです。

4-3-3. 平衡化

最後に平衡化ステップを加えることによる違いです。

平衡化の計算自体はシミュレーション本番と同様に、simluationに対して必要な時間に相当するstep()を行えばOKです。

初速分布を与えて平衡化(イラストはいらすとやさんより)

ただし上図のように、OpenMM Setupで作成したコードではさらにsetVelocitiesToTemerature()というものをcontextに与えています。

これは「系の全粒子に対して、設定の温度ボルツマン分布に従うようにランダムな速度を与える」という役割を担います。エネルギー最小化後の系に初速を与えたと考えればよさそうです*5

contextが速度を含む系の状態を保持するクラスだったことをふまえれば、このような追加方法になるのも納得できそうです。

また、平衡化のstep()の計算が終わったのち、シミュレーション本番のstep()に入る前に以下のコードを実行しています。

simulation.currentStep = 0

平衡化の状態は保持しつつ「ステップのインデックスcurrentStep)をゼロ」にすることで、シミュレーション本番のトラジェクトリ出力(reporter)に合わせていると考えられそうです。

以上がコードにおける主な違いでした。

5. クラスの振り返り

サンプルコード、およびOpenMM Setupで作成したコードの振り返りが終わりました。

記事全体のまとめも兼ねてコードにでてきたOpenMMクラスのダイアグラムを作ってみましょう。

こんな感じでしょうか?

(すこしアップデートした)簡略化クラスダイアグラム

一番最初に引用したOpenMM Class Diagramの半分くらいは抑えられたかな?、という印象です。

設定項目とクラスの種類が多いので関係性がわかるとコードを書きやすくなりそうですね。

6. (おまけ)OpenMMのアーキテクチャ

最後に脱線ですが、OpenMMの構成(アーキテクチャ)の図を引用しておきます*6

OpenMMのアーキテクチャ

OpenMMは計算を実行するためのコードの部分では、計算機の環境(プラットフォーム)を意識しなくて良いように設計されているそうです。

利用者はアーキテクチャ上部のPublic Interface(Public API)の使い方を把握しておけばよく、これはプラットフォームと独立なコードになっています。このPublic APIが上で見たSystemContextなどです*7

プラットフォームに依存した部分は、アーキテクチャー下部のPlatform Abstraction Layer(Low Level API)やComputational Kernelsです。最下層が実際に計算を実行するハードウェアに関する部分だそうです。

独自の計算手法を開発する場合などは下層の理解が必要となるそうですが、そもそもそんな上級者はこんなブログ読みませんよね。。。

以上のように、プラットフォーム依存の部分独立の部分分かれた構成となっているので、(貧弱な)自分のパソコン上でもGoogle ColabのパワフルなGPUを使った計算のためのプログラム作成が簡単にできたのでしょうか???

よくわかりませんがアーキテクチャって言ってみたかっただけです。ごめんなさい。

7. おわりに

以上、今回はOpenMMでシミュレーションを実行するためのコードの書き方について、「① (OpenMM HPの)サンプルコード」と「② OpenMM Setupで作成したコード」を例に眺めてみました。

また、コードの中身を眺めるに先立って、OpenMM Class Diagramを利用して各クラス間の相関関係を大まかに把握し、OpenMMプログラム全体の概観の理解を試みました。

OpenMMは(ばりばりの計算の専門家でなくても)シミュレーションを実行できるように、Pythonで簡単にシミュレーションプログラムをかけるよう設計されているそうです。

実際に、MD計算に必要な設定項目の多さ多段階工程のややこしさからすれば、シンプルなコードを並べるだけで計算を実行できるようになっていました。

・・・とはいうものの、やはりある程度は複雑になってしまいます。

今回の記事で見たように、クラスダイアグラムのような形式で各操作の関係性がわかっていると動作するコードをちゃんとかけるようになるかもしれませんね?

以上、今回も間違いが多そうです。ご指摘いただければ幸いです。

ではでは!

*1: Class Diagram資料の公開日付は2014年のようです。2022年現在はOpenMM 7.7が最新版となっています。

*2:SimTKは医学生物学計算コミュニティのための無料のプロジェクトホスティングプラットホームだそうです(Whit is SimTK?)。

*3:Lee-Ping Wang博士によるワークショップ(YouTube)「Creating and Customizing Force Fields in OpenMM」での説明などを参考にしています。

*4:利用可能な力場のXMLファイルについてはOpenMM User Guide : 3.6.2. Force Fieldsなどを参照してください。

*5:サンプルコードで初速を設定していない理由はよくわからないです…。NPTアンサンブルのMonteCarloBarostatにも温度がある関係?

*6:OpenMM User Guide: 7.4. Architectural Overview

*7:OpenMM User Guide: 7.5. The OpenMM Public API