RDKitの計算を並列化する

　はじめましてtky_cowと申します。某内資系製薬企業で研究職をしています。技術メモ的な内容を書き記していけたらなと思いブログを開設してみました。初の記事はタイトルの通りなのですが、非常にニッチな内容ですね笑。まあ気にせず行きましょう。

RDKitの記述子計算

　Chemistry分野では、化合物の特性を表す記述子と呼ばれる各種指標を計算することが良くあります。分子量(MW)、logP、tPSAなどが代表的ですね。RDKitでこれらを計算することができるわけですが、千単位の化合物数を計算すると時間がかかってちょっとストレスを感じることがあります。
　RDKitの記述子計算はデフォルトだと1コアで計算を実行するようです。そこで、マルチコアで実行するオプションでもあるのかなと調べましたが、特になさそうでした。公式ドキュメントでもmultiprocessingなどの他の機能で補って的なことが書かれていたので実装してみます。

multiprocessingで並列化

　multiprocessingはPythonの標準ライブラリーの一種です。使い方に関しては以下の記事を参考にさせていただきました。記事によると、今回のケースではPoolよりProcessの方が正しい気がしますが、そこまで厳密な制御を必要としていないので、実装がシンプルなPoolを使います。
Pythonの並列処理・並行処理をしっかり調べてみた - Qiita
Python並列処理(multiprocessingとJoblib) - Qiita

実際のコード

　ここでは実装の全体を示した後、細かいところを少し解説します。最後にTHE DRUG REPURPOSING HUBからダウンロードした10421化合物のデータセットで比較した結果もお示しします。

インポートと設定

import sys
from datetime import datetime
from multiprocessing import Pool
import multiprocessing as multi
from rdkit import Chem, rdBase
from rdkit.Chem import AllChem, Descriptors
from rdkit.ML.Descriptors import MoleculeDescriptors
rdBase.DisableLog('rdApp.warning')

　最後のrdBase.DisableLog('rdApp.warning')は立体フラグが不明瞭などの理由で結構warningが出たりするので、それをオフにしています。市販化合物などのデータではよく出るwarningです。

記述子計算部分

def calc_desclist( smi ):
	# MW, tPSA, AROM, HBA, HBD, ROTB, LogP, RING
	desclist = ['MolWt', 'TPSA', 'NumAromaticRings', 'NumHAcceptors',\
				 'NumHDonors', 'NumRotatableBonds', 'MolLogP', 'RingCount']
	calculator = MoleculeDescriptors.MolecularDescriptorCalculator( desclist )

	#Calculation
	res = []
	try:
		mol = Chem.MolFromSmiles(smi)
		descs = calculator.CalcDescriptors(mol)
		for num in range(0, len(desclist)):
			if descs[num]:
				res.append(str(descs[num]))
			else:
				res.append("")
		resrow = "\t".join(res)
	except Exception:
		ex, ms, tb = sys.exc_info()
		sys.stderr.write("ERROR : " + str(ms) + "\n")

	return resrow

　よく使う8種類の記述子を計算するように設定したcalculatorを作り、CalcDescriptorsメソッドで計算します。for文以降はファイルに書き出すために文字列に置き換え、1行分をまとめているだけです。

並列処理

def run_calc_multi(inpfile, num_cpus=multi.cpu_count()):
	with open(inpfile, 'r') as inF:
		print("Calculation started")
		counter = 0
		smilist = []
		for l in inF:
			if counter != 0:
				smi = l.rstrip()
				smilist.append(smi)
			counter += 1

		print("Compounds processing :", len(smilist))

		# Multiprocessing
		p = Pool(num_cpus)
		result = p.map(calc_desclist, smilist)
		p.close()

		# Output preparation
		pref = inpfile.split(".")[0]
		with open(pref + '_calc.tsv', 'w') as out:
			out.write("SMILES	MW	TPSA	AROM	HBA	" +\
				"HBD	ROTB	LOGP	RING\n")
			
			for smi, line in zip(smilist, result):
				out.write(smi + "\t" + line + "\n")

　本題の並列化部分です。後々のことを考えてnum_cpus引数で使うCPU数を指定できるようにしておきます。デフォルトでは利用可能なCPUをすべて使います。まず、計算する構造のリストを作成し、実行する関数と共に渡します。並列化する部分はこれだけで、非常に簡単ですね！

並列化の効果

　さて、並列化の効果はどの程度あるのでしょうか。CPU数を4→2→1と振って計算時間を比較してみます。結果は以下のようになりました。

CPU数	4	2	1
計算時間(秒)	6.77	8.14	14.0

1コアでも早いやないかい！例が悪かったですね。。記述子計算はそもそも計算速いようです。アラート構造を判定する部分構造検索などが含まれていると遅くなるのですが、その場合でも同じ並列化は有効ですので効果的に使えると思います。今回のコードやサンプルファイルはこちらにまとめてあります。