Clarify return float/int type for core.Composition.reduced_composition and siblings, minor type clean up for core.Composition (#4265)

* add some types

* more types

* fix factor type

* fix IUPAC writing

* fix another factor type

* fix factor

* fix type

* add test for composition error

* add test for reduce formula

Author: DanielYang59

src/pymatgen/core/composition.py

@@ -24,7 +24,7 @@
 from pymatgen.util.string import Stringify, formula_double_format
 
 if TYPE_CHECKING:
-    from collections.abc import Generator, ItemsView, Iterator
+    from collections.abc import Generator, ItemsView, Iterator, Mapping
     from typing import Any, ClassVar, Literal
 
     from typing_extensions import Self
@@ -115,11 +115,11 @@ class Composition(collections.abc.Hashable, collections.abc.Mapping, MSONable, S
     # Tolerance in distinguishing different composition amounts.
     # 1e-8 is fairly tight, but should cut out most floating point arithmetic
     # errors.
-    amount_tolerance = 1e-8
-    charge_balanced_tolerance = 1e-8
+    amount_tolerance: ClassVar[float] = 1e-8
+    charge_balanced_tolerance: ClassVar[float] = 1e-8
 
     # Special formula handling for peroxides and certain elements. This is so
-    # that formula output does not write LiO instead of Li2O2 for example.
+    # that formula output does not write "LiO" instead of "Li2O2" for example.
     special_formulas: ClassVar[dict[str, str]] = {
         "LiO": "Li2O2",
         "NaO": "Na2O2",
@@ -134,7 +134,8 @@ class Composition(collections.abc.Hashable, collections.abc.Mapping, MSONable, S
         "H": "H2",
     }
 
-    oxi_prob = None  # prior probability of oxidation used by oxi_state_guesses
+    # Prior probability of oxidation used by oxi_state_guesses
+    oxi_prob: ClassVar[dict | None] = None
 
     def __init__(self, *args, strict: bool = False, **kwargs) -> None:
         """Very flexible Composition construction, similar to the built-in Python
@@ -417,35 +418,38 @@ def get_reduced_composition_and_factor(self) -> tuple[Self, float]:
         """Calculate a reduced composition and factor.
 
         Returns:
-            A normalized composition and a multiplicative factor, i.e.,
-            Li4Fe4P4O16 returns (Composition("LiFePO4"), 4).
+            tuple[Composition, float]: Normalized Composition and multiplicative factor,
+            i.e. "Li4Fe4P4O16" returns (Composition("LiFePO4"), 4).
         """
-        factor = self.get_reduced_formula_and_factor()[1]
+        factor: float = self.get_reduced_formula_and_factor()[1]
         return self / factor, factor
 
     def get_reduced_formula_and_factor(self, iupac_ordering: bool = False) -> tuple[str, float]:
         """Calculate a reduced formula and factor.
 
         Args:
             iupac_ordering (bool, optional): Whether to order the
-                formula by the iupac "electronegativity" series, defined in
+                formula by the IUPAC "electronegativity" series, defined in
                 Table VI of "Nomenclature of Inorganic Chemistry (IUPAC
                 Recommendations 2005)". This ordering effectively follows
-                the groups and rows of the periodic table, except the
+                the groups and rows of the periodic table, except for
                 Lanthanides, Actinides and hydrogen. Note that polyanions
                 will still be determined based on the true electronegativity of
                 the elements.
 
         Returns:
-            A pretty normalized formula and a multiplicative factor, i.e.,
-            Li4Fe4P4O16 returns (LiFePO4, 4).
+            tuple[str, float]: Normalized formula and multiplicative factor,
+                i.e., "Li4Fe4P4O16" returns (LiFePO4, 4).
         """
-        all_int = all(abs(val - round(val)) < type(self).amount_tolerance for val in self.values())
+        all_int: bool = all(abs(val - round(val)) < type(self).amount_tolerance for val in self.values())
         if not all_int:
             return self.formula.replace(" ", ""), 1
-        el_amt_dict = {key: round(val) for key, val in self.get_el_amt_dict().items()}
+
+        el_amt_dict: dict[str, int] = {key: round(val) for key, val in self.get_el_amt_dict().items()}
+        factor: float
         formula, factor = reduce_formula(el_amt_dict, iupac_ordering=iupac_ordering)
 
+        # Do not "completely reduce" certain formulas
         if formula in type(self).special_formulas:
             formula = type(self).special_formulas[formula]
             factor /= 2
@@ -458,10 +462,10 @@ def get_integer_formula_and_factor(
         """Calculate an integer formula and factor.
 
         Args:
-            max_denominator (int): all amounts in the el:amt dict are
-                first converted to a Fraction with this maximum denominator
+            max_denominator (int): all amounts in the el_amt dict are
+                first converted to a Fraction with this maximum denominator.
             iupac_ordering (bool, optional): Whether to order the
-                formula by the iupac "electronegativity" series, defined in
+                formula by the IUPAC "electronegativity" series, defined in
                 Table VI of "Nomenclature of Inorganic Chemistry (IUPAC
                 Recommendations 2005)". This ordering effectively follows
                 the groups and rows of the periodic table, except the
@@ -470,13 +474,14 @@ def get_integer_formula_and_factor(
                 the elements.
 
         Returns:
-            A pretty normalized formula and a multiplicative factor, i.e.,
-            Li0.5O0.25 returns (Li2O, 0.25). O0.25 returns (O2, 0.125)
+            A normalized formula and a multiplicative factor, i.e.,
+            "Li0.5O0.25" returns (Li2O, 0.25). "O0.25" returns (O2, 0.125)
         """
-        el_amt = self.get_el_amt_dict()
-        _gcd = gcd_float(list(el_amt.values()), 1 / max_denominator)
+        el_amt: dict[str, float] = self.get_el_amt_dict()
+        _gcd: float = gcd_float(list(el_amt.values()), 1 / max_denominator)
 
-        dct = {key: round(val / _gcd) for key, val in el_amt.items()}
+        dct: dict[str, int] = {key: round(val / _gcd) for key, val in el_amt.items()}
+        factor: float
         formula, factor = reduce_formula(dct, iupac_ordering=iupac_ordering)
         if formula in type(self).special_formulas:
             formula = type(self).special_formulas[formula]
@@ -485,8 +490,8 @@ def get_integer_formula_and_factor(
 
     @property
     def reduced_formula(self) -> str:
-        """A pretty normalized formula, i.e., LiFePO4 instead of
-        Li4Fe4P4O16.
+        """A normalized formula, i.e., "LiFePO4" instead of
+        "Li4Fe4P4O16".
         """
         return self.get_reduced_formula_and_factor()[0]
 
@@ -1055,7 +1060,7 @@ def _get_oxi_state_guesses(
                 raise ValueError(f"Composition {comp} cannot accommodate max_sites setting!")
 
         # Load prior probabilities of oxidation states, used to rank solutions
-        if not type(self).oxi_prob:
+        if type(self).oxi_prob is None:
             all_data = loadfn(f"{MODULE_DIR}/../analysis/icsd_bv.yaml")
             type(self).oxi_prob = {Species.from_str(sp): data for sp, data in all_data["occurrence"].items()}
         oxi_states_override = oxi_states_override or {}
@@ -1319,38 +1324,38 @@ def _parse_chomp_and_rank(match, formula: str, m_dict: dict[str, float], m_point
 
 
 def reduce_formula(
-    sym_amt: dict[str, float] | dict[str, int],
+    sym_amt: Mapping[str, float],
     iupac_ordering: bool = False,
-) -> tuple[str, float]:
-    """Helper function to reduce a sym_amt dict to a reduced formula and factor.
+) -> tuple[str, int]:
+    """Helper function to reduce a symbol-amount mapping.
 
     Args:
-        sym_amt (dict): {symbol: amount}.
+        sym_amt (dict[str, float]): Symbol to amount mapping.
         iupac_ordering (bool, optional): Whether to order the
-            formula by the iupac "electronegativity" series, defined in
+            formula by the IUPAC "electronegativity" series, defined in
             Table VI of "Nomenclature of Inorganic Chemistry (IUPAC
             Recommendations 2005)". This ordering effectively follows
-            the groups and rows of the periodic table, except the
+            the groups and rows of the periodic table, except for
             Lanthanides, Actinides and hydrogen. Note that polyanions
             will still be determined based on the true electronegativity of
             the elements.
 
     Returns:
-        tuple[str, float]: reduced formula and factor.
+        tuple[str, int]: reduced formula and factor.
     """
-    syms = sorted(sym_amt, key=lambda x: [get_el_sp(x).X, x])
+    syms: list[str] = sorted(sym_amt, key=lambda x: [get_el_sp(x).X, x])
 
     syms = list(filter(lambda x: abs(sym_amt[x]) > Composition.amount_tolerance, syms))
 
-    factor = 1
-    # Enforce integers for doing gcd.
+    # Enforce integer for calculating greatest common divisor
+    factor: int = 1
     if all(int(i) == i for i in sym_amt.values()):
         factor = abs(gcd(*(int(i) for i in sym_amt.values())))
 
-    poly_anions = []
-    # if the composition contains a poly anion
+    # If the composition contains polyanion
+    poly_anions: list[str] = []
     if len(syms) >= 3 and get_el_sp(syms[-1]).X - get_el_sp(syms[-2]).X < 1.65:
-        poly_sym_amt = {syms[i]: sym_amt[syms[i]] / factor for i in [-2, -1]}
+        poly_sym_amt: dict[str, float] = {syms[i]: sym_amt[syms[i]] / factor for i in (-2, -1)}
         poly_form, poly_factor = reduce_formula(poly_sym_amt, iupac_ordering=iupac_ordering)
 
         if poly_factor != 1:
@@ -1369,9 +1374,17 @@ def reduce_formula(
     return "".join([*reduced_form, *poly_anions]), factor
 
 
+class CompositionError(Exception):
+    """Composition exceptions."""
+
+
 class ChemicalPotential(dict, MSONable):
-    """Represent set of chemical potentials. Can be: multiplied/divided by a Number
-    multiplied by a Composition (returns an energy) added/subtracted with other ChemicalPotentials.
+    """Represent set of chemical potentials.
+
+    Can be:
+        - multiplied/divided by a Number
+        - multiplied by a Composition (returns an energy)
+        - added/subtracted with other ChemicalPotentials
     """
 
     def __init__(self, *args, **kwargs) -> None:
@@ -1424,7 +1437,3 @@ def get_energy(self, composition: Composition, strict: bool = True) -> float:
         if strict and (missing := set(composition) - set(self)):
             raise ValueError(f"Potentials not specified for {missing}")
         return sum(self.get(key, 0) * val for key, val in composition.items())
-
-
-class CompositionError(Exception):
-    """Exception class for composition errors."""

src/pymatgen/core/ion.py

@@ -141,7 +141,7 @@ def get_reduced_formula_and_factor(
 
         Args:
             iupac_ordering (bool, optional): Whether to order the
-                formula by the iupac "electronegativity" series, defined in
+                formula by the IUPAC "electronegativity" series, defined in
                 Table VI of "Nomenclature of Inorganic Chemistry (IUPAC
                 Recommendations 2005)". This ordering effectively follows
                 the groups and rows of the periodic table, except the

tests/core/test_composition.py

@@ -12,7 +12,7 @@
 from pytest import approx
 
 from pymatgen.core import Composition, DummySpecies, Element, Species
-from pymatgen.core.composition import ChemicalPotential
+from pymatgen.core.composition import ChemicalPotential, CompositionError, reduce_formula
 from pymatgen.util.testing import PymatgenTest
 
 
@@ -316,7 +316,7 @@ def test_reduced_formula(self):
         all_formulas = [c.reduced_formula for c in self.comps]
         assert all_formulas == correct_reduced_formulas
 
-        # test iupac reduced formula (polyanions should still appear at the end)
+        # test IUPAC reduced formula (polyanions should still appear at the end)
         all_formulas = [c.get_reduced_formula_and_factor(iupac_ordering=True)[0] for c in self.comps]
         assert all_formulas == correct_reduced_formulas
         assert Composition("H6CN").get_integer_formula_and_factor(iupac_ordering=True)[0] == "CNH6"
@@ -347,7 +347,7 @@ def test_integer_formula(self):
         assert formula == "Li(BH)6"
         assert factor == approx(1 / 6)
 
-        # test iupac reduced formula (polyanions should still appear at the end)
+        # test IUPAC reduced formula (polyanions should still appear at the end)
         all_formulas = [c.get_integer_formula_and_factor(iupac_ordering=True)[0] for c in self.comps]
         assert all_formulas == correct_reduced_formulas
         assert Composition("H6CN0.5").get_integer_formula_and_factor(iupac_ordering=True) == ("C2NH12", 0.5)
@@ -860,6 +860,18 @@ def test_isotopes(self):
         assert "Deuterium" in [elem.long_name for elem in composition.elements]
 
 
+def test_reduce_formula():
+    assert reduce_formula({"Li": 2, "Mn": 4, "O": 8}) == ("LiMn2O4", 2)
+    assert reduce_formula({"Li": 4, "O": 4}) == ("LiO", 4)
+    assert reduce_formula({"Zn": 2, "O": 2, "H": 2}) == ("ZnHO", 2)
+
+
+def test_composition_error():
+    error = CompositionError("Composition error")
+    assert isinstance(error, CompositionError)
+    assert str(error) == "Composition error"
+
+
 class TestChemicalPotential:
     def test_init(self):
         dct = {"Fe": 1, Element("Fe"): 1}