diff --git a/dpdata/plugins/vasp.py b/dpdata/plugins/vasp.py index 84d3142b..a1014c8c 100644 --- a/dpdata/plugins/vasp.py +++ b/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 diff --git a/tests/test_vasp_outcar.py b/tests/test_vasp_outcar.py index ec514fe1..744a8605 100644 --- a/tests/test_vasp_outcar.py +++ b/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