Optionally zip files at the end of jobs

src/atomate2/common/files.py

@@ -363,3 +363,31 @@ def get_zfile(
         return None
 
     raise FileNotFoundError(f"Could not find {base_name} or {base_name}.gz file.")
+
+
+def gzip_output_folder(
+    directory: str | Path, setting: bool | str, files_list: list[str]
+):
+    """
+    Zip the content of the output folder based on the specific code setting.
+
+    Parameters
+    ----------
+    directory: str or Path or None
+        Directory in which to gzip files.
+    setting: bool or str
+        the setting determining which files to zip. If True all the files in
+        the directory will be zipped, if "atomate" only the files in
+        files_list, if False no file will be zipped.
+    files_list: list of str
+        list of files to be zipped in case setting is "atomate"
+    """
+    if setting == "atomate":
+        gzip_files(
+            directory=directory,
+            include_files=files_list,
+            allow_missing=True,
+            force=True,
+        )
+    elif setting:
+        gzip_files(directory=directory, force=True)

src/atomate2/cp2k/jobs/base.py

@@ -8,7 +8,6 @@
 
 from jobflow import Maker, Response, job
 from monty.serialization import dumpfn
-from monty.shutil import gzip_dir
 from pymatgen.alchemy.materials import TransformedStructure
 from pymatgen.alchemy.transmuters import StandardTransmuter
 from pymatgen.core.trajectory import Trajectory
@@ -19,6 +18,8 @@
 from pymatgen.electronic_structure.dos import DOS, CompleteDos, Dos
 from pymatgen.io.common import VolumetricData
 
+from atomate2 import SETTINGS
+from atomate2.common.files import gzip_files, gzip_output_folder
 from atomate2.common.utils import get_transformations
 from atomate2.cp2k.files import (
     cleanup_cp2k_outputs,
@@ -46,6 +47,9 @@
 ]
 
 
+_FILES_TO_ZIP = ["cp2k.inp", "cp2k.out"]
+
+
 def cp2k_job(method: Callable):
     """
     Decorate the ``make`` method of CP2K job makers.
@@ -179,7 +183,16 @@
         cleanup_cp2k_outputs(directory=Path.cwd())
 
         # gzip folder
-        gzip_dir(".")
+        gzip_output_folder(
+            directory=Path.cwd(),
+            setting=SETTINGS.CP2K_ZIP_FILES,
+            files_list=_FILES_TO_ZIP,
+        )
+
+        if SETTINGS.CP2K_ZIP_FILES == "atomate":
+            gzip_files(include_files=_FILES_TO_ZIP, allow_missing=True, force=True)
+        elif SETTINGS.CP2K_ZIP_FILES:
+            gzip_files(force=True)
 
         return Response(
             stop_children=stop_children,

src/atomate2/lobster/jobs.py

@@ -7,12 +7,13 @@
 from pathlib import Path
 
 from jobflow import Maker, job
-from monty.shutil import gzip_dir
 from pymatgen.electronic_structure.cohp import CompleteCohp
 from pymatgen.electronic_structure.dos import LobsterCompleteDos
 from pymatgen.io.lobster import Lobsterin
 
-from atomate2.lobster.files import copy_lobster_files
+from atomate2 import SETTINGS
+from atomate2.common.files import gzip_output_folder
+from atomate2.lobster.files import LOBSTEROUTPUT_FILES, copy_lobster_files
 from atomate2.lobster.run import run_lobster
 from atomate2.lobster.schemas import LobsterTaskDocument
 
@@ -21,6 +22,9 @@
 logger = logging.getLogger(__name__)
 
 
+_FILES_TO_ZIP = [*LOBSTEROUTPUT_FILES, "lobsterin"]
+
+
 @dataclass
 class LobsterMaker(Maker):
     """
@@ -87,7 +91,11 @@ def make(
         run_lobster(**self.run_lobster_kwargs)
 
         # gzip folder
-        gzip_dir(".")
+        gzip_output_folder(
+            directory=Path.cwd(),
+            setting=SETTINGS.LOBSTER_ZIP_FILES,
+            files_list=_FILES_TO_ZIP,
+        )
 
         # parse lobster outputs
         return LobsterTaskDocument.from_directory(

src/atomate2/settings.py

@@ -1,7 +1,7 @@
 """Settings for atomate2."""
 
 from pathlib import Path
-from typing import Optional, Tuple, Union
+from typing import Literal, Optional, Tuple, Union
 
 from pydantic import BaseSettings, Field, root_validator
 
@@ -53,7 +53,7 @@ class Atomate2Settings(BaseSettings):
         description="Maximum volume change allowed in VASP relaxations before the "
         "calculation is tagged with a warning",
     )
-    VASP_HANDLE_UNSUCCESSFUL: Union[str, bool] = Field(
+    VASP_HANDLE_UNSUCCESSFUL: Union[bool, Literal["error"]] = Field(
         "error",
         description="Three-way toggle on what to do if the job looks OK but is actually"
         " unconverged (either electronic or ionic). - True: mark job as COMPLETED, but "
@@ -77,6 +77,13 @@ class Atomate2Settings(BaseSettings):
         "Requires the bader executable to be on the path.",
     )
 
+    VASP_ZIP_FILES: Union[bool, Literal["atomate"]] = Field(
+        "atomate",
+        description="Determine if the files in folder are being compressed. If True "
+        "all the files are compressed. If 'atomate' only a selection of files related "
+        "to the simulation will be compressed. If False no file is compressed.",
+    )
+
     LOBSTER_CMD: str = Field(
         default="lobster", description="Command to run standard version of VASP."
     )
@@ -85,6 +92,13 @@ class Atomate2Settings(BaseSettings):
         5, description="Maximum number of errors to correct before custodian gives up"
     )
 
+    LOBSTER_ZIP_FILES: Union[bool, Literal["atomate"]] = Field(
+        "atomate",
+        description="Determine if the files in folder are being compressed. If True "
+        "all the files are compressed. If 'atomate' only a selection of files related "
+        "to the simulation will be compressed. If False no file is compressed.",
+    )
+
     CP2K_CMD: str = Field(
         "cp2k.psmp", description="Command to run the MPI version of cp2k"
     )
@@ -125,6 +139,13 @@ class Atomate2Settings(BaseSettings):
         "parsing CP2K directories useful for storing duplicate of FW.json",
     )
 
+    CP2K_ZIP_FILES: Union[bool, Literal["atomate"]] = Field(
+        True,
+        description="Determine if the files in folder are being compressed. If True "
+        "all the files are compressed. If 'atomate' only a selection of files related "
+        "to the simulation will be compressed. If False no file is compressed.",
+    )
+
     # Elastic constant settings
     ELASTIC_FITTING_METHOD: str = Field(
         "finite_difference", description="Elastic constant fitting method"

src/atomate2/vasp/jobs/base.py

@@ -10,7 +10,6 @@
 from emmet.core.tasks import TaskDoc
 from jobflow import Maker, Response, job
 from monty.serialization import dumpfn
-from monty.shutil import gzip_dir
 from pymatgen.core.trajectory import Trajectory
 from pymatgen.electronic_structure.bandstructure import (
     BandStructure,
@@ -20,6 +19,7 @@
 from pymatgen.io.vasp import Chgcar, Locpot, Wavecar
 
 from atomate2 import SETTINGS
+from atomate2.common.files import gzip_output_folder
 from atomate2.vasp.files import copy_vasp_outputs, write_vasp_input_set
 from atomate2.vasp.run import run_vasp, should_stop_children
 from atomate2.vasp.sets.base import VaspInputGenerator
@@ -45,6 +45,69 @@
     "normalmode_eigenvecs",
 ]
 
+# Input files. Partially from https://www.vasp.at/wiki/index.php/Category:Input_files
+# Exclude those that are also outputs
+_INPUT_FILES = [
+    "DYNMATFULL",
+    "ICONST",
+    "INCAR",
+    "KPOINTS",
+    "KPOINTS OPT",
+    "ML_AB",
+    "ML_FF",
+    "PENALTYPOT",
+    "POSCAR",
+    "POTCAR",
+    "QPOINTS",
+]
+
+# Output files. Partially from https://www.vasp.at/wiki/index.php/Category:Output_files
+_OUTPUT_FILES = [
+    "AECCAR0",
+    "AECCAR1",
+    "AECCAR2",
+    "BSEFATBAND",
+    "CHG",
+    "CHGCAR",
+    "CONTCAR",
+    "DOSCAR",
+    "EIGENVAL",
+    "ELFCAR",
+    "HILLSPOT",
+    "IBZKPT",
+    "LOCPOT",
+    "ML_ABN",
+    "ML_FFN",
+    "ML_HIS",
+    "ML_LOGFILE",
+    "ML_REG",
+    "OSZICAR",
+    "OUTCAR",
+    "PARCHG",
+    "PCDAT",
+    "POT",
+    "PROCAR",
+    "PROOUT",
+    "REPORT",
+    "TMPCAR",
+    "vasprun.xml",
+    "vaspout.h5",
+    "vaspwave.h5",
+    "W*.tmp",
+    "WAVECAR",
+    "WAVEDER",
+    "WFULL*.tmp",
+    "XDATCAR",
+]
+
+# Files to zip: inputs, outputs and additionally generated files
+_FILES_TO_ZIP = (
+    _INPUT_FILES
+    + _OUTPUT_FILES
+    + [f"{name}.orig" for name in _INPUT_FILES]
+    + ["vasp.out", "custodian.json"]
+)
+
 
 def vasp_job(method: Callable):
     """
@@ -158,7 +221,11 @@ def make(self, structure: Structure, prev_vasp_dir: str | Path | None = None):
         stop_children = should_stop_children(task_doc, **self.stop_children_kwargs)
 
         # gzip folder
-        gzip_dir(".")
+        gzip_output_folder(
+            directory=Path.cwd(),
+            setting=SETTINGS.VASP_ZIP_FILES,
+            files_list=_FILES_TO_ZIP,
+        )
 
         return Response(
             stop_children=stop_children,

tests/vasp/conftest.py

@@ -34,8 +34,8 @@ def mock_vasp(monkeypatch, vasp_test_dir):
     reference files will be copied into the directory instead. As we do not want to
     test whether VASP is giving the correct output rather that the calculation inputs
     are generated correctly and that the outputs are parsed properly, this should be
-    sufficient for our needs. An other potential issue is that the POTCAR files
-    distribute with VASP are not present on the testing server due to licensing
+    sufficient for our needs. Another potential issue is that the POTCAR files
+    distributed with VASP are not present on the testing server due to licensing
     constraints. Accordingly, VaspInputSet.write_inputs will fail unless the
     "potcar_spec" option is set to True, in which case a POTCAR.spec file will be
     written instead. This fixture solves both of these issues.