90 lines
3.5 KiB
Python
90 lines
3.5 KiB
Python
import os
|
|
from pymatgen.core import Structure
|
|
from pymatgen.core.periodic_table import Element
|
|
# 导入你的CSM计算工具库 (根据 provided context [11])
|
|
try:
|
|
from utils.analyze_env_st import extract_sites, export_envs
|
|
except ImportError:
|
|
print("Error: 找不到 utils.analyze_env_st 模块，请检查 utils 文件夹。")
|
|
exit()
|
|
|
|
from tqdm import tqdm
|
|
|
|
# ================= 配置区域 =================
|
|
# 输入目录：使用筛选后的目录，只计算符合要求的材料
|
|
INPUT_DIR = "../../solidstate-tools/corner-sharing/data/1209/input"
|
|
# 输出目录
|
|
OUTPUT_DIR = "../output/CSM"
|
|
# 分析参数
|
|
TARGET_ELEMENT = 'Na'
|
|
ENV_TYPE = 'both' # 可选 'tet', 'oct', 'both'
|
|
# ===========================================
|
|
|
|
def run_csm_analysis():
|
|
"""
|
|
遍历 after_screening 文件夹，计算 CSM 并生成 .dat 文件到 output/CSM
|
|
"""
|
|
if not os.path.exists(INPUT_DIR):
|
|
print(f"输入目录不存在: {INPUT_DIR}，请先运行筛选步骤。")
|
|
return
|
|
|
|
# 收集所有需要处理的 CIF 文件
|
|
cif_files = []
|
|
for root, dirs, files in os.walk(INPUT_DIR):
|
|
for file in files:
|
|
if file.endswith(".cif"):
|
|
# 保存完整路径
|
|
cif_files.append(os.path.join(root, file))
|
|
|
|
print(f"开始进行 CSM 分析，共找到 {len(cif_files)} 个筛选后的材料...")
|
|
|
|
for cif_path in tqdm(cif_files, desc="Calculating CSM"):
|
|
try:
|
|
# 1. 确定输出路径，保持目录结构
|
|
# 获取相对路径 (例如: S/195819.cif 或 S+O/S/195819.cif)
|
|
rel_path = os.path.relpath(cif_path, INPUT_DIR)
|
|
# 获取所在文件夹 (例如: S 或 S+O/S)
|
|
rel_dir = os.path.dirname(rel_path)
|
|
# 获取文件名 (例如: 195819)
|
|
file_base = os.path.splitext(os.path.basename(cif_path))[0]
|
|
|
|
# 构建目标文件夹: ../output/CSM/S/
|
|
target_dir = os.path.join(OUTPUT_DIR, rel_dir)
|
|
if not os.path.exists(target_dir):
|
|
os.makedirs(target_dir)
|
|
|
|
# 构建目标文件路径: ../output/CSM/S/195819.dat
|
|
target_dat_path = os.path.join(target_dir, f"{file_base}.dat")
|
|
|
|
# 2. 如果已经存在，跳过 (可选，视需求而定，这里默认覆盖)
|
|
# if os.path.exists(target_dat_path):
|
|
# continue
|
|
|
|
# 3. 读取结构
|
|
struct = Structure.from_file(cif_path)
|
|
|
|
# 检查是否包含目标元素 (Li)
|
|
if Element(TARGET_ELEMENT) not in struct.composition.elements:
|
|
# print(f"Skipping {file_base}: No {TARGET_ELEMENT}")
|
|
continue
|
|
|
|
# 4. 计算 CSM (引用 utils 中的函数)
|
|
# extract_sites 返回环境列表
|
|
env_list = extract_sites(struct, sp=TARGET_ELEMENT, envtype=ENV_TYPE)
|
|
|
|
# 5. 导出结果 (引用 utils 中的函数)
|
|
# export_envs 将结果写入 .dat 文件
|
|
if env_list:
|
|
export_envs(env_list, sp=TARGET_ELEMENT, envtype=ENV_TYPE, fname=target_dat_path)
|
|
else:
|
|
# 如果没有提取到环境（例如没有配位环境），生成一个空文件或记录日志
|
|
with open(target_dat_path, 'w') as f:
|
|
f.write("No environments found.")
|
|
|
|
except Exception as e:
|
|
print(f"处理出错 {cif_path}: {e}")
|
|
|
|
print(f"CSM 分析完成，结果已保存至 {OUTPUT_DIR}")
|
|
|
|
if __name__ == "__main__":
|
|
run_csm_analysis() |