Add lattice velocities to Poscar (#3428)

* Add lattice velocities to Poscar

* use auto-cleaned PymatgenTest.tmp_path over tmp_file.unlink()

---------

Co-authored-by: Janosh Riebesell <janosh.riebesell@gmail.com>

pymatgen/io/vasp/inputs.py

@@ -89,6 +89,7 @@ def __init__(
         velocities: ArrayLike | None = None,
         predictor_corrector: ArrayLike | None = None,
         predictor_corrector_preamble: str | None = None,
+        lattice_velocities: ArrayLike | None = None,
         sort_structure: bool = False,
     ):
         """
@@ -108,6 +109,8 @@ def __init__(
                 Typically parsed in MD runs. Defaults to None.
             predictor_corrector_preamble (str | None, optional): Preamble to the predictor
                 corrector. Defaults to None.
+            lattice_velocities (ArrayLike | None, optional): Lattice velocities and current
+                lattice for the POSCAR. Available in MD runs with variable cell. Defaults to None.
             sort_structure (bool, optional): Whether to sort the structure. Useful if species
                 are not grouped properly together. Defaults to False.
         """
@@ -137,6 +140,9 @@ def __init__(
             self.comment = structure.formula if comment is None else comment
             if predictor_corrector_preamble:
                 self.structure.properties["predictor_corrector_preamble"] = predictor_corrector_preamble
+
+            if lattice_velocities is not None and np.any(lattice_velocities):
+                self.structure.properties["lattice_velocities"] = np.asarray(lattice_velocities)
         else:
             raise ValueError("Disordered structure with partial occupancies cannot be converted into POSCAR!")
 
@@ -162,6 +168,11 @@ def predictor_corrector_preamble(self):
         """Predictor corrector preamble in Poscar."""
         return self.structure.properties.get("predictor_corrector_preamble")
 
+    @property
+    def lattice_velocities(self):
+        """Lattice velocities in Poscar (including the current lattice vectors)."""
+        return self.structure.properties.get("lattice_velocities")
+
     @velocities.setter  # type: ignore
     def velocities(self, velocities):
         """Setter for Poscar.velocities."""
@@ -182,6 +193,11 @@ def predictor_corrector_preamble(self, predictor_corrector_preamble):
         """Setter for Poscar.predictor_corrector."""
         self.structure.properties["predictor_corrector"] = predictor_corrector_preamble
 
+    @lattice_velocities.setter  # type: ignore
+    def lattice_velocities(self, lattice_velocities: ArrayLike) -> None:
+        """Setter for Poscar.lattice_velocities."""
+        self.structure.properties["lattice_velocities"] = np.asarray(lattice_velocities)
+
     @property
     def site_symbols(self) -> list[str]:
         """
@@ -422,6 +438,15 @@ def from_str(data, default_names=None, read_velocities=True):
         )
 
         if read_velocities:
+            # Parse the lattice velocities and current lattice, if present.
+            # The header line should contain "Lattice velocities and vectors"
+            # There is no space between the coordinates and this section, so
+            # it appears in the lines of the first chunk
+            lattice_velocities = []
+            if len(lines) > ipos + n_sites + 1 and lines[ipos + n_sites + 1].lower().startswith("l"):
+                for line in lines[ipos + n_sites + 3 : ipos + n_sites + 9]:
+                    lattice_velocities.append([float(tok) for tok in line.split()])
+
             # Parse velocities if any
             velocities = []
             if len(chunks) > 1:
@@ -450,7 +475,7 @@ def from_str(data, default_names=None, read_velocities=True):
                     d3 = [float(tok) for tok in lines[st + 2 * n_sites].split()]
                     predictor_corrector.append([d1, d2, d3])
         else:
-            velocities = predictor_corrector = predictor_corrector_preamble = None
+            velocities = predictor_corrector = predictor_corrector_preamble = lattice_velocities = None
 
         return Poscar(
             struct,
@@ -460,6 +485,7 @@ def from_str(data, default_names=None, read_velocities=True):
             velocities=velocities,
             predictor_corrector=predictor_corrector,
             predictor_corrector_preamble=predictor_corrector_preamble,
+            lattice_velocities=lattice_velocities,
         )
 
     @np.deprecate(message="Use get_str instead")
@@ -514,6 +540,15 @@ def get_str(self, direct: bool = True, vasp4_compatible: bool = False, significa
             line += " " + site.species_string
             lines.append(line)
 
+        if self.lattice_velocities is not None:
+            try:
+                lines.append("Lattice velocities and vectors")
+                lines.append("  1")
+                for v in self.lattice_velocities:
+                    lines.append(" ".join(format_str.format(i) for i in v))
+            except Exception:
+                warnings.warn("Lattice velocities are missing or corrupted.")
+
         if self.velocities:
             try:
                 lines.append("")

tests/files/CONTCAR.MD.npt

@@ -0,0 +1,61 @@
+Si8                                     
+   1.00000000000000     
+     5.6029196758839159   -0.0603693499677626    0.1110429263825206
+    -0.0000000000000100    5.3747130151091689    0.0034014855406440
+     0.0000000000000068   -0.0000000000000120    5.4692974031883859
+   Si
+     8
+Direct
+  0.6539918994283490  0.8969172295968553  0.0676681342693018
+ -0.2312437400904284  0.5904002014657854  0.3939526918410268
+  0.8235113603041605  0.3274668347998835  0.7376879470910200
+  0.0234540394811329 -0.0006436675535249 -0.1173920341194730
+  0.4494229071161935  0.8354718770893421  0.6950789454078377
+  0.5785762144444008  0.4746645739302746  0.0005135261381124
+  0.0419959757024323  0.1291469345761019  0.3339785656099613
+  0.5094723570400191  0.1350290791632852  0.4745741352855470
+Lattice velocities and vectors
+           1
+  0.11376865E-02 -0.20054010E-02  0.10745440E-02
+ -0.80980926E-03 -0.54988058E-03 -0.11593411E-02
+  0.40755213E-03 -0.91838934E-03  0.10978311E-02
+  0.56062799E+01 -0.68862342E-01  0.11555075E+00
+ -0.98606785E-14  0.53730639E+01 -0.27835157E-02
+  0.64062547E-14 -0.12698742E-13  0.54725901E+01
+
+ -0.26486997E-01  0.15289665E-01 -0.24183306E-01
+ -0.21835373E-01 -0.48466524E-02  0.34963571E-02
+ -0.36389021E-02 -0.28571877E-02 -0.18926241E-02
+  0.49389009E-02  0.28594559E-02 -0.48140896E-02
+  0.20044174E-01 -0.32290274E-03  0.67317795E-02
+  0.51447245E-02  0.75356806E-02 -0.13128815E-01
+ -0.12148614E-01 -0.19846176E-01  0.16721524E-01
+ -0.73911890E-02  0.29807035E-02 -0.42831313E-02
+
+           1
+   3.00000000000000     
+  0.10000000E+01  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.63981833E+00  0.90527242E+00  0.54714689E-01
+  0.75707184E+00  0.58754437E+00  0.39611461E+00
+  0.82156413E+00  0.32584659E+00  0.73669073E+00
+  0.26096915E-01  0.98675463E-03  0.87991397E+00
+  0.46014883E+00  0.83542905E+00  0.69854272E+00
+  0.58132923E+00  0.47890733E+00  0.99326052E+00
+  0.35495080E-01  0.11798269E+00  0.34327666E+00
+  0.50551723E+00  0.13664264E+00  0.47231051E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00
+  0.00000000E+00  0.00000000E+00  0.00000000E+00

tests/io/vasp/test_inputs.py

@@ -4,7 +4,7 @@
 import os
 import pickle
 import unittest
-from pathlib import Path
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pytest
@@ -32,6 +32,9 @@
 )
 from pymatgen.util.testing import TEST_FILES_DIR, PymatgenTest
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
 
 class TestPoscar(PymatgenTest):
     def test_init(self):
@@ -265,24 +268,42 @@ def test_str(self):
     def test_from_md_run(self):
         # Parsing from an MD type run with velocities and predictor corrector data
         poscar = Poscar.from_file(f"{TEST_FILES_DIR}/CONTCAR.MD", check_for_potcar=False)
-        assert np.sum(np.array(poscar.velocities)) == approx(0.0065417961324)
+        assert np.sum(poscar.velocities) == approx(0.0065417961324)
         assert poscar.predictor_corrector[0][0][0] == 0.33387820e00
         assert poscar.predictor_corrector[0][1][1] == -0.10583589e-02
+        assert poscar.lattice_velocities is None
+
+        # Parsing from an MD type run with velocities, predictor corrector data and lattice velocities
+        poscar = Poscar.from_file(f"{TEST_FILES_DIR}/CONTCAR.MD.npt", check_for_potcar=False)
+        assert np.sum(poscar.velocities) == approx(-0.06193299494)
+        assert poscar.predictor_corrector[0][0][0] == 0.63981833
+        assert poscar.lattice_velocities.sum() == approx(16.49411358474)
 
     def test_write_md_poscar(self):
         # Parsing from an MD type run with velocities and predictor corrector data
         # And writing a new POSCAR from the new structure
         poscar = Poscar.from_file(f"{TEST_FILES_DIR}/CONTCAR.MD", check_for_potcar=False)
 
-        path = Path("POSCAR.testing.md")
+        path = f"{self.tmp_path}/POSCAR.testing.md"
+        poscar.write_file(path)
+        p3 = Poscar.from_file(path)
+
+        assert_allclose(poscar.structure.lattice.abc, p3.structure.lattice.abc, 5)
+        assert_allclose(poscar.velocities, p3.velocities, 5)
+        assert_allclose(poscar.predictor_corrector, p3.predictor_corrector, 5)
+        assert poscar.predictor_corrector_preamble == p3.predictor_corrector_preamble
+
+        # Same as above except also has lattice velocities
+        poscar = Poscar.from_file(f"{TEST_FILES_DIR}/CONTCAR.MD.npt", check_for_potcar=False)
+
         poscar.write_file(path)
         p3 = Poscar.from_file(path)
 
         assert_allclose(poscar.structure.lattice.abc, p3.structure.lattice.abc, 5)
         assert_allclose(poscar.velocities, p3.velocities, 5)
         assert_allclose(poscar.predictor_corrector, p3.predictor_corrector, 5)
         assert poscar.predictor_corrector_preamble == p3.predictor_corrector_preamble
-        path.unlink()
+        assert_allclose(poscar.lattice_velocities, p3.lattice_velocities, 5)
 
     def test_setattr(self):
         filepath = f"{TEST_FILES_DIR}/POSCAR"
@@ -395,7 +416,7 @@ def test_selective_dynamics(self):
         ]
 
 
-class TestIncar(unittest.TestCase):
+class TestIncar(PymatgenTest):
     def setUp(self):
         file_name = f"{TEST_FILES_DIR}/INCAR"
         self.incar = Incar.from_file(file_name)
@@ -564,11 +585,10 @@ def test_as_dict_and_from_dict(self):
         assert incar3["MAGMOM"] == [Magmom([1, 2, 3])]
 
     def test_write(self):
-        tempfname = Path("INCAR.testing")
-        self.incar.write_file(tempfname)
-        i = Incar.from_file(tempfname)
-        assert i == self.incar
-        tempfname.unlink()
+        tmp_file = f"{self.tmp_path}/INCAR.testing"
+        self.incar.write_file(tmp_file)
+        incar = Incar.from_file(tmp_file)
+        assert incar == self.incar
 
     def test_get_str(self):
         s = self.incar.get_str(pretty=True, sort_keys=True)