fix(vasp): collect OUTCAR files from multiple directories (#1034)

Author: njzjz-bot

dpdata/plugins/vasp.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -82,6 +83,36 @@ def to_system(self, data, frame_idx=0, **kwargs):
 @Format.register("outcar")
 @Format.register("vasp/outcar")
 class VASPOutcarFormat(Format):
+    def from_multi_systems(self, directory, **kwargs):
+        """Find conventionally named OUTCAR files below ``directory``.
+
+        VASP calculations are commonly stored one calculation per directory,
+        so the objects consumed by :meth:`from_labeled_system` are the OUTCAR
+        files themselves rather than the calculation directories.  Searching
+        recursively also supports grouping calculations below intermediate
+        directories such as workflow stages or temperatures.
+
+        Parameters
+        ----------
+        directory : str or os.PathLike
+            Root directory containing VASP calculation directories.
+        **kwargs : dict
+            Additional format options. They are consumed later when each
+            discovered OUTCAR is loaded.
+
+        Returns
+        -------
+        list[str]
+            Deterministically ordered paths to files named ``OUTCAR``.
+        """
+        outcar_files = []
+        for root, _, files in os.walk(directory):
+            # Match VASP's canonical output name so unrelated files in the
+            # calculation tree are not accidentally parsed as OUTCAR data.
+            if "OUTCAR" in files:
+                outcar_files.append(os.path.join(root, "OUTCAR"))
+        return sorted(outcar_files)
+
     @Format.post("rot_lower_triangular")
     def from_labeled_system(
         self, file_name, begin=0, step=1, convergence_check=True, **kwargs

tests/test_vasp_outcar.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
 import io
+import os
+import shutil
+import tempfile
 import unittest
 import warnings
 
@@ -202,6 +205,52 @@ def test(self):
         self.assertEqual(ss.get_nframes(), 2)
 
 
+class TestVaspOUTCARMultiSystems(unittest.TestCase):
+    def test_loads_outcars_recursively_from_nested_directories(self):
+        """The CLI ``-m`` path must discover OUTCAR files, not parse directories."""
+
+        def nonzero_composition(system):
+            """Normalize away MultiSystems' shared zero-count type entries."""
+            return frozenset(
+                (name, count)
+                for name, count in zip(
+                    system["atom_names"], system["atom_numbs"], strict=True
+                )
+                if count
+            )
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            calculation_dirs = [
+                os.path.join(tmpdir, "Bond_calc", "calculation-000"),
+                os.path.join(tmpdir, "heating", "temperature-300", "calculation-000"),
+            ]
+            source_outcars = [
+                os.path.join("poscars", "OUTCAR.Ge.vdw"),
+                os.path.join("poscars", "Ti-O-Ti-v6", "OUTCAR"),
+            ]
+            expected_compositions = {
+                nonzero_composition(dpdata.LabeledSystem(source, fmt="vasp/outcar"))
+                for source in source_outcars
+            }
+            for calculation_dir, source_outcar in zip(
+                calculation_dirs, source_outcars, strict=True
+            ):
+                os.makedirs(calculation_dir)
+                shutil.copy(source_outcar, os.path.join(calculation_dir, "OUTCAR"))
+
+            systems = dpdata.MultiSystems.from_file(tmpdir, fmt="vasp/outcar")
+
+        self.assertEqual(len(systems), 2)
+        self.assertEqual(systems.get_nframes(), 2)
+        self.assertEqual({system.get_nframes() for system in systems}, {1})
+        # MultiSystems aligns type maps and element order across systems, so
+        # compare the non-zero compositions rather than display formulas.
+        self.assertEqual(
+            {nonzero_composition(system) for system in systems},
+            expected_compositions,
+        )
+
+
 class TestVaspAtomNamesV6(unittest.TestCase):
     def test(self):
         # in vasp v6, the key TITEL is removed. check if the atom names