diff --git a/pycomsia/src/generate_3d_structures.py b/pycomsia/src/generate_3d_structures.py new file mode 100644 index 0000000..3c9e4a6 --- /dev/null +++ b/pycomsia/src/generate_3d_structures.py @@ -0,0 +1,182 @@ +import os +import subprocess +import multiprocessing as mp +import pandas as pd +import time +from pathlib import Path +from rdkit import Chem +from rdkit.Chem import AllChem + +def generate_rdkit_sdf(csv_file): + """ + 读取 CSV 中的所有字段, + 利用 RDKit 生成分子对象(2D,添加氢并嵌入一个初步构象), + 并保存为 SDF 文件到当前目录下的 “rdkit_output” 文件夹中, + 文件名为 .sdf,同时在分子属性中写入 CSV 中所有信息。 + + 返回一个列表,每个元素为 (compound, smi, rdkit_sdf_path)。 + """ + data = pd.read_csv(csv_file) + required_cols = {"SMILES", "Compound"} + if not required_cols.issubset(set(data.columns)): + raise ValueError("CSV 文件必须包含 'SMILES' 和 'Compound' 列。") + + rdkit_folder = Path.cwd() / "rdkit_output" + rdkit_folder.mkdir(exist_ok=True) + + file_list = [] + for _, row in data.iterrows(): + smi = row["SMILES"] + compound = str(row["Compound"]) + mol = Chem.MolFromSmiles(smi) + if mol is None: + print(f"Warning: 无法从 SMILES {smi} 生成分子对象。") + continue + + # 遍历 CSV 这一行的所有列,添加到分子属性中 + for key, value in row.items(): + mol.SetProp(str(key), str(value)) + + # 添加氢并生成初步构象(2D) + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol) + + out_file = rdkit_folder / f"{compound}.sdf" + writer = Chem.SDWriter(str(out_file)) + writer.write(mol) + writer.close() + file_list.append((compound, smi, out_file)) + return file_list + +def fallback_rdkit_generate(compound, smi): + """ + 使用 RDKit 的 ETKDGv3 方法尝试生成 3D 构象,并将生成的构象保存到 balloon_output 文件夹中, + 文件名为 fallback_.sdf,同时打印日志。 + """ + mol = Chem.MolFromSmiles(smi) + if mol is None: + print(f"Warning: fallback 无法从 SMILES {smi} 生成分子对象。") + return None + # 添加 CSV 中所有属性,这里至少添加 Compound 信息 + mol.SetProp("Compound", compound) + + mol = Chem.AddHs(mol) + params = AllChem.ETKDGv3() + params.randomSeed = 10 + params.useMacrocycleTorsions = True + params.maxAttempts = 10000 + _ = AllChem.EmbedMultipleConfs(mol, numConfs=5, params=params) + if mol.GetNumConformers() == 0: + print(f"Warning: fallback 分子 {smi} 未生成构象。") + return None + energies = [] + for conf_id in range(mol.GetNumConformers()): + props = AllChem.MMFFGetMoleculeProperties(mol) + ff = AllChem.MMFFGetMoleculeForceField(mol, props, confId=conf_id) + ff.Minimize() + energies.append(ff.CalcEnergy()) + min_conf = energies.index(min(energies)) + conf_ids = [conf.GetId() for conf in mol.GetConformers()] + for cid in conf_ids: + if cid != min_conf: + mol.RemoveConformer(cid) + AllChem.MMFFOptimizeMolecule(mol) + + print(f"Fallback: 使用 RDKit 成功生成构象:{compound}") + + # 保存生成的构象到 balloon_output 文件夹 + balloon_folder = Path.cwd() / "balloon_output" + balloon_folder.mkdir(exist_ok=True) + output_sdf = balloon_folder / f"fallback_{compound}.sdf" + writer = Chem.SDWriter(str(output_sdf)) + writer.write(mol) + writer.close() + + return mol + +def balloon_worker(args): + """ + 多进程 worker: + 参数 args 为 (compound, smi, input_sdf, balloon_path)。 + 使用 Balloon 对 input_sdf 文件进行 3D 构象生成,输出写入当前目录下的 “balloon_output” 文件夹, + 文件名为 .sdf。 + 增加重试机制:如果一次调用失败(即在最多等待 5 秒内输出文件未生成),则重试最多 10 次, + 每次重试时打印重试日志信息;若10次都失败,则调用 fallback_rdkit_generate() 进行后备生成。 + 返回生成的分子对象(或 None)。 + """ + compound, smi, input_sdf, balloon_path = args + balloon_folder = Path.cwd() / "balloon_output" + balloon_folder.mkdir(exist_ok=True) + output_sdf = balloon_folder / f"{compound}.sdf" + + max_retries = 10 # 重试 + success = False + for attempt in range(1, max_retries + 1): + # 调用 Balloon 生成构象 + cmd = [ + balloon_path, + "--nconfs", "1", + "--strict", + str(input_sdf), + "-o", "sdf", + str(output_sdf) + ] + try: + subprocess.check_output(cmd, stderr=subprocess.STDOUT, universal_newlines=True) + except subprocess.CalledProcessError as e: + print(f"Error: Balloon 处理 {compound} 第 {attempt} 次出错,输出:\n{e.output}") + + # 最多等待 5 秒检查输出文件是否生成 + wait_time = 0 + while wait_time < 5: + if output_sdf.exists(): + success = True + break + time.sleep(1) + wait_time += 1 + + if success: + break + else: + print(f"重试 {attempt}: {compound} 的输出文件 {output_sdf} 仍未生成。") + + if not success: + print(f"Warning: 尝试 {max_retries} 次后,{compound} 的 Balloon 输出仍不存在,采用 fallback。") + return fallback_rdkit_generate(compound, smi) + + with open(str(output_sdf), "r") as f: + sdf_content = f.read() + mol = Chem.MolFromMolBlock(sdf_content, sanitize=True, removeHs=False) + if mol is None: + print(f"Warning: Balloon 未能生成有效构象:{compound}") + return fallback_rdkit_generate(compound, smi) + return mol + +def generate_3d_mols_from_balloon(csv_file, balloon_path): + """ + 主函数: + 1. 调用 RDKit 生成 SDF 文件,保存在 “rdkit_output” 文件夹中(文件名为 .sdf), + 分子对象中包含 CSV 中所有字段信息,便于溯源。 + 2. 使用多进程调用 Balloon 对每个 SDF 文件生成 3D 构象,输出到 “balloon_output” 文件夹(文件名为 .sdf)。 + 若 Balloon 多次重试后仍失败,则使用 RDKit 的 ETKDGv3 生成 3D 构象,并保存 fallback 版本。 + 3. 返回一个列表,每个元素为 (mol, True)。 + """ + rdkit_files = generate_rdkit_sdf(csv_file) + if not rdkit_files: + print("没有生成任何 RDKit 的 SDF 文件。") + return [] + + tasks = [(compound, smi, rdkit_sdf, balloon_path) for compound, smi, rdkit_sdf in rdkit_files] + with mp.Pool() as pool: + results = pool.map(balloon_worker, tasks) + + mols = [mol for mol in results if mol is not None and mol.GetNumConformers() > 0] + aligned_results = [(mol, True) for mol in mols] + return aligned_results + +if __name__ == "__main__": + csv_file = "/root/project/qsar/1d-qsar/data_smi.csv" # 替换为实际 CSV 文件路径 + balloon_path = "/root/project/qsar/pycomsia/src/balloon" # Balloon 二进制程序路径 + + aligned_results = generate_3d_mols_from_balloon(csv_file, balloon_path) + print(f"共生成 {len(aligned_results)} 个3D构象的分子。") diff --git a/pycomsia/src/balloon_rdkit_pipeline.py b/pycomsia/src/train_qsar_model.py similarity index 100% rename from pycomsia/src/balloon_rdkit_pipeline.py rename to pycomsia/src/train_qsar_model.py