From 300a33e2ab63c4b2a0fd3b7730bcb0909a6e9928 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Se=C3=A1n=20Kavanagh?=
 <51478689+kavanase@users.noreply.github.com>
Date: Tue, 31 Dec 2024 17:15:11 -0500
Subject: [PATCH] Fix DOS parsing for SOC calculations (#4239)

* Add informative error message to SymmetryUndeterminedError

* Add debug message for DOS parsing

* Fix `vasprun.xml` DOS parsing for SOC calculations

* Add test for `vasprun.xml` DOS parsing for SOC calculations

* pre-commit auto-fixes

* Typing fix

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 src/pymatgen/io/vasp/outputs.py |  5 +++++
 tests/io/vasp/test_outputs.py   | 13 ++++++++++++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/src/pymatgen/io/vasp/outputs.py b/src/pymatgen/io/vasp/outputs.py
index 4694183b8c5..2bb73cd89ac 100644
--- a/src/pymatgen/io/vasp/outputs.py
+++ b/src/pymatgen/io/vasp/outputs.py
@@ -1631,10 +1631,13 @@ def _parse_dos(elem: XML_Element) -> tuple[Dos, Dos, list[dict]]:
         tdensities = {}
         idensities = {}
 
+        soc_run = len(elem.find("total").find("array").find("set").findall("set")) > 2  # type: ignore[union-attr]
         for s in elem.find("total").find("array").find("set").findall("set"):  # type: ignore[union-attr]
             data = np.array(_parse_vasp_array(s))
             energies = data[:, 0]
             spin = Spin.up if s.attrib["comment"] == "spin 1" else Spin.down
+            if spin != Spin.up and soc_run:  # other 'spins' are x,y,z SOC projections
+                continue
             tdensities[spin] = data[:, 1]
             idensities[spin] = data[:, 2]
 
@@ -1649,6 +1652,8 @@ def _parse_dos(elem: XML_Element) -> tuple[Dos, Dos, list[dict]]:
 
                 for ss in s.findall("set"):
                     spin = Spin.up if ss.attrib["comment"] == "spin 1" else Spin.down
+                    if spin != Spin.up and soc_run:  # other 'spins' are x,y,z SOC projections
+                        continue
                     data = np.array(_parse_vasp_array(ss))
                     _n_row, n_col = data.shape
                     for col_idx in range(1, n_col):
diff --git a/tests/io/vasp/test_outputs.py b/tests/io/vasp/test_outputs.py
index 837884c22f0..17916375dae 100644
--- a/tests/io/vasp/test_outputs.py
+++ b/tests/io/vasp/test_outputs.py
@@ -54,6 +54,18 @@
 
 
 class TestVasprun(PymatgenTest):
+    def test_vasprun_soc(self):
+        # Test that SOC vaspruns are parsed appropriately, giving just Spin.Up tdos, idos and pdos
+        vasp_run = Vasprun(f"{VASP_OUT_DIR}/vasprun.int_Te_SOC.xml.gz")
+        dos_density_dicts_to_check = [vasp_run.complete_dos.densities, vasp_run.tdos.densities, vasp_run.idos.densities]
+        dos_density_dicts_to_check += [
+            densities for orbital_dict in vasp_run.complete_dos.pdos.values() for densities in orbital_dict.values()
+        ]
+        for i, dos_density_dict in enumerate(dos_density_dicts_to_check):
+            assert set(dos_density_dict.keys()) == {Spin.up}, f"Failed spin keys check for {i}th dos obj!"
+
+        assert vasp_run.complete_dos.spin_polarization is None
+
     def test_vasprun_ml(self):
         # Test for ML MD simulation
         # The trajectory data is stored in md_data
@@ -79,7 +91,6 @@ def test_vasprun_md(self):
     def test_vasprun_ediffg_set_to_0(self):
         # Test for case where EDIFFG is set to 0. This should pass if all ionic steps
         # complete and are electronically converged.
-        print(list(os.walk(VASP_OUT_DIR)))
         vasp_run = Vasprun(f"{VASP_OUT_DIR}/vasprun.ediffg_set_to_0.xml.gz")
         assert len(vasp_run.ionic_steps) == 3
         assert vasp_run.final_energy == approx(-34.60164204)