57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
#!/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 |