CSM计算

corner-sharing/utils/analyze_env_st.py

@@ -47,28 +47,19 @@ class HiddenPrints:
 
 
 def non_elements(struct, sp='Li'):
-    '''
-    struct : structure object from Pymatgen
+    """
+    struct : 必须是一个有序结构
     sp : the mobile specie
-    returns the structure with all of the mobile specie (Li) removed
-    '''
-    num_li = struct.species.count(Element(sp))
-    species = list(set(struct.species))
-    try:
-        species.remove(Element("O"))
-    except ValueError:
-        print("没有O")
-    try:
-        species.remove(Element("S"))
-    except ValueError:
-        print("没有S")
-    try:
-        species.remove(Element("N"))
-    except ValueError:
-        print("没有N")
+    returns a new structure containing only the framework anions (O, S, N).
+    """
+    anions_to_keep = {"O", "S", "N"}
     stripped = struct.copy()
-    stripped.remove_species(species)
-    stripped = stripped.get_sorted_structure(reverse=True)
+    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
 
 
@@ -160,22 +151,43 @@ def site_env(coord, struct, sp="Li", envtype='both'):
 
 
 def extract_sites(struct, sp="Li", envtype='both'):
-    '''
+    """
     struct : structure object from Pymatgen
     envtype : 'tet', 'oct', or 'both'
     sp : target element to analyze environment
-
-    '''
+    """
     envlist = []
-    for i in range(len(struct.sites)):
-        if struct.sites[i].specie != Element(sp):
-            continue
-        site = struct.sites[i]
-        singleenv = site_env(site.frac_coords, struct, sp, envtype)
-        envlist.append({'frac_coords': site.frac_coords, 'type': singleenv['type'], 'csm': singleenv['csm'],
-                        'volume': singleenv['vol']})
-    return envlist
 
+    # --- 关键修改：直接遍历原始结构，即使它是无序的 ---
+    # 我们不再调用 get_sorted_structure()
+    # 我们只关心那些含有目标元素 sp 的位点
+
+    # 遍历每一个位点 (site)
+    for i, site in enumerate(struct):
+        # 检查当前位点的组分(site.species)中是否包含我们感兴趣的元素(sp)
+        # site.species.elements 返回该位点上的元素列表，例如 [Element Li, Element Fe]
+        # [el.symbol for el in site.species.elements] 将其转换为符号列表 ['Li', 'Fe']
+        site_elements = [el.symbol for el in site.species.elements]
+
+        if sp in site_elements:
+            # 如果找到了Li，我们就对这个位点进行环境分析
+            # 注意：我们将原始的、可能无序的 struct 传递给 site_env
+            # 因为 site_env 内部的函数 (如 LocalGeometryFinder) 知道如何处理它
+
+            # 为了让下游函数（特别是 non_elements）能够工作，
+            # 我们在这里创建一个一次性的、临时的有序结构副本给它
+            # 这可以避免我们之前遇到的所有 'ordered structures only' 错误
+            temp_ordered_struct = struct.get_sorted_structure()
+
+            singleenv = site_env(site.frac_coords, temp_ordered_struct, sp, envtype)
+
+            envlist.append({'frac_coords': site.frac_coords, 'type': singleenv['type'], 'csm': singleenv['csm'],
+                            'volume': singleenv['vol']})
+
+    if not envlist:
+        print(f"警告: 在结构中未找到元素 {sp} 的占位。")
+
+    return envlist
 
 def export_envs(envlist, sp='Li', envtype='both', fname=None):
     '''
@@ -193,6 +205,6 @@ def export_envs(envlist, sp='Li', envtype='both', fname=None):
             f.write("Site index " + str(index) + ": " + str(i) + '\n')
 
 
-struct = Structure.from_file("../data/31960.cif")
-site_info = extract_sites(struct, envtype="both")
-export_envs(site_info, sp="Li", envtype="both")
+# struct = Structure.from_file("../data/0921/wjy_475.cif")
+# site_info = extract_sites(struct, envtype="both")
+# export_envs(site_info, sp="Li", envtype="both")