Improve onsite density matrix parsing and test

materialsproject · Jun 25, 2019 · 717e99f · 717e99f
1 parent 93380e5
commit 717e99f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/pymatgen/io/vasp/outputs.py b/pymatgen/io/vasp/outputs.py
@@ -1690,6 +1690,13 @@ def __init__(self, filename):
             self.read_nmr_efg()
             self.read_nmr_efg_tensor()
 
+        self.has_onsite_density_matrices = False
+        self.read_pattern({"has_onsite_density_matrices": r"onsite density matrix"},
+                          terminate_on_match=True)
+        if "has_onsite_density_matrices" in self.data:
+            self.has_onsite_density_matrices = True
+            self.read_onsite_density_matrices()
+
         # Store the individual contributions to the final total energy
         final_energy_contribs = {}
         for k in ["PSCENC", "TEWEN", "DENC", "EXHF", "XCENC", "PAW double counting",
@@ -2722,6 +2729,13 @@ def as_dict(self):
             d.update({"nmr_efg": {"raw": self.data["unsym_efg_tensor"],
                                   "parameters": self.data["efg"]}})
 
+        if self.has_onsite_density_matrices:
+            # cast Spin to str for consistency with electronic_structure
+            # TODO: improve handling of Enum (de)serialization in monty
+            onsite_density_matrices = [{str(k): v for k, v in d.items()}
+                                       for d in self.data["onsite_density_matrices"]]
+            d.update({"onsite_density_matrices": onsite_density_matrices})
+
         return d
 
     def read_fermi_contact_shift(self):

diff --git a/pymatgen/io/vasp/tests/test_outputs.py b/pymatgen/io/vasp/tests/test_outputs.py
@@ -927,11 +927,11 @@ def test_mag_electrostatic_error(self):
 
     def test_onsite_density_matrix(self):
         outcar = Outcar(self.TEST_FILES_DIR / "OUTCAR.LinearResponseU.gz")
-        outcar.read_onsite_density_matrices()
         matrices = outcar.data["onsite_density_matrices"]
         self.assertEqual(matrices[0][Spin.up][0][0], 1.0227)
         self.assertEqual(len(matrices[0][Spin.up]), 5)
         self.assertEqual(len(matrices[0][Spin.up][0]), 5)
+        self.assertTrue("onsite_density_matrices" in outcar.as_dict())
 
 
 class BSVasprunTest(PymatgenTest):