diff --git a/script/optimize_sdf.py b/script/optimize_sdf.py new file mode 100644 index 0000000..28e4ba5 --- /dev/null +++ b/script/optimize_sdf.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- +''' +@file :optimize_sdf.py +@Description: : 使用rdkit 将2d转3d,增加基本的坐标优化功能 +@Date :2024/09/07 09:34:41 +@Author :lyzeng +@Email :pylyzeng@gmail.com +@version :1.0 +''' + + + +import os +import click +from rdkit import Chem +from rdkit.Chem import AllChem + +def generate_3d_coordinates(input_sdf: str, output_sdf: str): + # 读取SDF文件 + suppl = Chem.SDMolSupplier(input_sdf) + writer = Chem.SDWriter(output_sdf) + + for mol in suppl: + if mol is None: + continue + + # 生成3D坐标 + mol = Chem.AddHs(mol) # 添加氢原子 + if AllChem.EmbedMolecule(mol) == 0: # 生成3D坐标 + AllChem.UFFOptimizeMolecule(mol) # 优化分子结构 + writer.write(mol) # 将分子写入输出文件 + + writer.close() + +@click.command() +@click.argument('input_sdf', type=click.File('r')) +@click.option('--output_dir', type=click.Path(), default='.', help='输出目录,默认为当前目录') +def optimize_sdf(input_sdf, output_dir): + """优化输入SDF文件中的3D坐标并输出到指定目录""" + # 获取输入文件名并构造输出文件路径 + input_filename = os.path.basename(input_sdf.name) + output_sdf = os.path.join(output_dir, f"optimized_{input_filename}") + + # 创建输出目录(如果不存在) + os.makedirs(output_dir, exist_ok=True) + + # 生成3D结构并保存 + generate_3d_coordinates(input_sdf.name, output_sdf) + + click.echo(f"优化的SDF文件已保存到: {output_sdf}") + +if __name__ == '__main__': + optimize_sdf() + +# python optimize_sdf.py /mnt/c/project/pymol/data/Structure2D_COMPOUND_CID_5280440.sdf --output_dir /mnt/c/project/pymol/data/output +# python optimize_sdf.py --help \ No newline at end of file