generate_3d_structures
This commit is contained in:
mm644706215
339
pycomsia/src/balloon_rdkit_pipeline.py
Normal file
339
pycomsia/src/balloon_rdkit_pipeline.py
Normal file
@@ -0,0 +1,339 @@
|
||||
#!/usr/bin/env python
|
||||
"""
|
||||
完整训练代码示例
|
||||
|
||||
功能：
|
||||
1. 从 CSV 文件中读取 SMILES 和目标变量，计算2D描述符（利用 RDKit 与 Mordred），并对描述符进行 SelectKBest 特征选择。
|
||||
2. 根据 CSV 文件中的 SMILES 生成分子的3D构象，然后利用 MolecularGridCalculator 与 MolecularFieldCalculator 计算3D‐QSAR场特征，
|
||||
将5个场（steric、electrostatic、hydrophobic、hbond_donor、hbond_acceptor）展平后合并为一个特征向量。
|
||||
3. 合并2D与3D特征，并使用随机森林和 XGBoost 进行回归训练，同时支持 Optuna 超参数调优。
|
||||
4. 通过命令行传递 CSV 文件路径和目标变量名称，使用 click 库实现，并提供详细帮助信息。
|
||||
|
||||
用法示例：
|
||||
python main.py --data-smi data_smi.csv --target MIC_LOG_ATCC25923
|
||||
"""
|
||||
|
||||
import os
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import click
|
||||
from rdkit import Chem
|
||||
from rdkit.Chem import Descriptors, AllChem
|
||||
from mordred import Calculator, descriptors
|
||||
from sklearn.feature_selection import SelectKBest, f_regression
|
||||
from sklearn.model_selection import train_test_split
|
||||
from sklearn.ensemble import RandomForestRegressor
|
||||
import xgboost as xgb
|
||||
from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error
|
||||
import optuna
|
||||
|
||||
# 导入3D相关模块（请确保 MolecularGridCalculator.py 和 MolecularFieldCalculator.py 在同一目录下）
|
||||
from MolecularGridCalculator import MolecularGridCalculator
|
||||
from MolecularFieldCalculator import MolecularFieldCalculator
|
||||
|
||||
# ------------------------ 辅助函数 ------------------------
|
||||
def safe_nunique(series):
|
||||
"""安全计算唯一值数量，防止非标量数据出错"""
|
||||
try:
|
||||
return series.nunique(dropna=False)
|
||||
except Exception:
|
||||
return series.apply(lambda x: tuple(x) if isinstance(x, (np.ndarray, list)) and not pd.isnull(x) else x).nunique(dropna=False)
|
||||
|
||||
# ------------------------ 2D描述符计算 ------------------------
|
||||
def get_rdkit_descriptors(smiles, missingVal=None):
|
||||
"""利用 RDKit 计算分子所有描述符"""
|
||||
mol = Chem.MolFromSmiles(smiles)
|
||||
res = {}
|
||||
if mol is None:
|
||||
for nm, _ in Descriptors._descList:
|
||||
res[nm] = missingVal
|
||||
return res
|
||||
for nm, fn in Descriptors._descList:
|
||||
try:
|
||||
res[nm] = fn(mol)
|
||||
except Exception:
|
||||
res[nm] = missingVal
|
||||
return res
|
||||
|
||||
def get_mordred_descriptors(smiles):
|
||||
"""利用 Mordred 计算分子描述符"""
|
||||
mol = Chem.MolFromSmiles(smiles)
|
||||
if mol is None:
|
||||
return {}
|
||||
calc = Calculator(descriptors)
|
||||
try:
|
||||
result = calc(mol)
|
||||
return result.asdict()
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
def compute_2d_features(csv_file):
|
||||
data = pd.read_csv(csv_file)
|
||||
# RDKit描述符
|
||||
rdkit_desc_list = data["SMILES"].apply(get_rdkit_descriptors)
|
||||
X_df_rdkit = pd.DataFrame(rdkit_desc_list.tolist())
|
||||
# 转换为数值类型并用均值填充（针对数值列）
|
||||
X_df_rdkit = X_df_rdkit.apply(pd.to_numeric, errors='coerce')
|
||||
X_df_rdkit = X_df_rdkit.fillna(X_df_rdkit.mean(numeric_only=True))
|
||||
|
||||
# Mordred描述符
|
||||
mordred_desc_list = []
|
||||
for smi in data["SMILES"]:
|
||||
mordred_desc_list.append(get_mordred_descriptors(smi))
|
||||
X_df_mordred = pd.DataFrame(mordred_desc_list)
|
||||
X_df_mordred.dropna(axis=1, how='all', inplace=True)
|
||||
invalid_features = [col for col in X_df_mordred.columns if safe_nunique(X_df_mordred[col]) <= 1]
|
||||
if invalid_features:
|
||||
print("Mordred 删除恒定特征:", invalid_features)
|
||||
X_df_mordred.drop(columns=invalid_features, inplace=True)
|
||||
X_df_mordred = X_df_mordred.apply(pd.to_numeric, errors='coerce')
|
||||
X_df_mordred = X_df_mordred.fillna(X_df_mordred.mean(numeric_only=True))
|
||||
|
||||
# 合并两部分描述符
|
||||
X_df = pd.concat([X_df_rdkit, X_df_mordred], axis=1)
|
||||
X_df = X_df.loc[:, ~X_df.columns.duplicated()]
|
||||
constant_features = [col for col in X_df.columns if safe_nunique(X_df[col]) <= 1]
|
||||
if constant_features:
|
||||
print("Combined 删除恒定特征:", constant_features)
|
||||
X_df.drop(columns=constant_features, inplace=True)
|
||||
|
||||
# 检查哪些列仍包含非数值数据（即转换时出现了 NaN，但原始数据非空）
|
||||
non_numeric_features = {}
|
||||
for col in X_df.columns:
|
||||
# 尝试将列转换为数值
|
||||
col_numeric = pd.to_numeric(X_df[col], errors='coerce')
|
||||
# 如果原始列中有非空值，但转换后对应位置为NaN，则说明存在非数值数据
|
||||
mask = X_df[col].notna() & col_numeric.isna()
|
||||
if mask.any():
|
||||
unique_vals = set(X_df.loc[mask, col].unique())
|
||||
non_numeric_features[col] = unique_vals
|
||||
if non_numeric_features:
|
||||
print("以下特征包含非数值数据（经过 set 去重）：")
|
||||
for col, uniq in non_numeric_features.items():
|
||||
print(f"{col}: {uniq}")
|
||||
else:
|
||||
print("所有特征均为数值类型。")
|
||||
|
||||
# 返回处理好的 DataFrame
|
||||
return X_df
|
||||
|
||||
def select_2d_features(X_df, y, k=10):
|
||||
"""
|
||||
使用 SelectKBest 选择 2D 描述符的前 k 个特征，
|
||||
返回转换后的特征矩阵和选中特征名称。
|
||||
"""
|
||||
selector = SelectKBest(score_func=f_regression, k=k)
|
||||
selector.fit(X_df, y)
|
||||
selected_features = X_df.columns[selector.get_support()].tolist()
|
||||
print("SelectKBest 选中的2D特征：", selected_features)
|
||||
X_selected = selector.transform(X_df)
|
||||
return X_selected, selected_features
|
||||
|
||||
# ------------------------ 3D-QSAR特征计算（从 CSV 中 SMILES 生成3D构象） ------------------------
|
||||
def generate_3d_mols_from_csv(csv_file, mmffVariant='MMFF94'):
|
||||
"""
|
||||
从 CSV 文件中读取 SMILES，生成分子的3D构象（添加氢原子、嵌入构象并进行MMFF能量最小化）。
|
||||
针对大环分子，启用了宏环扭转角优化并增大了嵌入尝试次数。
|
||||
|
||||
mmffVariant：MMFF94或MMFF94S，默认为 MMFF94
|
||||
|
||||
返回值：分子列表，每个元素为 (mol, True)
|
||||
"""
|
||||
import pandas as pd
|
||||
from rdkit import Chem
|
||||
from rdkit.Chem import AllChem
|
||||
|
||||
data = pd.read_csv(csv_file)
|
||||
mols = []
|
||||
# 距离几何+ETKDG生成3D构象
|
||||
for smi in data["SMILES"]:
|
||||
mol = Chem.MolFromSmiles(smi)
|
||||
if mol is None:
|
||||
print(f"Warning: 无法从 SMILES {smi} 生成分子。")
|
||||
continue
|
||||
m3d = Chem.AddHs(mol)
|
||||
AllChem.EmbedMolecule(m3d, randomSeed=10, useMacrocycleTorsions=True)
|
||||
# MMFF生成3D构象 优化
|
||||
if m3d.GetNumConformers() > 0:
|
||||
AllChem.MMFFOptimizeMolecule(m3d)
|
||||
mols.append(m3d)
|
||||
else:
|
||||
print(f"Warning: 分子 {smi} 未生成构象。")
|
||||
|
||||
aligned_results = [(mol, True) for mol in mols if mol.GetNumConformers() > 0]
|
||||
return aligned_results
|
||||
|
||||
|
||||
def compute_3d_features_with_params_from_csv(csv_file, grid_spacing, padding, alpha):
|
||||
"""
|
||||
根据 CSV 文件中的 SMILES 生成分子的3D构象后，利用 MolecularGridCalculator 与 MolecularFieldCalculator 计算3D-QSAR场特征，
|
||||
参数 grid_spacing、padding 和 alpha 可调。
|
||||
返回一个 shape=(n_samples, feature_dim) 的特征矩阵。
|
||||
"""
|
||||
# 生成3D分子
|
||||
aligned_results = generate_3d_mols_from_csv(csv_file)
|
||||
|
||||
grid_calc = MolecularGridCalculator()
|
||||
field_calc = MolecularFieldCalculator()
|
||||
|
||||
grid_spacing_tuple, grid_dimensions, grid_origin = grid_calc.generate_grid(aligned_results, resolution=grid_spacing, padding=padding)
|
||||
# 修改高斯衰减参数 alpha
|
||||
field_calc.ALPHA = alpha
|
||||
fields_dict = field_calc.calc_field(aligned_results, grid_spacing_tuple, grid_dimensions, grid_origin)
|
||||
|
||||
selected_field_names = ["steric_field", "electrostatic_field", "hydrophobic_field", "hbond_donor_field", "hbond_acceptor_field"]
|
||||
X_3d_list = []
|
||||
n_mols = len(aligned_results)
|
||||
for i in range(n_mols):
|
||||
feat_vec = []
|
||||
for field in selected_field_names:
|
||||
field_vec = fields_dict['train_fields'][field][i]
|
||||
feat_vec.extend(field_vec)
|
||||
X_3d_list.append(feat_vec)
|
||||
X_3d = np.array(X_3d_list)
|
||||
return X_3d
|
||||
|
||||
def split_fields_from_X3d(X_3d_all, field_dims):
|
||||
"""
|
||||
根据各场的维度将3D特征矩阵 X_3d_all 拆分为列表，
|
||||
field_dims 为各场展平后特征的维度列表。
|
||||
"""
|
||||
fields = []
|
||||
start = 0
|
||||
for dim in field_dims:
|
||||
fields.append(X_3d_all[:, start:start+dim])
|
||||
start += dim
|
||||
return fields
|
||||
|
||||
# ------------------------ 模型训练与评估 ------------------------
|
||||
def evaluate_model(X, y, random_state=42):
|
||||
"""
|
||||
将数据划分为80:20，训练随机森林和 XGBoost 回归模型，
|
||||
返回各模型的 R²、RMSE、MAE 指标字典。
|
||||
"""
|
||||
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=random_state)
|
||||
|
||||
models = {
|
||||
"RandomForest": RandomForestRegressor(random_state=random_state),
|
||||
"XGBoost": xgb.XGBRegressor(random_state=random_state, verbosity=0)
|
||||
}
|
||||
|
||||
results = {}
|
||||
for name, model in models.items():
|
||||
model.fit(X_train, y_train)
|
||||
y_pred = model.predict(X_test)
|
||||
r2 = r2_score(y_test, y_pred)
|
||||
rmse = np.sqrt(mean_squared_error(y_test, y_pred))
|
||||
mae = mean_absolute_error(y_test, y_pred)
|
||||
results[name] = {"R2": r2, "RMSE": rmse, "MAE": mae}
|
||||
return results
|
||||
|
||||
# ------------------------ 超参数调优（Optuna） ------------------------
|
||||
def objective(trial):
|
||||
# 超参数：选择使用哪些3D场
|
||||
use_steric = trial.suggest_categorical("use_steric", [True, False])
|
||||
use_electrostatic = trial.suggest_categorical("use_electrostatic", [True, False])
|
||||
use_hydrophobic = trial.suggest_categorical("use_hydrophobic", [True, False])
|
||||
use_hbond_donor = trial.suggest_categorical("use_hbond_donor", [True, False])
|
||||
use_hbond_acceptor = trial.suggest_categorical("use_hbond_acceptor", [True, False])
|
||||
|
||||
# 对每个场设置权重（若未选中则为0）
|
||||
weight_steric = trial.suggest_float("weight_steric", 0.0, 2.0) if use_steric else 0.0
|
||||
weight_electrostatic = trial.suggest_float("weight_electrostatic", 0.0, 2.0) if use_electrostatic else 0.0
|
||||
weight_hydrophobic = trial.suggest_float("weight_hydrophobic", 0.0, 2.0) if use_hydrophobic else 0.0
|
||||
weight_hbond_donor = trial.suggest_float("weight_hbond_donor", 0.0, 2.0) if use_hbond_donor else 0.0
|
||||
weight_hbond_acceptor = trial.suggest_float("weight_hbond_acceptor", 0.0, 2.0) if use_hbond_acceptor else 0.0
|
||||
|
||||
# 网格与衰减参数
|
||||
grid_spacing = trial.suggest_float("grid_spacing", 0.5, 2.0)
|
||||
alpha = trial.suggest_float("alpha", 0.1, 1.0)
|
||||
|
||||
# 随机森林超参数
|
||||
n_estimators = trial.suggest_int("n_estimators", 50, 300)
|
||||
max_depth = trial.suggest_int("max_depth", 3, 15)
|
||||
|
||||
# 计算3D特征（从 CSV 中生成3D构象）
|
||||
csv_file = click.get_current_context().params.get("data_smi")
|
||||
X_3d_all = compute_3d_features_with_params_from_csv(csv_file, grid_spacing=grid_spacing, padding=3, alpha=alpha)
|
||||
# 假设每个场展平后的维度为 total_dim/5
|
||||
dim_per_field = X_3d_all.shape[1] // 5
|
||||
field_dims = [dim_per_field] * 5
|
||||
X_3d_fields = split_fields_from_X3d(X_3d_all, field_dims)
|
||||
|
||||
# 根据选择情况与权重组合3D特征
|
||||
selected_fields = []
|
||||
for flag, weight, field in zip(
|
||||
[use_steric, use_electrostatic, use_hydrophobic, use_hbond_donor, use_hbond_acceptor],
|
||||
[weight_steric, weight_electrostatic, weight_hydrophobic, weight_hbond_donor, weight_hbond_acceptor],
|
||||
X_3d_fields):
|
||||
if flag:
|
||||
selected_fields.append(field * weight)
|
||||
if selected_fields:
|
||||
X_3d_selected = np.hstack(selected_fields)
|
||||
else:
|
||||
X_3d_selected = np.zeros((X_3d_all.shape[0], 1))
|
||||
|
||||
# 读取2D特征（SelectKBest后的结果）
|
||||
data = pd.read_csv(csv_file)
|
||||
target = click.get_current_context().params.get("target")
|
||||
y = data[target].values
|
||||
X_df = compute_2d_features(csv_file)
|
||||
X_2d_selected, _ = select_2d_features(X_df, y, k=10)
|
||||
|
||||
# 合并2D与3D特征
|
||||
X_combined = np.hstack((X_2d_selected, X_3d_selected))
|
||||
|
||||
# 划分训练/验证集并评估模型
|
||||
X_train, X_val, y_train, y_val = train_test_split(X_combined, y, test_size=0.2, random_state=42)
|
||||
model = RandomForestRegressor(n_estimators=n_estimators, max_depth=max_depth, random_state=42)
|
||||
model.fit(X_train, y_train)
|
||||
y_pred = model.predict(X_val)
|
||||
rmse = np.sqrt(mean_squared_error(y_val, y_pred))
|
||||
return rmse
|
||||
|
||||
# ------------------------ 主程序入口（Click命令行） ------------------------
|
||||
@click.command()
|
||||
@click.option('--data-smi', required=True, type=click.Path(exists=True),
|
||||
help="包含 SMILES 和目标变量的 CSV 文件路径。")
|
||||
@click.option('--target', default="MIC_LOG_ATCC25923",
|
||||
help="目标变量列名，默认：MIC_LOG_ATCC25923。")
|
||||
def cli(data_smi, target):
|
||||
"""
|
||||
通过命令行启动QSAR模型训练和超参数调优流程。
|
||||
|
||||
示例：
|
||||
python main.py --data-smi data_smi.csv --target MIC_LOG_ATCC25923
|
||||
"""
|
||||
ctx = click.get_current_context()
|
||||
ctx.params["data_smi"] = data_smi
|
||||
ctx.params["target"] = target
|
||||
|
||||
# 计算2D描述符和SelectKBest特征
|
||||
data = pd.read_csv(data_smi)
|
||||
y = data[target].values
|
||||
X_df = compute_2d_features(data_smi)
|
||||
X_2d_selected, selected_feats = select_2d_features(X_df, y, k=10)
|
||||
click.echo("2D特征选择完毕：{}".format(selected_feats))
|
||||
|
||||
# 计算3D特征（从 CSV 中 SMILES生成3D构象）
|
||||
X_3d = compute_3d_features_with_params_from_csv(data_smi, grid_spacing=1.0, padding=3, alpha=0.3)
|
||||
click.echo("3D特征计算完毕。")
|
||||
|
||||
# 合并2D与3D特征
|
||||
X_combined = np.hstack((X_2d_selected, X_3d))
|
||||
|
||||
click.echo("开始初步模型训练评估（不调超参）...")
|
||||
results_2d = evaluate_model(X_2d_selected, y)
|
||||
results_combined = evaluate_model(X_combined, y)
|
||||
click.echo("仅2D描述符模型结果：{}".format(results_2d))
|
||||
click.echo("2D + 3D描述符模型结果：{}".format(results_combined))
|
||||
|
||||
click.echo("开始使用Optuna进行超参数调优...")
|
||||
study = optuna.create_study(direction="minimize")
|
||||
study.optimize(objective, n_trials=50)
|
||||
click.echo("Optuna最佳试验结果：")
|
||||
click.echo(study.best_trial)
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli()
|
||||
# python main.py --data-smi /root/project/qsar/1d-qsar/data_smi.csv --target MIC_LOG_ATCC25923
|
||||
Reference in New Issue
Block a user