CSM及TET，CS

py/update_tet_occupancy.py

@@ -0,0 +1,129 @@
+import os
+import pandas as pd
+from tqdm import tqdm
+
+# ================= 配置区域 =================
+# CSV 所在的根目录
+CSV_ROOT_DIR = "../output"
+# CSM .dat 文件所在的根目录
+CSM_ROOT_DIR = "../output/CSM"
+
+
+# ===========================================
+
+def calculate_tet_ratio_from_dat(dat_path):
+    """
+    解析 .dat 文件，计算四面体位 Li 的占比。
+    返回: float (0.0 - 1.0) 或 None (如果文件不存在或为空)
+    """
+    if not os.path.exists(dat_path):
+        return None
+
+    tet_count = 0
+    total_count = 0
+
+    try:
+        with open(dat_path, 'r', encoding='utf-8') as f:
+            lines = f.readlines()
+
+            # 简单检查文件是否包含 "No environments found"
+            if len(lines) > 0 and "No environments found" in lines[0]:
+                return None
+
+            for line in lines:
+                # 根据截图，每行是一个位点的信息
+                # 简单字符串匹配，这比 eval 更安全且足够快
+                if "'type': 'tet'" in line:
+                    tet_count += 1
+                    total_count += 1
+                elif "'type': 'oct'" in line:
+                    total_count += 1
+                # 如果还有其他类型，可以在这里加，或者只要是位点行都算进 total
+
+        if total_count == 0:
+            return 0.0
+
+        return round(tet_count / total_count, 4)
+
+    except Exception as e:
+        print(f"解析出错 {dat_path}: {e}")
+        return None
+
+
+def process_single_csv(csv_path, group_name, anion_name):
+    """
+    读取 CSV -> 寻找对应的 CSM dat 文件 -> 计算比例 -> 更新 CSV
+    """
+    print(f"正在更新 CSV: {csv_path}")
+
+    # 读取 CSV，确保 ID 是字符串
+    try:
+        df = pd.read_csv(csv_path, dtype={'Filename': str})
+    except Exception as e:
+        print(f"读取 CSV 失败: {e}")
+        return
+
+    tet_ratios = []
+
+    # 遍历 CSV 中的每一行
+    for index, row in tqdm(df.iterrows(), total=df.shape[0], desc="Updating Occupancy"):
+        material_id = str(row['Filename']).replace('.0', '')
+
+        # 构建对应的 .dat 文件路径
+        # 路径逻辑: ../output/CSM/Group/Anion/ID.dat
+        # 注意: 这里的 Group/Anion 结构必须与 analyze_csm.py 生成的一致
+
+        if group_name == anion_name:
+            # 单一阴离子: ../output/CSM/S/123.dat
+            dat_rel_path = os.path.join(group_name, f"{material_id}.dat")
+        else:
+            # 混合阴离子: ../output/CSM/S+O/S/123.dat
+            dat_rel_path = os.path.join(group_name, anion_name, f"{material_id}.dat")
+
+        dat_path = os.path.join(CSM_ROOT_DIR, dat_rel_path)
+
+        # 计算比例
+        ratio = calculate_tet_ratio_from_dat(dat_path)
+        tet_ratios.append(ratio)
+
+    # 添加或更新列
+    df['Tet_Li_Ratio'] = tet_ratios
+
+    # 保存
+    df.to_csv(csv_path, index=False)
+    print(f"  - 已保存更新后的数据到: {csv_path}")
+
+
+def run_update():
+    """
+    主程序：遍历 output 目录下的 CSV
+    """
+    if not os.path.exists(CSV_ROOT_DIR):
+        print(f"CSV 目录不存在: {CSV_ROOT_DIR}")
+        return
+
+    for root, dirs, files in os.walk(CSV_ROOT_DIR):
+        for file in files:
+            if file.endswith(".csv"):
+                csv_path = os.path.join(root, file)
+
+                # 解析路径获取 Group 和 Anion
+                # root: ../output/S  --> rel: S
+                rel_root = os.path.relpath(root, CSV_ROOT_DIR)
+                path_parts = rel_root.split(os.sep)
+
+                if len(path_parts) == 1:
+                    group_name = path_parts[0]
+                    anion_name = path_parts[0]
+                elif len(path_parts) >= 2:
+                    group_name = path_parts[0]
+                    anion_name = path_parts[1]
+                else:
+                    continue
+
+                # 只有当 CSM 目录里有对应的文件夹时才处理（可选）
+                process_single_csv(csv_path, group_name, anion_name)
+
+
+if __name__ == "__main__":
+    run_update()