CSM及TET，CS

2025-12-07 20:08:19 +08:00
parent 08f5a51fc4
commit 3d44b31194
8 changed files with 1087 additions and 0 deletions
--- a/py/csm.py
+++ b/py/csm.py
@@ -0,0 +1,90 @@
+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 = "../data/after_screening"
+# 输出目录
+OUTPUT_DIR = "../output/CSM"
+# 分析参数
+TARGET_ELEMENT = 'Li'
+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()