Enable a list of specie for pymatgen

bjmorgan · Dec 13, 2022 · dc998b6 · dc998b6
1 parent f121ed9
commit dc998b6
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/kinisi/parser.py b/kinisi/parser.py
@@ -302,7 +302,10 @@ def get_structure_coords_latt(
     def get_indices(
         structure: "pymatgen.core.structure.Structure", specie: Union["pymatgen.core.periodic_table.Element",
                                                                       "pymatgen.core.periodic_table.Specie",
-                                                                      "pymatgen.core.periodic_table.Species"]
+                                                                      "pymatgen.core.periodic_table.Species",
+                                                                      List["pymatgen.core.periodic_table.Element"],
+                                                                      List["pymatgen.core.periodic_table.Specie"],
+                                                                      List["pymatgen.core.periodic_table.Species"]]
     ) -> Tuple[np.ndarray, np.ndarray]:
         """
         Determine framework and non-framework indices for a :py:mod:`pymatgen` compatible file.
@@ -315,8 +318,10 @@ def get_indices(
         """
         indices = []
         framework_indices = []
+        if not isinstance(specie, List):
+            specie = [specie]
         for i, site in enumerate(structure):
-            if site.specie.__str__() == specie:
+            if site.specie.__str__() in specie:
                 indices.append(i)
             else:
                 framework_indices.append(i)