Merge pull request #388 from hmlli/data_type_expand

Data type expand
materialsvirtuallab · Jul 15, 2024 · ae07ac2 · ae07ac2
2 parents 652264f + b082355
commit ae07ac2
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/pymatgen/analysis/diffusion/utils/edge_data_from_sc.py b/pymatgen/analysis/diffusion/utils/edge_data_from_sc.py
@@ -28,7 +28,7 @@ def add_edge_data_from_sc(
     mg: MigrationGraph,
     i_sc: Structure,
     e_sc: Structure,
-    data_array: list,
+    data_array: list | str | float,
     key: str = "custom_key",
     use_host_sg: bool = True,
 ) -> None:
@@ -56,7 +56,10 @@ def add_edge_data_from_sc(
     isite, esite = i_wi[0], e_wi[0]
     uhop_index, mh_from_sc = get_unique_hop(mg, i_sc, isite, esite, use_host_sg)
     add_dict = {key: data_array}
-    mg.add_data_to_similar_edges(target_label=uhop_index, data=add_dict, m_hop=mh_from_sc)
+    if isinstance(data_array, list):
+        mg.add_data_to_similar_edges(target_label=uhop_index, data=add_dict, m_hop=mh_from_sc)
+    else:
+        mg.add_data_to_similar_edges(target_label=uhop_index, data=add_dict)
 
 
 def get_uc_pos(

diff --git a/pymatgen/analysis/diffusion/utils/tests/test_edge_data_from_sc.py b/pymatgen/analysis/diffusion/utils/tests/test_edge_data_from_sc.py
@@ -43,6 +43,16 @@ def test_add_edge_data_from_sc():
         key=test_key,
     )
 
+    test_single_key = "test_single_key"
+    test_single = 100
+    add_edge_data_from_sc(
+        mg_Li,
+        i_sc=input_struct_i,
+        e_sc=input_struct_e,
+        data_array=test_single,
+        key=test_single_key,
+    )
+
     edge_data = []
     for _u, _v, d in mg_Li.m_graph.graph.edges(data=True):
         edge_data.append(d)