store DDB + add LDA PSP

Author: VicTrqt

src/atomate2/abinit/jobs/base.py

@@ -16,12 +16,14 @@
 from atomate2.abinit.files import write_abinit_input_set
 from atomate2.abinit.run import run_abinit
 from atomate2.abinit.schemas.calculation import TaskState
+from atomate2.abinit.schemas.mrgddb import DdbFileStr
 from atomate2.abinit.schemas.task import AbinitTaskDoc
 from atomate2.abinit.utils.common import UnconvergedError
 from atomate2.abinit.utils.history import JobHistory
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
+    from typing import Callable
 
     from abipy.flowtk.events import AbinitCriticalWarning
     from pymatgen.core.structure import Structure
@@ -99,6 +101,53 @@ def setup_job(
     )
 
 
+_DATA_OBJECTS = [  # either str (TaskDoc fields) or MSONable class
+    # BandStructure,
+    # BandStructureSymmLine,
+    # DOS,
+    # Dos,
+    # CompleteDos,
+    # VolumetricData,
+    # Trajectory,
+    DdbFileStr,
+]
+
+
+def abinit_job(method: Callable) -> job:
+    """
+    Decorate the ``make`` method of CP2K job makers.
+
+    This is a thin wrapper around :obj:`~jobflow.core.job.job` that configures common
+    settings for all CP2K jobs. For example, it ensures that large data objects
+    (band structures, density of states, Cubes, etc) are all stored in the
+    atomate2 data store. It also configures the output schema to be a CP2K
+    :obj:`.TaskDocument`.
+
+    Any makers that return CP2K jobs (not flows) should decorate the ``make`` method
+    with @cp2k_job. For example:
+
+    .. code-block:: python
+
+        class MyCp2kMaker(BaseCp2kMaker):
+            @cp2k_job
+            def make(structure):
+                # code to run Cp2k job.
+                pass
+
+    Parameters
+    ----------
+    method : callable
+        A BaseCp2kMaker.make method. This should not be specified directly and is
+        implied by the decorator.
+
+    Returns
+    -------
+    callable
+        A decorated version of the make function that will generate Cp2k jobs.
+    """
+    return job(method, data=_DATA_OBJECTS, output_schema=AbinitTaskDoc)
+
+
 @dataclass
 class BaseAbinitMaker(Maker):
     """
@@ -136,7 +185,7 @@ def calc_type(self) -> str:
         """Get the type of calculation for this maker."""
         return self.input_set_generator.calc_type
 
-    @job
+    @abinit_job
     def make(
         self,
         structure: Structure | None = None,

src/atomate2/abinit/jobs/core.py

@@ -12,9 +12,8 @@
     RelaxConvergenceWarning,
     ScfConvergenceWarning,
 )
-from jobflow import Job, job
 
-from atomate2.abinit.jobs.base import BaseAbinitMaker
+from atomate2.abinit.jobs.base import BaseAbinitMaker, abinit_job
 from atomate2.abinit.sets.core import (
     LineNonSCFSetGenerator,
     NonSCFSetGenerator,
@@ -27,6 +26,7 @@
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
+    from jobflow import Job
     from pymatgen.core.structure import Structure
 
     from atomate2.abinit.sets.base import AbinitInputGenerator
@@ -114,7 +114,7 @@ class NonSCFMaker(BaseAbinitMaker):
         NscfConvergenceWarning,
     )
 
-    @job
+    @abinit_job
     def make(
         self,
         structure: Structure | None = None,

src/atomate2/abinit/jobs/mrgddb.py

@@ -16,10 +16,12 @@
 from atomate2.abinit.jobs.base import setup_job
 from atomate2.abinit.run import run_mrgddb
 from atomate2.abinit.schemas.calculation import TaskState
-from atomate2.abinit.schemas.mrgddb import MrgddbTaskDoc
+from atomate2.abinit.schemas.mrgddb import DdbFileStr, MrgddbTaskDoc
 from atomate2.abinit.sets.mrgddb import MrgddbInputGenerator
 
 if TYPE_CHECKING:
+    from typing import Callable
+
     from atomate2.abinit.utils.history import JobHistory
 
 logger = logging.getLogger(__name__)
@@ -28,6 +30,45 @@
     "MrgddbMaker",
 ]
 
+_MRGDDB_DATA_OBJECTS = [
+    DdbFileStr,
+]
+
+
+def mrgddb_job(method: Callable) -> job:
+    """
+    Decorate the ``make`` method of CP2K job makers.
+
+    This is a thin wrapper around :obj:`~jobflow.core.job.job` that configures common
+    settings for all CP2K jobs. For example, it ensures that large data objects
+    (band structures, density of states, Cubes, etc) are all stored in the
+    atomate2 data store. It also configures the output schema to be a CP2K
+    :obj:`.TaskDocument`.
+
+    Any makers that return CP2K jobs (not flows) should decorate the ``make`` method
+    with @cp2k_job. For example:
+
+    .. code-block:: python
+
+        class MyCp2kMaker(BaseCp2kMaker):
+            @cp2k_job
+            def make(structure):
+                # code to run Cp2k job.
+                pass
+
+    Parameters
+    ----------
+    method : callable
+        A BaseCp2kMaker.make method. This should not be specified directly and is
+        implied by the decorator.
+
+    Returns
+    -------
+    callable
+        A decorated version of the make function that will generate Cp2k jobs.
+    """
+    return job(method, data=_MRGDDB_DATA_OBJECTS, output_schema=MrgddbTaskDoc)
+
 
 @dataclass
 class MrgddbMaker(Maker):
@@ -54,7 +95,7 @@ def calc_type(self) -> str:
         """Get the type of calculation for this maker."""
         return self.input_set_generator.calc_type
 
-    @job
+    @mrgddb_job
     def make(
         self,
         prev_outputs: list[str] | None = None,

src/atomate2/abinit/jobs/response.py

@@ -14,7 +14,7 @@
 )
 from jobflow import Flow, Job, Response, job
 
-from atomate2.abinit.jobs.base import BaseAbinitMaker
+from atomate2.abinit.jobs.base import BaseAbinitMaker, abinit_job
 from atomate2.abinit.powerups import update_user_abinit_settings
 from atomate2.abinit.sets.response import (
     DdeSetGenerator,
@@ -65,7 +65,7 @@ class ResponseMaker(BaseAbinitMaker):
         ScfConvergenceWarning,
     )
 
-    @job
+    @abinit_job
     def make(
         self,
         structure: Structure | None = None,

src/atomate2/abinit/schemas/calculation.py

@@ -6,9 +6,10 @@
 import os
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 from abipy.abio.outputs import AbinitOutputFile
+from abipy.dfpt.ddb import DdbFile
 from abipy.electrons.gsr import GsrFile
 from abipy.flowtk import events
 from abipy.flowtk.utils import File
@@ -18,6 +19,7 @@
 from pymatgen.core import Structure
 from typing_extensions import Self
 
+from atomate2.abinit.schemas.mrgddb import DdbFileStr
 from atomate2.abinit.utils.common import (
     LOG_FILE_NAME,
     MPIABORTFILE,
@@ -44,6 +46,7 @@ class AbinitObject(ValueEnum):
     ELECTRON_DENSITY = "electron_density"  # e_density
     WFN = "wfn"  # Wavefunction file
     TRAJECTORY = "trajectory"
+    DDBFILESTR = "ddbfilestr"  # DDB file as string
 
 
 class CalculationOutput(BaseModel):
@@ -251,6 +254,7 @@ def from_abinit_files(
         abinit_log_file: Path | str = LOG_FILE_NAME,
         abinit_abort_file: Path | str = MPIABORTFILE,
         abinit_out_file: Path | str = OUTPUT_FILE_NAME,
+        abinit_outddb_file: Path | str = "out_DDB",
     ) -> tuple[Self, dict[AbinitObject, dict]]:
         """
         Create an Abinit calculation document from a directory and file paths.
@@ -278,6 +282,14 @@ def from_abinit_files(
         abinit_log_file = dir_name / abinit_log_file
         abinit_abort_file = dir_name / abinit_abort_file
         abinit_out_file = dir_name / abinit_out_file
+        abinit_outddb_file = dir_name / abinit_outddb_file
+
+        abinit_objects: dict[AbinitObject, Any] = {}
+        if abinit_outddb_file.exists():
+            abinit_outddb = DdbFile.from_file(abinit_outddb_file)
+            abinit_objects[AbinitObject.DDBFILESTR] = DdbFileStr.from_ddbfile(  # type: ignore[index]
+                abinit_outddb
+            )
 
         output_doc = None
         if abinit_out_file.exists():
@@ -329,4 +341,4 @@ def from_abinit_files(
             event_report=report,
         )
 
-        return instance, None  # abinit_objects,
+        return (instance, abinit_objects)