return self in SiteCollection.remove_oxidation_states()

materialsproject · Jan 23, 2024 · b06a0ad · b06a0ad
1 parent f9e2830
commit b06a0ad
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/pymatgen/core/structure.py b/pymatgen/core/structure.py
@@ -578,7 +578,7 @@ def add_oxidation_state_by_site(self, oxidation_states: list[float]) -> None:
                 new_sp[Species(sym, ox)] = occu
             site.species = Composition(new_sp)
 
-    def remove_oxidation_states(self) -> None:
+    def remove_oxidation_states(self) -> SiteCollection:
         """Removes oxidation states from a structure."""
         for site in self:
             new_sp: dict[Element, float] = collections.defaultdict(float)
@@ -587,6 +587,8 @@ def remove_oxidation_states(self) -> None:
                 new_sp[Element(sym)] += occu
             site.species = Composition(new_sp)
 
+        return self
+
     def add_oxidation_state_by_guess(self, **kwargs) -> None:
         """Decorates the structure with oxidation state, guessing
         using Composition.oxi_state_guesses(). If multiple guesses are found