224 lines
8.1 KiB
Python
224 lines
8.1 KiB
Python
import os
|
||
import sys
|
||
import numpy as np
|
||
import argparse
|
||
from tqdm import tqdm
|
||
from scipy.spatial import ConvexHull
|
||
from pymatgen.core import Structure
|
||
from pymatgen.core.periodic_table import Element
|
||
from pymatgen.analysis.chemenv.coordination_environments.coordination_geometry_finder import LocalGeometryFinder
|
||
|
||
# ================= 配置区域 =================
|
||
# 建议使用绝对路径，避免找不到文件夹
|
||
INPUT_DIR = "../../solidstate-tools/corner-sharing/data/1209/input" # 请确保这里有你的 .cif 文件
|
||
OUTPUT_DIR = "../output/CSM"
|
||
TARGET_ELEMENT = 'Li'
|
||
ENV_TYPE = 'both'
|
||
|
||
|
||
# ===========================================
|
||
|
||
class HiddenPrints:
|
||
'''用于隐藏 pymatgen 繁杂的输出'''
|
||
|
||
def __enter__(self):
|
||
self._original_stdout = sys.stdout
|
||
sys.stdout = open(os.devnull, 'w')
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
sys.stdout.close()
|
||
sys.stdout = self._original_stdout
|
||
|
||
|
||
def non_elements(struct):
|
||
"""
|
||
【关键修复】保留卤素(F, Cl, Br, I) 和其他阴离子，防止氯化物结构被清空。
|
||
"""
|
||
# 这里加入了 F, Cl, Br, I, P, Se, Te 等
|
||
anions_to_keep = {"O", "S", "N", "F", "Cl", "Br", "I", "P", "Se", "Te", "As", "Sb", "C"}
|
||
stripped = struct.copy()
|
||
species_to_remove = [el.symbol for el in stripped.composition.elements
|
||
if el.symbol not in anions_to_keep]
|
||
if species_to_remove:
|
||
stripped.remove_species(species_to_remove)
|
||
return stripped
|
||
|
||
|
||
def site_env(coord, struct, sp="Li", envtype='both'):
|
||
stripped = non_elements(struct)
|
||
|
||
# 如果剥离后结构为空（例如纯金属锂），直接返回
|
||
if len(stripped) == 0:
|
||
return {'csm': np.nan, 'vol': np.nan, 'type': 'Error_NoAnions'}
|
||
|
||
with_li = stripped.copy()
|
||
# 插入一个探测用的 Li 原子
|
||
with_li.append(sp, coord, coords_are_cartesian=False, validate_proximity=False)
|
||
|
||
# 尝试排序，如果因为部分占据导致排序失败，则使用原始顺序
|
||
try:
|
||
with_li = with_li.get_sorted_structure()
|
||
except:
|
||
pass
|
||
|
||
tet_oct_competition = []
|
||
|
||
# ---------------- 四面体 (Tet) 检测 ----------------
|
||
if envtype == 'both' or envtype == 'tet':
|
||
for dist in np.linspace(1, 4, 601): # 扫描距离 1A 到 4A
|
||
neigh = with_li.get_neighbors(with_li.sites[0], dist)
|
||
if len(neigh) < 4:
|
||
continue
|
||
elif len(neigh) > 4:
|
||
break
|
||
|
||
neigh_coords = [i.coords for i in neigh]
|
||
try:
|
||
with HiddenPrints():
|
||
lgf = LocalGeometryFinder(only_symbols=["T:4"])
|
||
lgf.setup_structure(structure=with_li)
|
||
lgf.setup_local_geometry(isite=0, coords=neigh_coords)
|
||
|
||
site_volume = ConvexHull(neigh_coords).volume
|
||
# 获取 CSM
|
||
csm_val = lgf.get_coordination_symmetry_measures()['T:4']['csm']
|
||
|
||
tet_env = {'csm': csm_val, 'vol': site_volume, 'type': 'tet'}
|
||
tet_oct_competition.append(tet_env)
|
||
except Exception:
|
||
pass
|
||
if len(neigh) == 4: break
|
||
|
||
# ---------------- 八面体 (Oct) 检测 ----------------
|
||
if envtype == 'both' or envtype == 'oct':
|
||
for dist in np.linspace(1, 4, 601):
|
||
neigh = with_li.get_neighbors(with_li.sites[0], dist)
|
||
if len(neigh) < 6:
|
||
continue
|
||
elif len(neigh) > 6:
|
||
break
|
||
|
||
neigh_coords = [i.coords for i in neigh]
|
||
try:
|
||
with HiddenPrints():
|
||
lgf = LocalGeometryFinder(only_symbols=["O:6"], permutations_safe_override=False)
|
||
lgf.setup_structure(structure=with_li)
|
||
lgf.setup_local_geometry(isite=0, coords=neigh_coords)
|
||
|
||
site_volume = ConvexHull(neigh_coords).volume
|
||
csm_val = lgf.get_coordination_symmetry_measures()['O:6']['csm']
|
||
|
||
oct_env = {'csm': csm_val, 'vol': site_volume, 'type': 'oct'}
|
||
tet_oct_competition.append(oct_env)
|
||
except Exception:
|
||
pass
|
||
if len(neigh) == 6: break
|
||
|
||
# ---------------- 结果判定 ----------------
|
||
if len(tet_oct_competition) == 0:
|
||
return {'csm': np.nan, 'vol': np.nan, 'type': 'Non_' + envtype}
|
||
elif len(tet_oct_competition) == 1:
|
||
return tet_oct_competition[0]
|
||
elif len(tet_oct_competition) >= 2:
|
||
return min(tet_oct_competition, key=lambda x: x['csm'])
|
||
|
||
|
||
def extract_sites(struct, sp="Li", envtype='both'):
|
||
envlist = []
|
||
# 遍历所有位点寻找 Li
|
||
for i, site in enumerate(struct):
|
||
site_elements = [el.symbol for el in site.species.elements]
|
||
if sp in site_elements:
|
||
try:
|
||
# 传入结构副本以防修改原结构
|
||
singleenv = site_env(site.frac_coords, struct.copy(), sp, envtype)
|
||
envlist.append({
|
||
'site_index': i,
|
||
'frac_coords': site.frac_coords,
|
||
'type': singleenv.get('type', 'unknown'),
|
||
'csm': singleenv.get('csm', np.nan),
|
||
'volume': singleenv.get('vol', np.nan)
|
||
})
|
||
except Exception as e:
|
||
# 捕捉单个位点计算错误，不中断程序
|
||
# print(f" [Warn] Site {i} calculation failed: {e}")
|
||
pass
|
||
return envlist
|
||
|
||
|
||
def export_envs(envlist, sp, envtype, fname):
|
||
with open(fname, 'w') as f:
|
||
f.write('List of environment information\n')
|
||
f.write(f'Species : {sp}\n')
|
||
f.write(f'Envtype : {envtype}\n')
|
||
for item in envlist:
|
||
# 格式化输出，确保没有数据也能看懂
|
||
f.write(f"Site index {item['site_index']}: {item}\n")
|
||
|
||
|
||
# ================= 主程序 =================
|
||
def run_csm_analysis():
|
||
# 1. 检查目录
|
||
if not os.path.exists(INPUT_DIR):
|
||
print(f"错误: 输入目录不存在 -> {os.path.abspath(INPUT_DIR)}")
|
||
return
|
||
|
||
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))
|
||
|
||
if not cif_files:
|
||
print(f"在 {INPUT_DIR} 中未找到 .cif 文件。")
|
||
return
|
||
|
||
print(f"开始分析 {len(cif_files)} 个文件 (目标元素: {TARGET_ELEMENT}, 包含阴离子: F,Cl,Br,I,O,S,N...)")
|
||
|
||
success_count = 0
|
||
|
||
for cif_path in tqdm(cif_files, desc="Calculating CSM"):
|
||
try:
|
||
# 准备路径
|
||
rel_path = os.path.relpath(cif_path, INPUT_DIR)
|
||
rel_dir = os.path.dirname(rel_path)
|
||
file_base = os.path.splitext(os.path.basename(cif_path))[0]
|
||
|
||
target_dir = os.path.join(OUTPUT_DIR, rel_dir)
|
||
if not os.path.exists(target_dir):
|
||
os.makedirs(target_dir)
|
||
|
||
target_dat_path = os.path.join(target_dir, f"{file_base}.dat")
|
||
|
||
# 如果文件已存在且不为空，可选择跳过
|
||
# if os.path.exists(target_dat_path) and os.path.getsize(target_dat_path) > 0:
|
||
# continue
|
||
|
||
# 读取结构
|
||
struct = Structure.from_file(cif_path)
|
||
|
||
# 检查是否含 Li
|
||
if Element(TARGET_ELEMENT) not in struct.composition.elements:
|
||
continue
|
||
|
||
# 计算环境
|
||
env_list = extract_sites(struct, sp=TARGET_ELEMENT, envtype=ENV_TYPE)
|
||
|
||
# 写入结果 (即使 env_list 为空也写入一个标记文件，方便debug)
|
||
if env_list:
|
||
export_envs(env_list, sp=TARGET_ELEMENT, envtype=ENV_TYPE, fname=target_dat_path)
|
||
success_count += 1
|
||
else:
|
||
with open(target_dat_path, 'w') as f:
|
||
f.write(f"No {TARGET_ELEMENT} environments found (Check connectivity or anion types).")
|
||
|
||
except Exception as e:
|
||
print(f"\n[Error] File: {os.path.basename(cif_path)} -> {e}")
|
||
continue
|
||
|
||
print(f"\n分析完成！成功生成 {success_count} 个文件。")
|
||
print(f"输出目录: {os.path.abspath(OUTPUT_DIR)}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
run_csm_analysis() |