Fix Procar.get_projection_on_elements for structures with multiple same-element ionic sites (#3261)

Add test for get_projction_on_elements() when structure have multiple ionic sites with the same elements.

Signed-off-by: Naoto kawaguchi <[email protected]>

* split test_get_projection_on_elements() into separate test

---------

Signed-off-by: Naoto kawaguchi <[email protected]>
Co-authored-by: Janosh Riebesell <[email protected]>

Author: Na-Kawa

pymatgen/io/vasp/outputs.py

@@ -3864,7 +3864,7 @@ def get_projection_on_elements(self, structure: Structure):
             name = structure.species[iat].symbol
             for spin, d in self.data.items():
                 for k, b in itertools.product(range(self.nkpoints), range(self.nbands)):
-                    dico[spin][b][k][name] = np.sum(d[k, b, iat, :])
+                    dico[spin][b][k][name] += np.sum(d[k, b, iat, :])
 
         return dico
 

tests/io/vasp/test_outputs.py

@@ -1592,14 +1592,6 @@ def test_init(self):
         assert p.nbands == 10
         assert p.nkpoints == 10
         assert p.nions == 3
-        lat = Lattice.cubic(3.0)
-        struct = Structure(
-            lat,
-            ["Li", "Na", "K"],
-            [[0.0, 0.0, 0.0], [0.25, 0.25, 0.25], [0.75, 0.75, 0.75]],
-        )
-        d = p.get_projection_on_elements(struct)
-        assert d[Spin.up][2][2] == approx({"Na": 0.042, "K": 0.646, "Li": 0.042})
         filepath = f"{TEST_FILES_DIR}/PROCAR"
         p = Procar(filepath)
         assert p.get_occupation(0, "dxy")[Spin.up] == approx(0.96214813853000025)
@@ -1621,6 +1613,21 @@ def test_phase_factors(self):
         p = Procar(filepath)
         assert p.phase_factors[Spin.up][0, 0, 0, 0] == approx(-0.13 + 0.199j)
 
+    def test_get_projection_on_elements(self):
+        filepath = f"{TEST_FILES_DIR}/PROCAR.simple"
+        p = Procar(filepath)
+        struct = Structure(
+            Lattice.cubic(3.0),
+            ["Li", "Na", "K"],
+            [[0.0, 0.0, 0.0], [0.25, 0.25, 0.25], [0.75, 0.75, 0.75]],
+        )
+        d = p.get_projection_on_elements(struct)
+        assert d[Spin.up][2][2] == approx({"Na": 0.042, "K": 0.646, "Li": 0.042})
+        # https://github.com/materialsproject/pymatgen/pull/3261
+        struct.replace_species({"K": "Na"})
+        d2 = p.get_projection_on_elements(struct)
+        assert d2[Spin.up][2][2] == approx({"Na": 0.688, "Li": 0.042})
+
 
 class TestXdatcar(PymatgenTest):
     def test_init(self):