V1

2025-12-07 13:56:33 +08:00
parent 49f54b04cd
commit c83985cd02
22 changed files with 2732 additions and 0 deletions
--- a/py/step6.py
+++ b/py/step6.py
@@ -0,0 +1,75 @@
+from pymatgen.core import Structure
+from pymatgen.analysis.structure_matcher import StructureMatcher
+import os
+import shutil
+
+# 定义函数，用于将 CIF 文件进行分类并输出
+
+def structure_classify(input_folder, output_folder):
+    """
+    分类 input_folder 中的 CIF 文件，并根据框架分组后存储到 output_folder 中。
+
+    :param input_folder: 输入文件夹路径，包含 .cif 文件。
+    :param output_folder: 输出文件夹路径，用于存储分类后的 .cif 文件。
+    """
+    # 检查输入文件夹是否存在
+    if not os.path.exists(input_folder):
+        print(f"输入文件夹 {input_folder} 不存在")
+        return
+
+    # 创建输出文件夹（如果不存在）
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
+
+    # 读取输入文件夹中的所有 .cif 文件
+    structures = []
+    file_map = {}  # 记录文件路径与结构的对应关系
+    for filename in os.listdir(input_folder):
+        if filename.endswith(".cif"):
+            file_path = os.path.join(input_folder, filename)
+            try:
+                structure = Structure.from_file(file_path)
+                structures.append(structure)
+                file_map[id(structure)] = filename  # 使用结构对象的 id 作为键
+            except Exception as e:
+                print(f"无法读取文件 {filename}：{e}")
+
+    # 检查是否成功加载了任何结构
+    if not structures:
+        print("未找到有效的 CIF 文件")
+        return
+
+    # 分组结构
+    matcher = StructureMatcher()
+    grouped_structures = []
+
+    for structure in structures:
+        matched = False
+        for group in grouped_structures:
+            if matcher.fit(structure, group[0]):  # 比较结构是否匹配
+                group.append(structure)
+                matched = True
+                break
+        if not matched:
+            grouped_structures.append([structure])
+
+    # 保存分组后的结构到输出文件夹
+    for group_index, group in enumerate(grouped_structures):
+        group_folder = os.path.join(output_folder, f"group_{group_index}")
+        os.makedirs(group_folder, exist_ok=True)
+
+        for structure_index, structure in enumerate(group):
+            output_file = os.path.join(group_folder, f"structure_{structure_index}.cif")
+            structure.to(filename=output_file)
+
+            # 获取原始文件名并复制到分组文件夹
+            original_filename = file_map[id(structure)]  # 使用结构对象的 id 获取文件名
+            original_file_path = os.path.join(input_folder, original_filename)
+            shutil.copy(original_file_path, group_folder)
+
+    print(f"处理完成，分类后的结构已保存到 {output_folder}")
+
+if __name__ == "__main__":
+    # 示例调用
+    structure_classify("../data/after_step5/S", "../data/after_step6/S")
+    structure_classify("../data/after_step5/O", "../data/after_step6/O")