allow waveder to be stored

Author: jmmshn

.gitignore

@@ -63,3 +63,5 @@ docs/reference/atomate2.*
 .vscode/
 
 .DS_Store
+
+env/

src/atomate2/vasp/jobs/base.py

@@ -16,7 +16,7 @@
     BandStructureSymmLine,
 )
 from pymatgen.electronic_structure.dos import DOS, CompleteDos, Dos
-from pymatgen.io.vasp import Chgcar, Locpot, Wavecar
+from pymatgen.io.vasp import Chgcar, Locpot, Wavecar, Waveder
 
 from atomate2.vasp.files import copy_vasp_outputs, write_vasp_input_set
 from atomate2.vasp.run import run_vasp, should_stop_children
@@ -35,6 +35,7 @@
     Locpot,
     Chgcar,
     Wavecar,
+    Waveder,
     Trajectory,
     "force_constants",
     "normalmode_eigenvecs",

src/atomate2/vasp/schemas/calculation.py

@@ -25,6 +25,7 @@
     PotcarSingle,
     Vasprun,
     VolumetricData,
+    Waveder,
 )
 
 from atomate2 import SETTINGS
@@ -80,6 +81,7 @@ class VaspObject(ValueEnum):
     LOCPOT = "locpot"
     OPTIC = "optic"
     PROCAR = "procar"
+    WAVEDER = "waveder"
 
 
 class PotcarSpec(BaseModel):
@@ -601,6 +603,7 @@ def from_vasp_files(
             Tuple[str]
         ] = SETTINGS.VASP_STORE_VOLUMETRIC_DATA,
         store_trajectory: bool = False,
+        store_waveder: bool = False,
         vasprun_kwargs: Optional[Dict] = None,
     ) -> Tuple["Calculation", Dict[VaspObject, Dict]]:
         """
@@ -656,6 +659,9 @@ def from_vasp_files(
             This can help reduce the size of DOS objects in systems with many atoms.
         store_volumetric_data
             Which volumetric files to store.
+        store_waveder
+            Whether to store the contents of the binary WAVEDER file.
+            Which is used to compute frequency dependent dielectric functions.
         store_trajectory
             Whether to store the ionic steps in a pymatgen Trajectory object. if `True`,
             :obj:'.CalculationOutput.ionic_steps' is set to None to reduce duplicating
@@ -728,6 +734,12 @@ def from_vasp_files(
             )
             vasp_objects[VaspObject.TRAJECTORY] = traj  # type: ignore
 
+        if store_waveder:
+            wavder = _parse_waveder(dir_name)
+            if wavder is None:
+                raise RuntimeError(f"WAVEDER file not found in directory {dir_name}.")
+            vasp_objects[VaspObject.WAVEDER] = wavder  # type: ignore
+
         # MD run
         if vasprun.parameters.get("IBRION", -1) == 0:
             if vasprun.parameters.get("NSW", 0) == vasprun.nionic_steps:
@@ -828,6 +840,26 @@ def _get_volumetric_data(
     return volumetric_data
 
 
+def _parse_waveder(dir_name: Path) -> Optional[Waveder]:
+    """
+    Parse the WAVEDER file.
+
+    Parameters
+    ----------
+    dir_name
+        The directory containing the WAVEDER file.
+
+    Returns
+    -------
+    Optional[Waveder]
+        The WAVEDER data.
+    """
+    try:
+        return Waveder.from_binary(dir_name / "WAVEDER")
+    except Exception:
+        return None
+
+
 def _parse_dos(parse_mode: Union[str, bool], vasprun: Vasprun) -> Optional[Dos]:
     """Parse DOS. See Calculation.from_vasp_files for supported arguments."""
     nsw = vasprun.incar.get("NSW", 0)