fix write Trajectory XDATACAR with variable lattice

materialsproject · Nov 25, 2021 · 9b8995f · 9b8995f
1 parent 37d8c1e
commit 9b8995f
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/pymatgen/core/tests/test_trajectory.py b/pymatgen/core/tests/test_trajectory.py
@@ -499,6 +499,13 @@ def test_variable_lattice(self):
             all(np.allclose(struct.lattice.matrix, structures[i].lattice.matrix) for i, struct in enumerate(traj))
         )
 
+        # Check if the file is written correctly when lattice is not constant.
+        traj.write_Xdatcar(filename="traj_test_XDATCAR")
+        # Load trajectory from written xdatcar and compare to original
+        written_traj = Trajectory.from_file("traj_test_XDATCAR", constant_lattice=False)
+        self._check_traj_equality(traj, written_traj)
+        os.remove("traj_test_XDATCAR")
+
     def test_to_from_dict(self):
         d = self.traj.as_dict()
         traj = Trajectory.from_dict(d)

diff --git a/pymatgen/core/trajectory.py b/pymatgen/core/trajectory.py
@@ -502,7 +502,7 @@ def write_Xdatcar(self, filename="XDATCAR", system=None, significant_figures=6):
 
         for si, frac_coords in enumerate(self.frac_coords):
             # Only print out the info block if
-            if self.constant_lattice and si == 0:
+            if si == 0 or not self.constant_lattice:
                 lines.extend([system, "1.0"])
 
                 if self.constant_lattice: