ABINIT DFPT SHG

pyproject.toml

@@ -38,7 +38,8 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-abinit = ["abipy>=0.9.3"]
+#abinit = ["abipy>=0.9.3"]
+abinit = ["abipy @ git+https://github.com/abinit/abipy.git"] # tmp
 amset = ["amset>=0.4.15", "pydash"]
 cclib = ["cclib"]
 mp = ["mp-api>=0.37.5"]

src/atomate2/abinit/files.py

@@ -2,29 +2,44 @@
 
 from __future__ import annotations
 
+import contextlib
 import logging
 import os
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 from abipy.flowtk.utils import abi_extensions
+from jobflow.core.job import job
 from monty.serialization import loadfn
 
+from atomate2 import SETTINGS
 from atomate2.abinit.utils.common import INDIR_NAME
+from atomate2.common.files import delete_files, gzip_files
 from atomate2.utils.file_client import FileClient, auto_fileclient
+from atomate2.utils.path import strip_hostname
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
-    from pathlib import Path
 
     from abipy.abio.inputs import AbinitInput
     from pymatgen.core.structure import Structure
 
+    from atomate2.abinit.sets.anaddb import AnaddbInputGenerator
     from atomate2.abinit.sets.base import AbinitInputGenerator
+    from atomate2.abinit.sets.mrgddb import MrgddbInputGenerator
 
+__all__ = [
+    "del_gzip_files",
+    "fname2ext",
+    "load_abinit_input",
+    "out_to_in",
+    "write_abinit_input_set",
+    "write_anaddb_input_set",
+    "write_mrgddb_input_set",
+]
 
 logger = logging.getLogger(__name__)
 
-
 ALL_ABIEXTS = abi_extensions()
 
 
@@ -37,6 +52,10 @@ def fname2ext(filepath: Path | str) -> None | str:
     if "_" not in filename:
         return None
     ext = filename.split("_")[-1].replace(".nc", "")
+    if "1WF" in ext:
+        ext = "1WF"
+    if "DEN" in ext and ext[-1].isdigit():
+        ext = "1DEN"
     if ext not in ALL_ABIEXTS:
         return None
     return ext
@@ -98,10 +117,11 @@ def load_abinit_input(
     AbinitInput
         The AbinitInput object.
     """
+    dirpath = strip_hostname(dirpath)  # TODO: to FileCLient?
     abinit_input_file = os.path.join(dirpath, f"{fname}")
     if not os.path.exists(abinit_input_file):
         raise NotImplementedError(
-            f"Cannot load AbinitInput from directory without {fname} file."
+            f"Cannot load AbinitInput from directory {dirpath} without {fname} file."
         )
 
     return loadfn(abinit_input_file)
@@ -139,3 +159,131 @@ def write_abinit_input_set(
         raise RuntimeError("AbinitInputSet is not valid.")
 
     ais.write_input(directory=directory, make_dir=True, overwrite=False)
+
+
+def write_mrgddb_input_set(
+    input_set_generator: MrgddbInputGenerator,
+    prev_outputs: str | Path | list[str] | None = None,
+    directory: str | Path = ".",
+) -> None:
+    """Write the mrgddb input using a given generator.
+
+    Parameters
+    ----------
+    input_set_generator
+        The input generator used to write the mrgddb inputs.
+    prev_outputs
+        The list of previous directories needed for the calculation.
+    directory
+        Directory in which to write the abinit inputs.
+    """
+    mrgis = input_set_generator.get_input_set(
+        prev_outputs=prev_outputs,
+        workdir=directory,
+    )
+    if not mrgis.validate():
+        raise RuntimeError(
+            "MrgddbInputSet is not valid. Some previous outputs \
+        do not exist."
+        )
+
+    mrgis.write_input(directory=directory, make_dir=True, overwrite=False)
+
+
+def write_anaddb_input_set(
+    structure: Structure,
+    input_set_generator: AnaddbInputGenerator,
+    prev_outputs: str | Path | list[str] | None = None,
+    directory: str | Path = ".",
+) -> None:
+    """Write the anaddb input using a given generator.
+
+    Parameters
+    ----------
+    input_set_generator
+        The input generator used to write the anaddb inputs.
+    prev_outputs
+        The list of previous directories needed for the calculation.
+    directory
+        Directory in which to write the abinit inputs.
+    """
+    anais = input_set_generator.get_input_set(
+        structure=structure,
+        prev_outputs=prev_outputs,
+    )
+    if not anais.validate():
+        raise RuntimeError(
+            "AnaddbInputSet is not valid. Some previous outputs \
+        do not exist."
+        )
+
+    anais.write_input(directory=directory, make_dir=True, overwrite=False)
+
+
+# TODO: does not take the directories of the store_inputs
+@job
+def del_gzip_files(
+    outputs: list,
+    exclude_files_from_zip: list[str | Path] | None = None,
+    delete: bool = True,
+    exclude_files_from_del: list[str | Path] | None = None,
+    include_files_to_del: list[str | Path] | None = None,
+) -> None:
+    r"""Delete and gzip files from a previous Job or Flow.
+
+    Parameters
+    ----------
+    outputs
+        Outputs of the Job or Flow that needs its files deleted/compressed.
+    exclude_files_from_zip
+        Filenames to exclude from the compression.
+        Supports glob file matching, e.g., "\*.dat".
+    delete: bool = True,
+        Activates the deletion of the files or not.
+    exclude_files_from_del
+        Filenames to exclude from the compression.
+        Supports glob file matching, e.g., "\*.dat".
+    include_files_to_del
+        Filenames to include as a list of str or Path objects given relative to
+        directory. Glob file paths are supported, e.g. "\*.dat". If ``None``, all files
+        in the directory will be deleted.
+    """
+    dirs_to_clean = []
+    for o in outputs:
+        with contextlib.suppress(TypeError, AttributeError):
+            dirs_to_clean.append(o.dir_name)
+        with contextlib.suppress(TypeError, AttributeError, KeyError):
+            dirs_to_clean.extend(o["dirs"])
+        with contextlib.suppress(TypeError, AttributeError, KeyError):
+            dirs_to_clean.append(o["dir_name"])  # to zip run_rf and generate_perts
+        with contextlib.suppress(TypeError, AttributeError, KeyError):
+            dirs_to_clean.extend(
+                o.history_dirs
+            )  # to clean dirs unconverged and restarted
+
+    recursiv_dirs_to_clean = []
+    for dz_ in dirs_to_clean:
+        dz = strip_hostname(dz_)
+        recursiv_dirs_to_clean.append(Path(dz))
+        for root, _, _ in os.walk(dz):
+            recursiv_dirs_to_clean.append(Path(root))
+    # to take its own directory into account
+    recursiv_dirs_to_clean.append(Path(os.getcwd()))
+
+    if delete:
+        if include_files_to_del is None:
+            include_files_to_del = SETTINGS.ABINIT_FILES_TO_DEL
+        for d in recursiv_dirs_to_clean:
+            delete_files(
+                directory=strip_hostname(d),
+                include_files=include_files_to_del,
+                exclude_files=exclude_files_from_del,
+                allow_missing=True,
+            )
+
+    for d in recursiv_dirs_to_clean:
+        gzip_files(
+            directory=strip_hostname(d),
+            exclude_files=exclude_files_from_zip,
+            force=True,
+        )