Updates for Vasprun with MD simulations (#3489)

* updates for vasp ML MD

* compress vasp md test file

* comment on MD tests

pymatgen/io/vasp/outputs.py

@@ -399,6 +399,8 @@ def _parse(self, stream, parse_dos, parse_eigen, parse_projected_eigen):
                         md_data[-1]["structure"] = self._parse_structure(elem)
                     elif tag == "varray" and elem.attrib.get("name") == "forces":
                         md_data[-1]["forces"] = _parse_vasp_array(elem)
+                    elif tag == "varray" and elem.attrib.get("name") == "stress":
+                        md_data[-1]["stress"] = _parse_vasp_array(elem)
                     elif tag == "energy":
                         d = {i.attrib["name"]: float(i.text) for i in elem.findall("i")}
                         if "kinetic" in d:
@@ -527,8 +529,13 @@ def converged_ionic(self) -> bool:
         Returns:
             bool: True if ionic step convergence has been reached, i.e. that vasp
                 exited before reaching the max ionic steps for a relaxation run.
+                In case IBRION=0 (MD) True if the max ionic steps are reached.
         """
         nsw = self.parameters.get("NSW", 0)
+        ibrion = self.parameters.get("IBRION", -1 if nsw in (-1, 0) else 0)
+        if ibrion == 0:
+            return nsw <= 1 or self.md_n_steps == nsw
+
         return nsw <= 1 or len(self.ionic_steps) < nsw
 
     @property
@@ -696,6 +703,14 @@ def is_spin(self) -> bool:
         """True if run is spin-polarized."""
         return self.parameters.get("ISPIN", 1) == 2
 
+    @property
+    def md_n_steps(self) -> int:
+        """Number of steps for md runs."""
+        # if ML enabled count all the actual MD steps
+        if self.md_data:
+            return len(self.md_data)
+        return self.nionic_steps
+
     def get_computed_entry(self, inc_structure=True, parameters=None, data=None, entry_id: str | None = None):
         """
         Returns a ComputedEntry or ComputedStructureEntry from the Vasprun.
@@ -1022,7 +1037,8 @@ def get_trajectory(self):
         from pymatgen.core.trajectory import Trajectory
 
         structs = []
-        for step in self.ionic_steps:
+        steps_list = self.md_data or self.ionic_steps
+        for step in steps_list:
             struct = step["structure"].copy()
             struct.add_site_property("forces", step["forces"])
             structs.append(struct)

tests/io/vasp/test_outputs.py

@@ -50,13 +50,26 @@
 
 class TestVasprun(PymatgenTest):
     def test_vasprun_ml(self):
+        # Test for ML MD simulation
+        # The trajectory data is stored in md_data
         vasp_run = Vasprun(f"{TEST_FILES_DIR}/vasprun.xml.ml_md")
         assert len(vasp_run.md_data) == 100
         for d in vasp_run.md_data:
             assert "structure" in d
             assert "forces" in d
             assert "energy" in d
         assert vasp_run.md_data[-1]["energy"]["total"] == approx(-491.51831988)
+        assert vasp_run.md_n_steps == 100
+        assert vasp_run.converged_ionic
+
+    def test_vasprun_md(self):
+        # Test for simple MD simulation (no ML).
+        # Does not generate the `md_data` attribute in Vasprun. Data based on `ionic_steps`
+        vasp_run = Vasprun(f"{TEST_FILES_DIR}/vasprun.xml.md.gz")
+        assert len(vasp_run.ionic_steps) == 10
+        assert vasp_run.final_energy == approx(-327.73014059)
+        assert vasp_run.md_n_steps == 10
+        assert vasp_run.converged_ionic
 
     def test_bad_random_seed(self):
         _ = Vasprun(f"{TEST_FILES_DIR}/vasprun.bad_random_seed.xml")