触媒反応収率予測:分子指標を用いた機械学習および深層学習モデルの実装

環境構築とデータの概要

本解析では、触媒反応の収率を予測するために、機械学習ライブラリであるscikit-learnと化学情報学ライブラリであるRDKit、および深層学習フレームワークのPyTorchを使用します。まず、必要なライブラリをインストールし、データセットの構造を理解します。

データセットは、反応物(Reactant1, Reactant2)、生成物(Product)、添加剤(Additive)、溶媒(Solvent)のSMILES文字列と、目的変数である収率(Yield)で構成されています。学習データには約23,500件、テストデータには約2,600件が含まれます。IDカラムは予測に寄与しないため、特徴量から除外します。

import pandas as pd
import numpy as np
import pickle
import rdkit
from rdkit import Chem
from rdkit.Chem import rdMolDescriptors
from sklearn.ensemble import RandomForestRegressor
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import re
import time

# RDKitのログを抑制
rdkit.RDLogger.DisableLog('rdApp.*')

特徴量エンジニアリング:分子フィンガープリントの生成

SMILES文字列を機械学習モデルが扱える数値ベクトルへ変換するため、RDKitを用いてMorganフィンガープリント(Circular Fingerprint)を生成します。これは分子の局所構造をビットベクトルとして表現する手法です。

以下では、SMILESリストを受け取り、対応するフィンガープリントの配列を返す処理を実装します。重複する分子がある場合、計算コストを削減するためにキャッシュ機能を利用します。

def generate_morgan_fingerprint(mol, radius=2, n_bits=2048):
    """分子オブジェクトからMorgan Fingerprintを生成する"""
    if mol is None:
        return np.zeros(n_bits)
    fp = rdMolDescriptors.GetMorganFingerprintAsBitVect(mol, radius=radius, nBits=n_bits)
    return np.array(list(fp))

def process_smiles_list(smiles_list):
    """SMILESのリストをフィンガープリントの行列に変換する"""
    unique_smiles = list(set(smiles_list))
    # 重複計算を避けるために辞書に保存
    cache = {}
    print("Processing SMILES to fingerprints...")
    for smi in tqdm(unique_smiles):
        mol = Chem.MolFromSmiles(smi)
        cache[smi] = generate_morgan_fingerprint(mol)
    cache[''] = np.zeros(2048) # 空文字列の処理
    
    # 元のリスト順序に合わせて行列を構築
    return np.array([cache[smi] for smi in smiles_list])

# データの読み込み
train_path = '../dataset/round1_train_data.csv'
test_path = '../dataset/round1_test_data.csv'

df_train = pd.read_csv(train_path)
df_test = pd.read_csv(test_path)

# 特徴量の抽出と結合
def prepare_features(df):
    rct1_fp = process_smiles_list(df['Reactant1'].tolist())
    rct2_fp = process_smiles_list(df['Reactant2'].tolist())
    add_fp = process_smiles_list(df['Additive'].tolist())
    sol_fp = process_smiles_list(df['Solvent'].tolist())
    # 次元1(特徴量方向)に結合
    return np.concatenate([rct1_fp, rct2_fp, add_fp, sol_fp], axis=1)

X_train = prepare_features(df_train)
y_train = df_train['Yield'].values
X_test = prepare_features(df_test)

機械学習モデルによる予測(ランダムフォレスト)

ベースラインモデルとして、ランダムフォレスト回帰を使用します。アンサンブル学習の一種であり、過学習を抑制しつつ非線形な関係を捉えるのに適しています。

# モデルの定義と学習
rf_params = {
    'n_estimators': 100,
    'max_depth': 15,
    'min_samples_split': 5,
    'n_jobs': -1,
    'random_state': 42
}

regressor = RandomForestRegressor(**rf_params)
regressor.fit(X_train, y_train)

# モデルの保存
with open('rf_yield_model.pkl', 'wb') as f:
    pickle.dump(regressor, f)

# 予測の実行
predictions = regressor.predict(X_test)

# 提出用ファイルの作成
def create_submission(predictions, filename='submission.csv'):
    with open(filename, 'w') as f:
        f.write("rxnid,Yield\n")
        for i, pred in enumerate(predictions):
            f.write(f"test{i+1},{pred:.4f}\n")

create_submission(predictions, 'rf_submission.csv')

ディープラーニングによるアプローチ:SMILESトークン化とRNN

より複雑なパターンを学習するため、PyTorchを用いたリカレントニューラルネットワーク(RNN)を実装します。SMILES文字列を文字単位(または正規表現によるトークン単位)で分解し、埋め込み層を経由してRNNに入力します。

トークナイザとデータセットの定義

SMILESを適切にトークン化するため、化学的に意味のある単位(塩素 'Cl'、臭素 'Br' など)で区切る正規表現を使用します。

class SmilesVocab:
    def __init__(self, vocab_file):
        self.pad_token = "<PAD>"
        self.unk_token = "<UNK>"
        self.stoi = {self.pad_token: 0, self.unk_token: 1}
        self.itos = {0: self.pad_token, 1: self.unk_token}
        
        with open(vocab_file, 'r') as f:
            for idx, line in enumerate(f):
                token = line.strip()
                self.stoi[token] = idx + 2
                self.itos[idx + 2] = token
    
    def encode(self, tokens):
        return [self.stoi.get(t, self.stoi[self.unk_token]) for t in tokens]

