99 lines
4.2 KiB
Python
99 lines
4.2 KiB
Python
import os
|
||
import shutil
|
||
from pymatgen.core import Structure
|
||
|
||
|
||
def get_anion_type(structure):
|
||
"""
|
||
简单判断阴离子类型。
|
||
返回: 'O', 'S', 'S+O' 等字符串
|
||
"""
|
||
# 定义常见的阴离子列表
|
||
valid_anions = {'O', 'S', 'Se', 'Te', 'F', 'Cl', 'Br', 'I', 'N', 'P'}
|
||
|
||
# 获取结构中的所有元素符号
|
||
elements = set([e.symbol for e in structure.composition.elements])
|
||
|
||
# 取交集找到当前结构包含的阴离子
|
||
found_anions = elements.intersection(valid_anions)
|
||
|
||
if not found_anions:
|
||
return "Unknown"
|
||
|
||
# 如果有多个阴离子，按字母顺序排序并用 '+' 连接 (模拟 step1 的逻辑)
|
||
sorted_anions = sorted(list(found_anions))
|
||
return "+".join(sorted_anions)
|
||
|
||
|
||
def organize_files_direct(input_folder, output_base):
|
||
if not os.path.exists(input_folder):
|
||
print(f"输入文件夹不存在: {input_folder}")
|
||
return
|
||
|
||
cif_files = [f for f in os.listdir(input_folder) if f.endswith(".cif")]
|
||
print(f"发现 {len(cif_files)} 个 CIF 文件，开始直接整理...")
|
||
|
||
for filename in cif_files:
|
||
file_path = os.path.join(input_folder, filename)
|
||
|
||
try:
|
||
# 仅读取结构用于分类
|
||
struct = Structure.from_file(file_path)
|
||
anion_type = get_anion_type(struct)
|
||
|
||
# 获取不带后缀的文件名 (ID)
|
||
file_base_name = os.path.splitext(filename)[0]
|
||
|
||
# --- 构建目标路径逻辑 ---
|
||
# 逻辑：../data/after_step1 / 阴离子类别 / ID / ID.cif
|
||
|
||
# 处理混合阴离子情况 (如 S+O)
|
||
if "+" in anion_type:
|
||
# 按照之前的逻辑，如果是混合阴离子，通常会有多层
|
||
# 但为了统一后续处理，我们这里将其放入组合名的文件夹下
|
||
# 比如: after_step1/S+O/S/123/123.cif (复杂)
|
||
# 或者简化为 after_step1/S+O/123/123.cif (简单)
|
||
# 根据你之前的 make_sh.py 和 extract_data.py，
|
||
# 只要是 Folder/ID/ID.cif 结构即可。
|
||
# 为了兼容 analyze_cs.py 的逻辑 (group_name, anion_name)，
|
||
# 这里我们采用 simplified 逻辑：
|
||
# 如果是混合，我们在第一层建 S+O，第二层建具体的 anion 文件夹(比如首字母排序第一个)
|
||
# 或者直接: after_step1/S+O/ID/ID.cif -> 这样 group=S+O, anion=ID (不对)
|
||
|
||
# 兼容旧代码的最佳实践：
|
||
# 对于混合 S+O，我们建立 S+O/S/ID/ID.cif 和 S+O/O/ID/ID.cif ?
|
||
# 不，原 Step1 是把一个文件复制了两份到不同文件夹。
|
||
# 这里为了简化，我们只复制一份到主阴离子文件夹，或者直接按组合命名。
|
||
|
||
# 让我们采用最稳妥的方式：如果是 S+O，放入 S+O/Mix/ID/ID.cif
|
||
# 这样 group=S+O, anion=Mix。
|
||
# 但为了让 CS_calc 正常工作，最好还是放入具体的元素文件夹。
|
||
# 这里我们简单处理：直接放入 S+O/Combined/ID/
|
||
# 或者根据你的 extract_data.py 逻辑：
|
||
# 它会遍历 top_dir (S+O) -> sub_anion (S, O)
|
||
|
||
# 策略：拆分放入。
|
||
sub_anions = anion_type.split("+")
|
||
for sub in sub_anions:
|
||
target_folder = os.path.join(output_base, anion_type, sub, file_base_name)
|
||
if not os.path.exists(target_folder):
|
||
os.makedirs(target_folder)
|
||
shutil.copy(file_path, os.path.join(target_folder, filename))
|
||
|
||
print(f"整理: {filename} -> {anion_type} (已复制到各子类)")
|
||
|
||
else:
|
||
# 单一阴离子: after_step1/S/ID/ID.cif
|
||
target_folder = os.path.join(output_base, anion_type, file_base_name)
|
||
if not os.path.exists(target_folder):
|
||
os.makedirs(target_folder)
|
||
|
||
shutil.copy(file_path, os.path.join(target_folder, filename))
|
||
print(f"整理: {filename} -> {anion_type}")
|
||
|
||
except Exception as e:
|
||
print(f"处理 {filename} 失败: {e}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
organize_files_direct("../data/input", "../data/after_step1") |