Raise ValueError for float('NaN') in Composition (#3519)

* raise ValueError for float('NaN') in Composition

* test err msg for float('NaN') in TestComposition.test_formula

Author: janosh

pymatgen/core/composition.py

@@ -11,6 +11,7 @@
 import warnings
 from functools import total_ordering
 from itertools import combinations_with_replacement, product
+from math import isnan
 from typing import TYPE_CHECKING, Union, cast
 
 from monty.fractions import gcd, gcd_float
@@ -122,17 +123,19 @@ def __init__(self, *args, strict: bool = False, **kwargs) -> None:
         if len(args) == 1 and isinstance(args[0], Composition):
             elem_map = args[0]
         elif len(args) == 1 and isinstance(args[0], str):
-            elem_map = self._parse_formula(args[0])  # type: ignore
+            elem_map = self._parse_formula(args[0])  # type: ignore[assignment]
+        elif len(args) == 1 and isinstance(args[0], float) and isnan(args[0]):
+            raise ValueError("float('NaN') is not a valid Composition, did you mean str('NaN')?")
         else:
             elem_map = dict(*args, **kwargs)  # type: ignore
         elem_amt = {}
-        self._natoms = 0
+        self._n_atoms = 0
         for key, val in elem_map.items():
             if val < -Composition.amount_tolerance and not self.allow_negative:
                 raise ValueError("Amounts in Composition cannot be negative!")
             if abs(val) >= Composition.amount_tolerance:
                 elem_amt[get_el_sp(key)] = val
-                self._natoms += abs(val)
+                self._n_atoms += abs(val)
         self._data = elem_amt
         if strict and not self.valid:
             raise ValueError(f"Composition is not valid, contains: {', '.join(map(str, self.elements))}")
@@ -201,8 +204,8 @@ def __add__(self, other: object) -> Composition:
 
         new_el_map: dict[SpeciesLike, float] = collections.defaultdict(float)
         new_el_map.update(self)
-        for k, v in other.items():
-            new_el_map[get_el_sp(k)] += v
+        for key, val in other.items():
+            new_el_map[get_el_sp(key)] += val
         return Composition(new_el_map, allow_negative=self.allow_negative)
 
     def __sub__(self, other: object) -> Composition:
@@ -219,8 +222,8 @@ def __sub__(self, other: object) -> Composition:
 
         new_el_map: dict[SpeciesLike, float] = collections.defaultdict(float)
         new_el_map.update(self)
-        for k, v in other.items():
-            new_el_map[get_el_sp(k)] -= v
+        for key, val in other.items():
+            new_el_map[get_el_sp(key)] -= val
         return Composition(new_el_map, allow_negative=self.allow_negative)
 
     def __mul__(self, other: object) -> Composition:
@@ -323,7 +326,7 @@ def fractional_composition(self) -> Composition:
         1.
         E.g. "Fe2 O3".fractional_composition = "Fe0.4 O0.6".
         """
-        return self / self._natoms
+        return self / self._n_atoms
 
     @property
     def reduced_composition(self) -> Composition:
@@ -362,7 +365,7 @@ def get_reduced_formula_and_factor(self, iupac_ordering: bool = False) -> tuple[
         all_int = all(abs(x - round(x)) < Composition.amount_tolerance for x in self.values())
         if not all_int:
             return self.formula.replace(" ", ""), 1
-        d = {k: int(round(v)) for k, v in self.get_el_amt_dict().items()}
+        d = {key: int(round(val)) for key, val in self.get_el_amt_dict().items()}
         formula, factor = reduce_formula(d, iupac_ordering=iupac_ordering)
 
         if formula in Composition.special_formulas:
@@ -395,7 +398,7 @@ def get_integer_formula_and_factor(
         el_amt = self.get_el_amt_dict()
         gcd = gcd_float(list(el_amt.values()), 1 / max_denominator)
 
-        dct = {k: round(v / gcd) for k, v in el_amt.items()}
+        dct = {key: round(val / gcd) for key, val in el_amt.items()}
         formula, factor = reduce_formula(dct, iupac_ordering=iupac_ordering)
         if formula in Composition.special_formulas:
             formula = Composition.special_formulas[formula]
@@ -439,7 +442,7 @@ def elements(self) -> list[Element | Species | DummySpecies]:
         return list(self)
 
     def __str__(self) -> str:
-        return " ".join(f"{k}{formula_double_format(v, ignore_ones=False)}" for k, v in self.as_dict().items())
+        return " ".join(f"{key}{formula_double_format(val, ignore_ones=False)}" for key, val in self.as_dict().items())
 
     def to_pretty_string(self) -> str:
         """
@@ -453,7 +456,7 @@ def num_atoms(self) -> float:
         """Total number of atoms in Composition. For negative amounts, sum
         of absolute values.
         """
-        return self._natoms
+        return self._n_atoms
 
     @property
     def weight(self) -> float:
@@ -469,7 +472,7 @@ def get_atomic_fraction(self, el: SpeciesLike) -> float:
         Returns:
             Atomic fraction for element el in Composition
         """
-        return abs(self[el]) / self._natoms
+        return abs(self[el]) / self._n_atoms
 
     def get_wt_fraction(self, el: SpeciesLike) -> float:
         """Calculate weight fraction of an Element or Species.
@@ -605,7 +608,7 @@ def valid(self) -> bool:
         return not any(isinstance(el, DummySpecies) for el in self.elements)
 
     def __repr__(self) -> str:
-        formula = " ".join(f"{k}{':' if hasattr(k, 'oxi_state') else ''}{v:g}" for k, v in self.items())
+        formula = " ".join(f"{key}{':' if hasattr(key, 'oxi_state') else ''}{val:g}" for key, val in self.items())
         cls_name = type(self).__name__
         return f"{cls_name}({formula!r})"
 
@@ -1250,20 +1253,20 @@ def __init__(self, *args, **kwargs) -> None:
             **kwargs: any valid dict init arguments.
         """
         dct = dict(*args, **kwargs)
-        super().__init__((get_el_sp(k), v) for k, v in dct.items())
+        super().__init__((get_el_sp(key), val) for key, val in dct.items())
         if len(dct) != len(self):
             raise ValueError("Duplicate potential specified")
 
     def __mul__(self, other: object) -> ChemicalPotential:
         if isinstance(other, (int, float)):
-            return ChemicalPotential({k: v * other for k, v in self.items()})
+            return ChemicalPotential({key: val * other for key, val in self.items()})
         return NotImplemented
 
     __rmul__ = __mul__
 
     def __truediv__(self, other: object) -> ChemicalPotential:
         if isinstance(other, (int, float)):
-            return ChemicalPotential({k: v / other for k, v in self.items()})
+            return ChemicalPotential({key: val / other for key, val in self.items()})
         return NotImplemented
 
     __div__ = __truediv__
@@ -1290,7 +1293,7 @@ def get_energy(self, composition: Composition, strict: bool = True) -> float:
         if strict and set(composition) > set(self):
             s = set(composition) - set(self)
             raise ValueError(f"Potentials not specified for {s}")
-        return sum(self.get(k, 0) * v for k, v in composition.items())
+        return sum(self.get(key, 0) * val for key, val in composition.items())
 
     def __repr__(self) -> str:
         return f"ChemPots: {super()!r}"

tests/core/test_composition.py

@@ -186,6 +186,10 @@ def test_formula(self):
 
         assert Composition("Na 3 Zr (PO 4) 3").reduced_formula == "Na3Zr(PO4)3"
 
+        assert Composition("NaN").reduced_formula == "NaN"
+        with pytest.raises(ValueError, match=r"float\('NaN'\) is not a valid Composition, did you mean str\('NaN'\)\?"):
+            Composition(float("NaN"))
+
         # test bad formulas raise ValueError
         for bad_formula in ("", " ", "  2", "4 2", "6123", "1.2", "1/2", "1.2.3"):
             with pytest.raises(ValueError, match=f"Invalid formula={bad_formula!r}"):