class SmilesDataset(Dataset):
    def __init__(self, df, vocab, max_len=300, is_train=True):
        self.df = df
        self.vocab = vocab
        self.max_len = max_len
        self.is_train = is_train
        # SMILESの正規表現パターン
        self.pattern = re.compile(r"(\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9])")

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        # 反応物と生成物を結合(例: Reactant1.Reactant2>Product)
        # 必要に応じてAdditiveやSolventも含めることが可能
        smiles = f"{row['Reactant1']}.{row['Reactant2']}>{row['Product']}"
        
        # トークン化
        tokens = self.pattern.findall(smiles)
        tokens = ["<CLS>"] + tokens + ["<SEP>"]
        
        # インデックス化とパディング
        indices = self.vocab.encode(tokens)
        if len(indices) > self.max_len:
            indices = indices[:self.max_len]
        else:
            indices = indices + [self.vocab.stoi["<PAD>"]] * (self.max_len - len(indices))
        
        x = torch.tensor(indices, dtype=torch.long)
        
        if self.is_train:
            y = torch.tensor(row['Yield'], dtype=torch.float)
            return x, y
        return x

# 辞書とデータセットの初期化
vocab = SmilesVocab("../vocab_full.txt")
train_dataset = SmilesDataset(df_train, vocab, is_train=True)
test_dataset = SmilesDataset(df_test, vocab, is_train=False)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

RNNモデルのアーキテクチャ

埋め込み層、双方向RNN層、そして全結合層からなるモデルを定義します。最終的な出力はSigmoid関数などを用いて0~1の範囲(収率)に変換されます。

class ReactionRNN(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, output_dim, n_layers, dropout):
        super(ReactionRNN, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.rnn = nn.RNN(embed_dim, hidden_dim, num_layers=n_layers, 
                          batch_first=True, bidirectional=True, dropout=dropout)
        
        # 全結合層の定義
        self.fc = nn.Sequential(
            nn.Linear(hidden_dim * 2 * n_layers, output_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(output_dim, 1),
            nn.Sigmoid() # 収率は0-1の範囲と仮定
        )

    def forward(self, x):
        # x: [batch, seq_len]
        embedded = self.embedding(x)
        # embedded: [batch, seq_len, embed_dim]
        
        output, hidden = self.rnn(embedded)
        # hidden: [num_layers * num_directions, batch, hidden_dim]
        
        # 隠れ状態を平坦化
        hidden = hidden.transpose(0, 1)
        hidden = hidden.reshape(hidden.shape[0], -1)
        
        return self.fc(hidden).squeeze(1)

学習の実行

損失関数にはMAE(L1Loss)またはMSEを使用し、最適化にはAdamを用います。勾配爆発を防ぐため、Gradient Clippingも実装します。

def run_training():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    model_params = {
        'vocab_size': len(vocab.stoi),
        'embed_dim': 256,
        'hidden_dim': 512,
        'output_dim': 256,
        'n_layers': 2,
        'dropout': 0.3
    }
    
    model = ReactionRNN(**model_params).to(device)
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.L1Loss() # MAE
    epochs = 50
    
    best_loss = float('inf')
    
    print("Start training...")
    for epoch in range(epochs):
        model.train()
        total_loss = 0
        for src, tgt in train_loader:
            src, tgt = src.to(device), tgt.to(device)
            
            optimizer.zero_grad()
            output = model(src)
            loss = criterion(output, tgt)
            loss.backward()
            
            # Gradient Clipping
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            
            optimizer.step()
            total_loss += loss.item()
        
        avg_loss = total_loss / len(train_loader)
        print(f"Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}")
        
        if avg_loss < best_loss:
            best_loss = avg_loss
            torch.save(model.state_dict(), 'best_rnn_model.pth')
            
    print("Training completed.")

if __name__ == '__main__':
    run_training()

テストデータによる予測

学習済みモデルを用いてテストデータの収率を予測し、結果をファイルに出力します。

def predict_with_rnn(model_path, output_csv):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # モデルの再構築
    model_params = {
        'vocab_size': len(vocab.stoi),
        'embed_dim': 256,
        'hidden_dim': 512,
        'output_dim': 256,
        'n_layers': 2,
        'dropout': 0.3
    }
    model = ReactionRNN(**model_params).to(device)
    model.load_state_dict(torch.load(model_path))
    model.eval()
    
    results = []
    with torch.no_grad():
        for batch in test_loader:
            src = batch.to(device)
            preds = model(src)
            results.extend(preds.cpu().tolist())
            
    # CSV書き出し
    with open(output_csv, 'w') as f:
        f.write("rxnid,Yield\n")
        for i, val in enumerate(results):
            f.write(f"test{i+1},{val:.4f}\n")
    print("Prediction saved.")

# predict_with_rnn('best_rnn_model.pth', 'rnn_submission.csv')

7月11日 21:38 投稿