171 lines
7.6 KiB
Python
171 lines
7.6 KiB
Python
from pathlib import Path
|
||
import pandas as pd
|
||
from rdkit import Chem
|
||
from rdkit.Chem import AllChem
|
||
import traceback
|
||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||
from rich.console import Console
|
||
from rich.table import Table
|
||
from rich.panel import Panel
|
||
|
||
console = Console()
|
||
|
||
def is_valid_sdf(sdf_path):
|
||
try:
|
||
# 尝试读取SDF文件为mol对象
|
||
suppl = Chem.SDMolSupplier(str(sdf_path), sanitize=False)
|
||
mols = [mol for mol in suppl if mol is not None]
|
||
return len(mols) > 0
|
||
except Exception:
|
||
return False
|
||
|
||
def smiles_to_3d_sdf(identifier, smiles, props, sdf_path, max_attempts=10):
|
||
try:
|
||
mol = Chem.MolFromSmiles(smiles)
|
||
if mol is None:
|
||
return identifier, False, "SMILES解析失败"
|
||
mol = Chem.AddHs(mol)
|
||
params = AllChem.ETKDGv3()
|
||
last_error = ""
|
||
for attempt in range(max_attempts):
|
||
try:
|
||
status = AllChem.EmbedMolecule(mol, params)
|
||
if status == 0:
|
||
AllChem.UFFOptimizeMolecule(mol)
|
||
if props:
|
||
for k, v in props.items():
|
||
mol.SetProp(str(k), str(v))
|
||
writer = Chem.SDWriter(str(sdf_path))
|
||
writer.write(mol)
|
||
writer.close()
|
||
return identifier, True, f"成功（第{attempt+1}次）"
|
||
except Exception as e:
|
||
last_error = f"3D生成异常: {e}"
|
||
continue
|
||
return identifier, False, last_error if last_error else f"3D构象生成失败（已重试{max_attempts}次）"
|
||
except Exception as e:
|
||
return identifier, False, "其它异常: " + traceback.format_exc(limit=1)
|
||
|
||
def smiles_to_3d_sdf_tuple(args):
|
||
return smiles_to_3d_sdf(*args)
|
||
|
||
def batch_csv_to_3d_sdf_parallel(csv_path, output_dir, smiles_col, id_col, n_jobs=4, max_attempts=10):
|
||
output_dir = Path(output_dir)
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
df = pd.read_csv(csv_path, sep=',', dtype=str)
|
||
tasks = []
|
||
skipped = []
|
||
for idx, row in df.iterrows():
|
||
smiles = row[smiles_col]
|
||
identifier = row[id_col]
|
||
props = row.to_dict()
|
||
sdf_file = output_dir / f"{identifier}.sdf"
|
||
if sdf_file.exists():
|
||
if is_valid_sdf(sdf_file):
|
||
# SDF存在且可读，跳过
|
||
skipped.append(identifier)
|
||
continue
|
||
else:
|
||
# SDF存在但不可读，认为损坏，先删除
|
||
try:
|
||
sdf_file.unlink()
|
||
console.print(f"[red]⚡发现损坏SDF文件 {sdf_file.name}，已删除，准备重新生成[/red]")
|
||
except Exception as e:
|
||
console.print(f"[bold magenta]❗无法删除损坏SDF: {sdf_file.name}, {e}[/]")
|
||
tasks.append((identifier, smiles, props, sdf_file, max_attempts))
|
||
|
||
console.rule(f"[bold green]共 {len(df)} 个分子，{len(skipped)} 个已存在且有效，{len(tasks)} 个待处理（使用 {n_jobs} 并行进程）[/]")
|
||
|
||
results = []
|
||
if tasks:
|
||
with ProcessPoolExecutor(max_workers=n_jobs) as executor:
|
||
future_to_identifier = {executor.submit(smiles_to_3d_sdf_tuple, task): task[0] for task in tasks}
|
||
for i, future in enumerate(as_completed(future_to_identifier), 1):
|
||
identifier, success, msg = future.result()
|
||
results.append((identifier, success, msg))
|
||
if success:
|
||
console.print(f"[bold green]✅ [{identifier}] 处理成功。[/][dim]{msg}[/]")
|
||
else:
|
||
if "SMILES解析失败" in msg:
|
||
console.print(f"[bold red]❌ [{identifier}] SMILES解析失败: {msg}[/]")
|
||
elif "3D" in msg:
|
||
console.print(f"[yellow]⚠️ [{identifier}] 3D构象生成失败: {msg}[/]")
|
||
else:
|
||
console.print(f"[magenta]❗ [{identifier}] 其它错误: {msg}[/]")
|
||
|
||
# 分类失败原因
|
||
failed = [r for r in results if not r[1]]
|
||
succeed = [r for r in results if r[1]]
|
||
failed_smiles = [r for r in failed if "SMILES解析失败" in r[2]]
|
||
failed_3d = [r for r in failed if ("3D" in r[2]) and ("SMILES解析失败" not in r[2])]
|
||
failed_other = [r for r in failed if r not in failed_smiles and r not in failed_3d]
|
||
|
||
# 展示 summary
|
||
table = Table(title="处理结果统计", show_lines=True)
|
||
table.add_column("状态", justify="center", style="cyan")
|
||
table.add_column("数量", justify="center")
|
||
table.add_row("成功", str(len(succeed)))
|
||
table.add_row("已跳过（已存在有效SDF）", str(len(skipped)))
|
||
table.add_row("SMILES解析失败", str(len(failed_smiles)))
|
||
table.add_row("3D构象失败", str(len(failed_3d)))
|
||
table.add_row("其它失败", str(len(failed_other)))
|
||
console.print(table)
|
||
|
||
# 输出详细失败信息
|
||
if failed:
|
||
fail_file = output_dir / "failed_smiles.txt"
|
||
with open(fail_file, "w", encoding="utf-8") as f:
|
||
for identifier, _, msg in failed:
|
||
f.write(f"{identifier}\t{msg}\n")
|
||
# 分类面板
|
||
if failed_smiles:
|
||
console.print(
|
||
Panel(
|
||
f"SMILES解析失败: [yellow]{', '.join([r[0] for r in failed_smiles])}[/yellow]",
|
||
title="[bold red]SMILES解析失败分子[/bold red]",
|
||
style="red"
|
||
)
|
||
)
|
||
if failed_3d:
|
||
console.print(
|
||
Panel(
|
||
f"3D构象失败: [yellow]{', '.join([r[0] for r in failed_3d])}[/yellow]",
|
||
title="[bold yellow]3D构象生成失败分子[/bold yellow]",
|
||
style="yellow"
|
||
)
|
||
)
|
||
if failed_other:
|
||
console.print(
|
||
Panel(
|
||
f"其它异常: [yellow]{', '.join([r[0] for r in failed_other])}[/yellow]",
|
||
title="[bold magenta]其它失败分子[/bold magenta]",
|
||
style="magenta"
|
||
)
|
||
)
|
||
console.print(
|
||
Panel(
|
||
f"共 [red]{len(failed)}[/red] 个分子失败，详情见: [bold]{fail_file.resolve()}[/bold]",
|
||
title="[bold red]失败分子统计[/bold red]",
|
||
style="red"
|
||
)
|
||
)
|
||
else:
|
||
console.print(Panel("[bold green]全部分子处理成功或已跳过！[/bold green]", style="green"))
|
||
|
||
if __name__ == '__main__':
|
||
import argparse
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('--csv', type=str, required=True, help='csv文件路径')
|
||
parser.add_argument('--outdir', type=str, default='./sdf_files', help='SDF输出文件夹')
|
||
parser.add_argument('--smiles_col', type=str, default='canonical_smiles', help='SMILES列名')
|
||
parser.add_argument('--id_col', type=str, default='identifier', help='标识符列名')
|
||
parser.add_argument('--n_jobs', type=int, default=4, help='并行进程数')
|
||
parser.add_argument('--max_attempts', type=int, default=10, help='最大尝试次数')
|
||
args = parser.parse_args()
|
||
|
||
batch_csv_to_3d_sdf_parallel(
|
||
args.csv, args.outdir, args.smiles_col, args.id_col,
|
||
n_jobs=args.n_jobs, max_attempts=args.max_attempts
|
||
)
|
||
# use example:
|
||
# python gen_sdf_parallel.py --csv coconut_data_info.csv --outdir ./sdf_files --n_jobs 8 --max_attempts 10 --smiles_col canonical_smiles --id_col identifier |