一、项目背景
- 任务: 预测化学反应的产率,具体来说是利用反应物、添加剂、溶剂等信息预测反应的产率(Yield)。
- 数据: 数据集包含化学分子结构的SMILES表示及对应的产率。
二、Baseline代码解析
1. 环境依赖与导入
!pip install pandas
!pip install -U scikit-learn
!pip install rdkit
# 导入库
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.*')
2. 特征提取
- 使用
rdkit
库从SMILES表示中提取分子的Morgan fingerprint。
def mfgen(mol,nBits=2048, radius=2):
# 返回分子的位向量形式的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)
3. 数据加载与特征构建
dataset_dir = '../dataset' # # 注:如果是在AI Studio上,将这里改为'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)}')
# 从csv中读取数据
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)
4. 模型选择与超参数调优
- 使用
GridSearchCV
对RandomForestRegressor
的参数进行调优。
# Model fitting
model = RandomForestRegressor(n_estimators=10,max_depth=10,min_samples_split=2,min_samples_leaf=1,n_jobs=-1) # 实例化模型,并指定重要参数
model.fit(train_x,train_y) # 训练模型
# 保存模型
with open('./random_forest_model.pkl', 'wb') as file:
pickle.dump(model, file)
5. 模型训练与预测
# 加载模型
with open('random_forest_model.pkl', 'rb') as file:
loaded_model = pickle.load(file)
# 预测\推理
test_pred = loaded_model.predict(test_x)
6. 结果保存与提交
- 将预测结果保存,以便提交比赛或进一步分析。
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))
三、总结
通过优化特征提取、模型选择和参数调优,可以提升模型的预测精度和稳定性。良好的代码结构和注释不仅有助于理解和维护,也为后续的改进和扩展提供了便利。