はじめに
SDFをサクっとCSVに変換するの逆バージョン。
なぜこんなことをしたいかというと、SMILESが含まれたCSVの化合物データに対し、SDF形式でしか表現できない情報(3次元座標や原子プロパティ等)を付加したいため。
仕様
- CSVの中のSMILESプロパティと、プロパティ項目を指定し、SDFとして出力
- 化学構造はSMILESを元に生成
- プロパティ項目は1つだけ指定可能(とりあえず機械学習の目的変数を1つ入れたいだけなので)
ソース
csv2sdf.py
import pandas as pd
from rdkit import Chem
import argparse
def main():
parser = argparse.ArgumentParser()
parser.add_argument("-input", type=str, required=True)
parser.add_argument("-output", type=str, required=True)
parser.add_argument("-smiles_col", type=str, default="SMILES")
parser.add_argument("-target_col", type=str, required=True)
args = parser.parse_args()
# CSVの読み込み
print(args.input)
df = pd.read_csv(args.input)
# SMILES列を取り出し
smiles_list = df[args.smiles_col].values
# target列の取得
targets = df[args.target_col].values
mols = []
for i, (smiles, target) in enumerate(zip(smiles_list, targets)):
# SMILESからmolを生成
mol = Chem.MolFromSmiles(smiles)
# Targetの値を設定
mol.SetProp(args.target_col, str(target))
mols.append(mol)
# SDFへの出力
sdf_wtr = Chem.SDWriter(args.output)
for mol in mols:
sdf_wtr.write(mol)
sdf_wtr.close()
解説
特になし。
3次元化したいとか、原子プロパティを埋め込みたい等したい場合は、ソース中のRDKitのmolオブジェクトに色々操作を加えてみてください。