add

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