RDKitでrdMolDraw2Dで生成した構造式画像をエクセルに出力する #RDKit

背景

RDKitのPandasToolsを使えば構造式画像をエクセルに出力することができる。しかしこの場合Drawモジュールという見た目がイケてない方により生成された画像が出力される。今回PandasToolsをハックしてrdMolDraw2Dというイケてる方でもエクセルに構造式画像を出力する方法を調べたためメモっておく。

環境

Windows10
Python3.6
RDKit 2018/9/2

方法

まず、必要モジュールのインポート。足りないものは事前にインストールしておく

from rdkit import Chem
from rdkit.Chem import Draw, AllChem
from rdkit.Chem.Draw import rdMolDraw2D
from IPython.display import SVG
from rdkit.Chem import PandasTools
import pandas as pd
import numpy as np
from io import BytesIO
import copy
import xlsxwriter
from PIL import Image, ImageFilter
from cairosvg import svg2png
import argparse

次にrdMolDraw2Dによる構造式画像をByteIO形式で生成する関数を用意する。

def generate_image(mol, size):

    image_data = BytesIO()
    view = rdMolDraw2D.MolDraw2DSVG(size[0], size[1])
    tm = rdMolDraw2D.PrepareMolForDrawing(mol)

    view.DrawMolecule(tm)
    view.FinishDrawing()
    svg = view.GetDrawingText()
    SVG(svg.replace('svg:', ''))
    svg2png(bytestring=svg, write_to='tmp/output.png')
    img = Image.open('tmp/output.png')
    img.save(image_data, format='PNG')

    return image_data

次に上の関数を用いてエクセルに出力する関数を用意する。PandasToolsの修正版である

def save_xls_from_frame(frame, out_file, mol_col='MolCol', size=(300, 300)):
    margin = 60

    cols = list(frame.columns)
    cols.remove(mol_col)

    data_types = dict(frame.dtypes)
    workbook = xlsxwriter.Workbook(out_file)  # New workbook
    worksheet = workbook.add_worksheet()  # New work sheet
    worksheet.set_column('A', (size[0] - margin + 20) / 6.)  # column width

    # Write first row with column names
    c2 = 0
    worksheet.write_string(0, c2,  mol_col)
    c2 += 1

    for x in cols:
        worksheet.write_string(0, c2, x)
        c2 += 1

    c = 1
    for _, row in frame.iterrows():

        worksheet.set_row(c, height=size[1] - margin)

        # 画像の生成
        image_data = generate_image(row[mol_col], size)
        worksheet.insert_image(c, 0, "f", {'image_data': image_data})

        c2 = 2
        for x in cols:
            if str(data_types[x]) == "object":
                # string length is limited in xlsx
                worksheet.write_string(c, c2, str(row[x])[:32000])
            elif ('float' in str(data_types[x])) or ('int' in str(data_types[x])):
                if (row[x] != np.nan) or (row[x] != np.inf):
                    worksheet.write_number(c, c2, row[x])
            elif 'datetime' in str(data_types[x]):
                worksheet.write_datetime(c, c2, row[x])
            c2 += 1
        c += 1

    workbook.close()
    image_data.close()

さあ、お膳立ては整った。あとは呼び出すだけ。

names = []
mols = []

for mol in Chem.SDMolSupplier("test.sdf", removeHs=False):
    mols.append(mol)
    names.append(mol.GetProp("_Name"))

df = pd.DataFrame(columns=["MolCol","NAME"])

df["NAME"] = names
df["MolCol"] = mols

save_xls_from_frame(df, "test.xlsx", mol_col="MolCol")

参考

PandasToolsソース