76 lines
2.9 KiB
Python
76 lines
2.9 KiB
Python
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")
|