👌 IMPROVE: Write subset of trajectory frames to disk. (aiidaplugins#22)

Co-authored-by: Chris Sewell <[email protected]>

.github/workflows/tests.yml

@@ -90,7 +90,7 @@ jobs:
         pytest  --lammps-exec lmp_serial --cov=aiida_lammps --cov-report=xml --cov-report=term-missing
 
     - name: Upload to Codecov
-      if: matrix.os == 'ubuntu-latest' && matrix.python-version == 3.7
+      if: matrix.python-version == 3.8
       uses: codecov/codecov-action@v1
       with:
         name: pytests-lammps

aiida_lammps/data/trajectory.py

@@ -169,13 +169,19 @@ def get_step_data(self, step_idx):
         """Return parsed data, for a specific trajectory step."""
         return parse_step(self.get_step_string(step_idx).splitlines())
 
-    def iter_step_strings(self):
-        """Yield the content string, for a each trajectory step."""
+    def iter_step_strings(self, steps=None):
+        """Yield the content string, for each trajectory step."""
+
+        if steps is None:
+            steps = range(self.number_steps)
+        elif isinstance(steps, int):
+            steps = range(0, self.number_steps, steps)
+
         with self.open(self.get_attribute("traj_filename"), mode="rb") as handle:
             with ZipFile(
                 handle, "r", self.get_attribute("compression_method")
             ) as zip_file:
-                for step_idx in range(self.number_steps):
+                for step_idx in steps:
                     zip_name = "{}{}".format(self.get_attribute("zip_prefix"), step_idx)
                     with zip_file.open(zip_name) as step_file:
                         content = step_file.read()
@@ -213,8 +219,11 @@ def get_step_structure(
             original_structure=original_structure,
         )
 
-    def write_as_lammps(self, handle):
-        """Write out the lammps trajectory to file."""
-        for string in self.iter_step_strings():
+    def write_as_lammps(self, handle, steps=None):
+        """Write out the lammps trajectory to file.
+        :param handle: a file handle, opened in "wb" mode
+        :param steps: a list of steps to write (default to all)
+        """
+        for string in self.iter_step_strings(steps=steps):
             handle.write(string)
-            handle.write("\n")
+            handle.write(b"\n")

aiida_lammps/tests/test_trajectory.py

@@ -73,3 +73,10 @@ def test_lammpstraj_timesteps(db_test_app):
         900,
         1000,
     ]
+
+
+def test_write_as_lammps(db_test_app, tmp_path):
+    path = os.path.join(TEST_DIR, "input_files", "trajectory.lammpstrj")
+    data = LammpsTrajectory(path)
+    with tmp_path.joinpath("trajectory.lammpstrj").open(mode="wb") as handle:
+        data.write_as_lammps(handle)

codecov.yml

@@ -0,0 +1,10 @@
+coverage:
+  status:
+    project:
+      default:
+        target: 75%
+        threshold: 0.2%
+    patch:
+      default:
+        target: 70%
+        threshold: 0.2%