Update fragment_processor.py

fragment_processor.py

@@ -1,50 +1,103 @@
-import os
 import pandas as pd
+import os
 from pathlib import Path
 from mmpdblib.fragment_io import read_fragment_records
 from rdkit import Chem
+class Index_Dummy:
+    """对 dummy 原子进行编号：变量和常量部分分别处理"""
+    def __init__(self, df):
+        self.df = df
+
+    def index_constant(self, constSmi, attachmentOrder):
+        count = -1
+        newConstSmi = ""
+        for idx, ichar in enumerate(constSmi):
+            if ichar == '*':
+                count += 1
+                # 注意：attachmentOrder 应为可迭代对象，这里假设传入的 attachmentOrder 为列表或可转换为列表
+                ichar = f"[*:{int(attachmentOrder[count]) + 1}]"
+            newConstSmi += ichar
+        return newConstSmi
+
+    def index_var(self, varSmi):
+        count = 0
+        newVarSmi = ""
+        for idx, ichar in enumerate(varSmi):
+            if ichar == '*':
+                count += 1
+                ichar = f"[*:{count}]"
+            newVarSmi += ichar
+        return newVarSmi
+
+    def add_index(self):
+        for idx, row in self.df.iterrows():
+            varSmi = row['variable_smiles']
+            constSmi = row['constant_smiles']
+            attachmentOrder = row['attachment_order']
+            self.df.loc[idx, 'variable_smiles'] = self.index_var(varSmi)
+            self.df.loc[idx, 'constant_smiles'] = self.index_constant(constSmi, attachmentOrder)
+        return self.df
+
+
+def count_heavy_atoms(smi):
+    mol = Chem.MolFromSmiles(smi)
+    if not mol:
+        return 0
+    heavy_count = len([atom for atom in mol.GetAtoms() if atom.GetAtomicNum() > 1])
+    return heavy_count
 
-def fragmentize_molecule(smiles_string, max_ratio=0.5):
-    # 创建临时文件名
+
+def fragmentize_molecule(smiles_string: str, max_ratio: float = 0.5) -> pd.DataFrame:
+    """
+    对单个分子进行 fragment 化处理：
+      1. 将 SMILES 字符串写入临时文件（同时写入标题信息）
+      2. 使用 mmpdb 工具 fragment 化分子
+      3. 读取 fragment 文件，并依据 heavy atom 个数筛选合适的 fragment
+      4. 对 fragment 中 dummy 原子添加编号
+      5. 最后返回 DataFrame 格式的 fragment 数据
+    """
+    # 定义临时文件名（这里保证文件名唯一性可根据需要进一步改进）
     input_file = "temp_input.smi"
     output_file = "temp_output.fragments"
-    
-    # 将SMILES字符串写入临时输入文件
-    with open(input_file, "w") as f:
-        f.write(smiles_string + "\t" + "Molecule" + "\n")
-    
-    # 使用mmpdb工具进行分子碎片化
-    os.system(f"mmpdb fragment {input_file} -o {output_file}")
-    
-    # 读取并处理碎片
-    fragment_reader = read_fragment_records(output_file)
-    fragment_list = []
-    
-    for record in fragment_reader:
-        for frag in record.fragments:
-            if count_heavy_atoms(frag.variable_smiles) < count_heavy_atoms(record.normalized_smiles) * max_ratio:
-                fragment_list.append({
-                    'variable_smiles': frag.variable_smiles,
-                    'constant_smiles': frag.constant_smiles,
-                    'record_id': record.id,
-                    'normalized_smiles': record.normalized_smiles,
-                    'attachment_order': frag.attachment_order
-                })
-    
-    # 删除临时文件
-    os.remove(input_file)
-    os.remove(output_file)
-    
-    # 返回碎片列表
-    return pd.DataFrame(fragment_list)
-
-def count_heavy_atoms(smiles):
-    # 使用RDKit计算重原子数
-    mol = Chem.MolFromSmiles(smiles)
-    return mol.GetNumHeavyAtoms() if mol else 0
-
-# 示例调用
-# smiles = "O=C1CCCC2=C1C1(CCS(=O)(=O)C1)N=C(Nc1nc3ccccc3o1)N2"
-# fragment_df = fragmentize_molecule(smiles)
-
-# print(fragment_df)
+
+    try:
+        # 将 SMILES 字符串写入临时输入文件（标题默认写 “Molecule”）
+        with open(input_file, "w") as f:
+            f.write(smiles_string + "\t" + "Molecule" + "\n")
+
+        # 使用 mmpdb 工具进行分子碎片化
+        ret = os.system(f"mmpdb fragment {input_file} -o {output_file}")
+        if ret != 0:
+            raise Exception("mmpdb fragment 命令执行失败，请确保 mmpdb 工具安装并配置正确。")
+
+        # 读取并处理碎片
+        fragment_reader = read_fragment_records(output_file)
+        frag_list = []
+        for record in fragment_reader:
+            # 打印或记录当前处理的 record 信息，可根据需要选择注释掉
+            print(f"Processing record: {record.id}, {record.normalized_smiles}")
+            for frag in record.fragments:
+                if count_heavy_atoms(frag.variable_smiles) < count_heavy_atoms(record.normalized_smiles) * max_ratio:
+                    frag_list.append({
+                        'variable_smiles': frag.variable_smiles,
+                        'constant_smiles': frag.constant_smiles,
+                        'record_id': record.id,
+                        'normalized_smiles': record.normalized_smiles,
+                        'attachment_order': frag.attachment_order
+                    })
+
+        if not frag_list:
+            raise Exception("未找到满足筛选条件的碎片。")
+
+        # 构造 DataFrame，并对 dummy 原子添加编号
+        df_frag = pd.DataFrame(frag_list)
+        index_dummy = Index_Dummy(df_frag)
+        df_frag = index_dummy.add_index()
+        return df_frag
+
+    finally:
+        # 删除临时文件，确保每次调用结束后文件被清理
+        if Path(input_file).exists():
+            os.remove(input_file)
+        if Path(output_file).exists():
+            os.remove(output_file)