Breaking: Return True for Element in Composition if Species.symbol matches Element (#3184)

* refactor get_el_sp()

* fix Composition.__contains__() by handling Species and DummySpecies correctly

* test all edge cases of Composition.__contains__

- Species in Composition
- Element in Composition with Species
- str in Composition with Element
- int (atomic number) in Composition
- float in Composition
- DummySpecies in Composition

* revert accidental removal of if Element("O") in comp:

* fix test_process_entry_with_oxidation_state

>       assert len(processed_entry.energy_adjustments) == 1
E       AssertionError: assert 2 == 1

* fix test_get_supercell_size()

expected wrong error msg

* fix expected err msg

Author: janosh

pymatgen/analysis/tests/test_structure_matcher.py

@@ -77,7 +77,7 @@ def test_get_supercell_size(self):
         assert sm._get_supercell_size(s2, s1) == (1, True)
 
         sm = StructureMatcher(supercell_size="wfieoh")
-        with pytest.raises(ValueError, match="Can't parse Element or String from str: wfieoh"):
+        with pytest.raises(ValueError, match="Can't parse Element or Species from str: wfieoh"):
             sm._get_supercell_size(s1, s2)
 
     def test_cmp_fstruct(self):

pymatgen/core/composition.py

@@ -151,7 +151,10 @@ def __iter__(self) -> Iterator[Species | Element | DummySpecies]:
     def __contains__(self, key) -> bool:
         try:
             sp = get_el_sp(key)
-            return sp in self._data
+            if isinstance(sp, Species):
+                return sp in self._data
+            # Element or str
+            return any(sp.symbol == s.symbol for s in self._data)
         except ValueError as exc:
             raise TypeError(f"Invalid {key=} for Composition") from exc
 
@@ -160,9 +163,9 @@ def __eq__(self, other: object) -> bool:
         if not isinstance(other, (Composition, dict)):
             return NotImplemented
 
-        #  elements with amounts < Composition.amount_tolerance don't show up
-        #  in the el_map, so checking len enables us to only check one
-        #  composition's elements
+        # elements with amounts < Composition.amount_tolerance don't show up
+        # in the el_map, so checking len enables us to only check one
+        # composition's elements
         if len(self) != len(other):
             return False
 

pymatgen/core/periodic_table.py

@@ -1540,15 +1540,15 @@ class DummySpecie(DummySpecies):
 
 
 @functools.lru_cache
-def get_el_sp(obj) -> Element | Species | DummySpecies:
-    """
-    Utility method to get an Element or Species from an input obj.
+def get_el_sp(obj: int | SpeciesLike) -> Element | Species | DummySpecies:
+    """Utility method to get an Element, Species or DummySpecies from any input.
+
     If obj is in itself an element or a specie, it is returned automatically.
-    If obj is an int or a string representing an integer, the Element
-    with the atomic number obj is returned.
-    If obj is a string, Species parsing will be attempted (e.g., Mn2+), failing
-    which Element parsing will be attempted (e.g., Mn), failing which
-    DummyElement parsing will be attempted.
+    If obj is an int or a string representing an integer, the Element with the
+    atomic number obj is returned.
+    If obj is a string, Species parsing will be attempted (e.g. Mn2+). Failing that
+    Element parsing will be attempted (e.g. Mn). Failing that DummyElement parsing
+    will be attempted.
 
     Args:
         obj (Element/Species/str/int): An arbitrary object. Supported objects
@@ -1560,28 +1560,28 @@ def get_el_sp(obj) -> Element | Species | DummySpecies:
         that can be determined.
 
     Raises:
-        ValueError if obj cannot be converted into an Element or Species.
+        ValueError: if obj cannot be converted into an Element or Species.
     """
     if isinstance(obj, (Element, Species, DummySpecies)):
         return obj
 
     try:
-        c = float(obj)
-        i = int(c)
-        i = i if i == c else None  # type: ignore
+        flt = float(obj)
+        integer = int(flt)
+        integer = integer if integer == flt else None  # type: ignore
+        return Element.from_Z(integer)
     except (ValueError, TypeError):
-        i = None
-
-    if i is not None:
-        return Element.from_Z(i)
+        pass
 
     try:
-        return Species.from_str(obj)
+        return Species.from_str(obj)  # type: ignore
     except (ValueError, KeyError):
-        try:
-            return Element(obj)
-        except (ValueError, KeyError):
-            try:
-                return DummySpecies.from_str(obj)
-            except Exception:
-                raise ValueError(f"Can't parse Element or String from {type(obj).__name__}: {obj}.")
+        pass
+    try:
+        return Element(obj)  # type: ignore
+    except (ValueError, KeyError):
+        pass
+    try:
+        return DummySpecies.from_str(obj)  # type: ignore
+    except Exception:
+        raise ValueError(f"Can't parse Element or Species from {type(obj).__name__}: {obj}.")

pymatgen/core/tests/test_composition.py

@@ -12,7 +12,7 @@
 from pytest import approx
 
 from pymatgen.core.composition import ChemicalPotential, Composition
-from pymatgen.core.periodic_table import Element, Species
+from pymatgen.core.periodic_table import DummySpecies, Element, Species
 from pymatgen.util.testing import PymatgenTest
 
 
@@ -52,6 +52,7 @@ def test_immutable(self):
         assert "'Composition' object does not support item deletion" in str(exc.value)
 
     def test_in(self):
+        # test the Composition.__contains__ magic method
         assert "Fe" in self.comps[0]
         assert "Fe" not in self.comps[2]
         assert Element("Fe") in self.comps[0]
@@ -62,6 +63,46 @@ def test_in(self):
         with pytest.raises(KeyError, match="Invalid key='Vac'"):
             self.comps[0]["Vac"]
 
+        # Test Species in Composition
+        comp = Composition({Species("Fe2+"): 2})
+        assert Species("Fe2+") in comp
+        assert Species("Fe3+") not in comp
+        assert "Fe" in comp
+        assert Element("Fe") in comp
+
+        # Test Element in Composition with Species
+        comp = Composition({Species("Fe2+"): 2})
+        assert Element("Fe") in comp
+        assert Element("O") not in comp
+        assert "Fe" in comp
+        assert "O" not in comp
+
+        # Test str in Composition with Element
+        comp = Composition({"Fe": 2})
+        assert "Fe" in comp
+        assert "O" not in comp
+        assert Element("Fe") in comp
+        assert Element("O") not in comp
+
+        # Test int (atomic number) in Composition
+        comp = Composition({Element("Fe"): 2})
+        assert 26 in comp  # atomic number for Fe
+        assert 8 not in comp  # atomic number for O
+
+        # Test float in Composition
+        comp = Composition({Element("Fe"): 2})
+        with pytest.raises(TypeError, match="expected string or bytes-like object"):
+            assert 1.5 in comp
+
+        # Test DummySpecies in Composition
+        comp = Composition({DummySpecies("X"): 2})
+        assert DummySpecies("X") in comp
+        assert DummySpecies("A") not in comp
+        assert "X" in comp
+        assert "Y" not in comp
+        assert Element("Fe") not in comp
+        assert Species("Fe2+") not in comp
+
     def test_hill_formula(self):
         c = Composition("CaCO3")
         assert c.hill_formula == "C Ca O3"
@@ -528,21 +569,12 @@ def test_oxi_state_guesses(self):
             {"V": 5, "O": -2},
         )
 
+        expected_oxi_guesses = dict(Li=1, Fe=2, P=5, O=-2)
         # max_sites for very large composition - should timeout if incorrect
-        assert Composition("Li10000Fe10000P10000O40000").oxi_state_guesses(max_sites=7)[0] == {
-            "Li": 1,
-            "Fe": 2,
-            "P": 5,
-            "O": -2,
-        }
+        assert Composition("Li10000Fe10000P10000O40000").oxi_state_guesses(max_sites=7)[0] == expected_oxi_guesses
 
         # max_sites for very large composition - should timeout if incorrect
-        assert Composition("Li10000Fe10000P10000O40000").oxi_state_guesses(max_sites=-1)[0] == {
-            "Li": 1,
-            "Fe": 2,
-            "P": 5,
-            "O": -2,
-        }
+        assert Composition("Li10000Fe10000P10000O40000").oxi_state_guesses(max_sites=-1)[0] == expected_oxi_guesses
 
         # negative max_sites less than -1 - should throw error if cannot reduce
         # to under the abs(max_sites) number of sites. Will also timeout if
@@ -572,7 +604,7 @@ def test_oxi_state_decoration(self):
         assert decorated.get(Species("Ni", 0)) == 1
         assert decorated.get(Species("Al", 0)) == 1
 
-    def test_Metallofullerene(self):
+    def test_metallofullerene(self):
         # Test: Parse Metallofullerene formula (e.g. Y3N@C80)
         formula = "Y3N@C80"
         sym_dict = {"Y": 3, "N": 1, "C": 80}

pymatgen/entries/compatibility.py

@@ -278,7 +278,7 @@ def get_correction(self, entry) -> ufloat:
                 correction += self.sulfide_correction[sf_type] * comp["S"]
 
         # Check for oxide, peroxide, superoxide, and ozonide corrections.
-        if True:
+        if Element("O") in comp:
             if self.correct_peroxide:
                 if entry.data.get("oxide_type"):
                     if entry.data["oxide_type"] in self.oxide_correction:

pymatgen/entries/tests/test_compatibility.py

@@ -1020,7 +1020,7 @@ def test_process_entry_with_oxidation_state(self):
         compat = MaterialsProject2020Compatibility(check_potcar=False)
         [processed_entry] = compat.process_entries(entry, clean=True, inplace=False)
 
-        assert len(processed_entry.energy_adjustments) == 1
+        assert len(processed_entry.energy_adjustments) == 2
 
 
 class MITCompatibilityTest(unittest.TestCase):

test_files/.pytest-split-durations

@@ -758,7 +758,7 @@
     "pymatgen/core/tests/test_bonds.py::FuncTest::test_obtain_all_bond_lengths": 0.0001646250020712614,
     "pymatgen/core/tests/test_composition.py::ChemicalPotentialTest::test_init": 0.00024487500195391476,
     "pymatgen/core/tests/test_composition.py::ChemicalPotentialTest::test_math": 0.0002614989934954792,
-    "pymatgen/core/tests/test_composition.py::CompositionTest::test_Metallofullerene": 0.0008804169774521142,
+    "pymatgen/core/tests/test_composition.py::CompositionTest::test_metallofullerene": 0.0008804169774521142,
     "pymatgen/core/tests/test_composition.py::CompositionTest::test_add": 0.0006662499945377931,
     "pymatgen/core/tests/test_composition.py::CompositionTest::test_almost_equals": 0.0006189589767018333,
     "pymatgen/core/tests/test_composition.py::CompositionTest::test_alphabetical_formula": 0.0006650839932262897,