store _POT files

Author: VicTrqt

src/atomate2/abinit/jobs/base.py

@@ -15,7 +15,7 @@
 from atomate2 import SETTINGS
 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.calculation import PotFileBStr, TaskState
 from atomate2.abinit.schemas.mrgddb import DdbFileStr
 from atomate2.abinit.schemas.task import AbinitTaskDoc
 from atomate2.abinit.utils.common import UnconvergedError
@@ -110,6 +110,7 @@ def setup_job(
     # VolumetricData,
     # Trajectory,
     DdbFileStr,
+    PotFileBStr,
 ]
 
 

src/atomate2/abinit/schemas/calculation.py

@@ -15,6 +15,7 @@
 from abipy.flowtk.utils import File
 from emmet.core.math import Matrix3D, Vector3D
 from jobflow.utils import ValueEnum
+from monty.json import MSONable
 from pydantic import BaseModel, Field
 from pymatgen.core import Structure
 from typing_extensions import Self
@@ -30,6 +31,27 @@
 logger = logging.getLogger(__name__)
 
 
+class PotFileBStr(MSONable):
+    """Object storing the raw binary string of a POT file."""
+
+    def __init__(
+        self, potfilepath: str | Path, potfilename: str, pot_as_bstr: bytes
+    ) -> None:
+        self.potfilepath = potfilepath
+        self.potfilename = potfilename
+        self.pot_as_bstr = pot_as_bstr
+
+    @classmethod
+    def from_potfilepath(cls, potfilepath: str | Path) -> Self:
+        """Create a PotFileBStr by directly reading the file."""
+        with open(potfilepath, "rb") as f:
+            pot_as_bstr = f.read()
+        potfilename = Path(potfilepath).name
+        return cls(
+            potfilepath=potfilepath, potfilename=potfilename, pot_as_bstr=pot_as_bstr
+        )
+
+
 class TaskState(ValueEnum):
     """Abinit calculation state."""
 
@@ -46,7 +68,8 @@ class AbinitObject(ValueEnum):
     ELECTRON_DENSITY = "electron_density"  # e_density
     WFN = "wfn"  # Wavefunction file
     TRAJECTORY = "trajectory"
-    DDBFILESTR = "ddbfilestr"  # DDB file as string
+    DDBFILESTR = "ddb"  # DDB file as string
+    POTFILEBSTR = "potential"  # POT file as b-string
 
 
 class CalculationOutput(BaseModel):
@@ -255,6 +278,7 @@ def from_abinit_files(
         abinit_abort_file: Path | str = MPIABORTFILE,
         abinit_out_file: Path | str = OUTPUT_FILE_NAME,
         abinit_outddb_file: Path | str = "out_DDB",
+        abinit_outpot_file: Path | str = "out_POT",
     ) -> tuple[Self, dict[AbinitObject, dict]]:
         """
         Create an Abinit calculation document from a directory and file paths.
@@ -283,13 +307,18 @@ def from_abinit_files(
         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_outpot_file = dir_name / abinit_outpot_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
             )
+        if abinit_outpot_file.exists():
+            abinit_objects[AbinitObject.POTFILEBSTR] = PotFileBStr.from_potfilepath(  # type: ignore[index]
+                abinit_outpot_file
+            )
 
         output_doc = None
         if abinit_out_file.exists():

src/atomate2/abinit/schemas/mrgddb.py

@@ -52,7 +52,7 @@ def from_ddbfile(cls, ddbfile: DdbFile) -> Self:
 class MrgddbObject(ValueEnum):
     """Types of Mrgddb data objects."""
 
-    DDBFILESTR = "ddbfilestr"  # DDB file as string
+    DDBFILESTR = "ddb"  # DDB file as string
 
 
 class CalculationOutput(BaseModel):

src/atomate2/abinit/schemas/task.py

@@ -381,6 +381,8 @@ def _get_task_files(files: list[Path], suffix: str = "") -> dict:
                 abinit_files["abinit_out_file"] = Path(file).relative_to(path)
             elif file.match(f"*outdata/out_DDB{suffix}*"):
                 abinit_files["abinit_outddb_file"] = Path(file).relative_to(path)
+            elif file.match(f"*outdata/out_POT{suffix}*"):
+                abinit_files["abinit_outpot_file"] = Path(file).relative_to(path)
 
         return abinit_files