催化反应产率预测

简介

这是 Datawhale 2024 年 AI 夏令营第三期的学习活动《从零入门AI for Science》中的一个示例，基于天池平台第二届世界科学智能大赛物质科学赛道：催化反应产率预测开展的实践学习。

这是一个简单的上手baseline，数据集包括测试集和训练集：

训练集包含23538条反应数据：
- 样本序号：
  - rxnid：样本序号
- 样本的特征：
  - Reactant1：反应物1，用SMILES表示。SMILES将化学分子用ASCII字符表示的方法，是化学信息学领域非常重要的工具
  - Reactant2：反应物2，用SMILES表示
  - Product：产物，用SMILES表示
  - dditive：催化剂，用SMILES表示
  - Solvent：溶剂，用SMILES表示
- 样本标签：
  - Yield：反应产率
测试集包含2616条反应数据，和训练集相比，没有样本标签。

分解步骤

1. 导入必要的库

#首先，导入库
import pickle
import pandas as pd
from tqdm import tqdm
from sklearn.ensemble import RandomForestRegressor
from rdkit.Chem import rdMolDescriptors
from rdkit import RDLogger,Chem
import numpy as np
RDLogger.DisableLog('rdApp.*')

pickle主要是用来保存训练的模型到本地，这样只需训练一次得到模型，后面可以直接使用，不需要重新训练。
rdkit是一个开源的化学信息python软件包。rdMolDescriptors是RDKit 库中的一个模块，用于量化分子的结构特征。

2. 从数据集提取数据

def mfgen(mol,nBits=2048, radius=2):
  '''
  Parameters
  ----------
  mol : mol
  RDKit mol object.
  nBits : int
  Number of bits for the fingerprint.
  radius : int
  Radius of the Morgan fingerprint.
  Returns
  -------
  mf_desc_map : ndarray
  ndarray of molecular fingerprint descriptors.
  '''
  # 返回分子的位向量形式的Morgan fingerprint
  fp = rdMolDescriptors.GetMorganFingerprintAsBitVect(mol,radius=radius,nBits=nBits)
  return np.array(list(map(eval,list(fp.ToBitString()))))
 
 # 加载数据
 def vec_cpd_lst(smi_lst):
  smi_set = list(set(smi_lst))
  smi_vec_map = {}
  for smi in tqdm(smi_set): # tqdm：显示进度条
  mol = Chem.MolFromSmiles(smi)
  smi_vec_map[smi] = mfgen(mol)
  smi_vec_map[''] = np.zeros(2048)

  vec_lst = [smi_vec_map[smi] for smi in smi_lst]
  return np.array(vec_lst)

mfgen函数用于生成 Morgan 指纹。

vec_cpd_lst函数将数据集中的各物质SMILES字符串转化为分子指纹向量构成的数组。

使用这两个函数，可以将原始的分子字符串数据集转化为可用于模型训练数据集：

# 使用Pandas读取数据集
dataset_dir = '../dataset' 
 
train_df = pd.read_csv(f'{dataset_dir}/round1_train_data.csv')
test_df = pd.read_csv(f'{dataset_dir}/round1_test_data.csv')
 
print(f'Training set size: {len(train_df)}, test set size: {len(test_df)}')
 
 # 使用Pandas读取数据集
train_rct1_smi = train_df['Reactant1'].to_list()
train_rct2_smi = train_df['Reactant2'].to_list()
train_add_smi = train_df['Additive'].to_list()
train_sol_smi = train_df['Solvent'].to_list()
 
 # 将SMILES转化为分子指纹
train_rct1_fp = vec_cpd_lst(train_rct1_smi)
train_rct2_fp = vec_cpd_lst(train_rct2_smi)
train_add_fp = vec_cpd_lst(train_add_smi)
train_sol_fp = vec_cpd_lst(train_sol_smi)
# 在dim=1维度进行拼接。即：将一条数据的Reactant1,Reactant2,Product,Additive,Solvent字段的morgan fingerprint拼接为一个特征向量。
train_x = np.concatenate([train_rct1_fp,train_rct2_fp,train_add_fp,train_sol_fp],axis=1)
train_y = train_df['Yield'].to_numpy()
 
 # 测试集也进行同样的操作
test_rct1_smi = test_df['Reactant1'].to_list()
test_rct2_smi = test_df['Reactant2'].to_list()
test_add_smi = test_df['Additive'].to_list()
test_sol_smi = test_df['Solvent'].to_list()
 
test_rct1_fp = vec_cpd_lst(test_rct1_smi)
test_rct2_fp = vec_cpd_lst(test_rct2_smi)
test_add_fp = vec_cpd_lst(test_add_smi)
test_sol_fp = vec_cpd_lst(test_sol_smi)
test_x = np.concatenate([test_rct1_fp,test_rct2_fp,test_add_fp,test_sol_fp],axis=1)

3. 训练模型

实例化并训练一个随机森林回归模型：

# 模型训练
model = RandomForestRegressor(n_estimators=10  # 决策树的个数，越多通常效果越好，但也容易过拟合
            ,max_depth=10  # 设置树的最大深度，与剪枝相关的参数，默认值None
            ,min_samples_split=2  # 根据属性划分节点时，最少的样本数，默认是2，调整模型是否过拟合或欠拟合
            ,min_samples_leaf=1  # 叶子节点最少的样本数，默认是1，也是调整模型是否过拟合或欠拟合
            ,n_jobs=-1)  # 并行job个数，-1表示使用所有cpu进行并行计算
model.fit(train_x,train_y)

保存模型，避免后面使用时重新训练模型：

with open('./random_forest_model.pkl', 'wb') as file:
  pickle.dump(model, file)</pre>

加载模型并预测：

with open('random_forest_model.pkl', 'rb') as file:
  loaded_model = pickle.load(file)
 # 测试集预测
test_pred = loaded_model.predict(test_x)</pre>

4. 生产结果文本

生成赛题要求的submit文件：

ans_str_lst = ['rxnid,Yield']
for idx,y in enumerate(test_pred):
  ans_str_lst.append(f'test{idx+1},{y:.4f}')
with open('./submit.txt','w') as fw:
  fw.writelines('\n'.join(ans_str_lst))