Use sets to handle mixed valence species and add test

Author: kavanase

pymatgen/analysis/defects/generators.py

@@ -103,10 +103,10 @@ def generate(
         Returns:
             Generator[Vacancy, None, None]: Generator that yields a list of ``Vacancy`` objects.
         """
-        all_species = [*map(_element_str, structure.composition.elements)]
-        rm_species = all_species if rm_species is None else [*map(str, rm_species)]
+        all_species = {*map(_element_str, structure.composition.elements)}
+        rm_species = all_species if rm_species is None else {*map(str, rm_species)}
 
-        if not set(rm_species).issubset(all_species):
+        if not rm_species.issubset(all_species):
             msg = f"rm_species({rm_species}) must be a subset of the structure's species ({all_species})."
             raise ValueError(
                 msg,
@@ -235,7 +235,7 @@ def generate(
             structure: The bulk structure the anti-site defects are generated from.
             **kwargs: Additional keyword arguments for the ``Substitution.generate`` function.
         """
-        all_species = [*map(_element_str, structure.composition.elements)]
+        all_species = {*map(_element_str, structure.composition.elements)}
         subs = collections.defaultdict(list)
         for u, v in combinations(all_species, 2):
             subs[u].append(v)

tests/conftest.py

@@ -1,5 +1,6 @@
 from collections import defaultdict
 from pathlib import Path
+import numpy as np
 
 import pytest
 from monty.serialization import loadfn
@@ -22,6 +23,14 @@ def test_dir():
 def gan_struct(test_dir):
     return Structure.from_file(test_dir / "GaN.vasp")
 
+@pytest.fixture(scope="session")
+def mixed_valence_struct(test_dir):
+    return Structure(
+        lattice=np.array([[3.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 3.0]]),
+        species=["Cu+", "Cu+", "Cu2+", "O2-"],
+        coords=[[0.0, 0.0, 0.0], [0.5, 0.5, 0.5], [0.5, 0.5, 0.0], [0.5, 0.0, 0.5]],
+    )
+
 
 @pytest.fixture(scope="session")
 def stable_entries_Mg_Ga_N(test_dir):

tests/test_generators.py

@@ -54,6 +54,12 @@ def test_antisite_generator(gan_struct) -> None:
     def_names = [defect.name for defect in anti_gen]
     assert sorted(def_names) == ["Ga_N", "N_Ga"]
 
+def test_mixed_valence_antisite_generator(mixed_valence_struct) -> None:
+    anti_gen = AntiSiteGenerator().get_defects(mixed_valence_struct)
+    def_names = [defect.name for defect in anti_gen]
+    assert "Cu_Cu" not in def_names
+    assert set(def_names) == {"Cu_O", "O_Cu"}
+
 
 def test_interstitial_generator(gan_struct) -> None:
     gen = InterstitialGenerator().get_defects(