Add support for vaspout.h5, improvements to potcar handling

dev_scripts/potcar_scrambler.py

@@ -27,7 +27,7 @@ class PotcarScrambler:
 
     Used to generate copyright-compliant POTCARs for PMG tests.
 
-    In case of questions, contact Aaron Kaplan <[email protected]>.
+    In case of questions, contact Aaron Kaplan <[email protected]>.
 
     Recommended use:
         PotcarScrambler.from_file(
@@ -199,3 +199,21 @@ def potcar_cleanser() -> None:
 if __name__ == "__main__":
     potcar_cleanser()
     # generate_fake_potcar_libraries()
+
+    """
+    Note that vaspout.h5 files also contain full POTCARs. While the
+    Vaspout class in `pymatgen.io.vasp.outputs` contains a method to
+    replace the POTCAR with its spec (`remove_potcar_and_write_file`),
+    for test purposes, its often useful to have a fake POTCAR in place
+    of the real one.
+
+    To use the scrambler on a vaspout.h5:
+    ```
+    vout = Vaspout("< path to vaspout.h5>")
+    scrambled = PotcarScrambler(vout.potcar)
+    vout.remove_potcar_and_write_file(
+        filename = "< path to output vaspout.h5>",
+        fake_potcar_str = scrambled.scrambled_potcars_str
+    )
+    ```
+    """

src/pymatgen/io/vasp/__init__.py

@@ -17,6 +17,7 @@
     Oszicar,
     Outcar,
     Procar,
+    Vaspout,
     Vasprun,
     VolumetricData,
     Wavecar,

src/pymatgen/io/vasp/inputs.py

@@ -2391,22 +2391,28 @@ def data_stats(data_list: Sequence) -> dict:
 
         data_match_tol: float = 1e-6
         for ref_psp in possible_potcar_matches:
-            key_match = all(
-                set(ref_psp["keywords"][key]) == set(self._summary_stats["keywords"][key]) for key in ["header", "data"]
-            )
-
-            data_diff = [
-                abs(ref_psp["stats"][key][stat] - self._summary_stats["stats"][key][stat])
-                for stat in ["MEAN", "ABSMEAN", "VAR", "MIN", "MAX"]
-                for key in ["header", "data"]
-            ]
-            data_match = all(np.array(data_diff) < data_match_tol)
-
-            if key_match and data_match:
+            if self.compare_potcar_stats(ref_psp, self._summary_stats, tolerance=data_match_tol):
                 return True
 
         return False
 
+    def spec(self, extra_spec: Sequence[str] | None = None) -> dict[str, Any]:
+        """
+        POTCAR spec used in vasprun.xml.
+
+        Args:
+            extra_spec : Sequence[str] or None (default)
+                A list of extra POTCAR fields to include in the spec.
+                If None, defaults to no extra spec.
+        Returns:
+            dict of POTCAR spec
+        """
+        extra_spec = extra_spec or []
+        spec = {"titel": self.TITEL, "hash": self.md5_header_hash, "summary_stats": self._summary_stats}
+        for attr in extra_spec:
+            spec[attr] = getattr(self, attr, None)
+        return spec
+
     def write_file(self, filename: str) -> None:
         """Write PotcarSingle to a file.
 
@@ -2517,6 +2523,47 @@ def verify_potcar(self) -> tuple[bool, bool]:
 
         return has_sha256, hash_is_valid
 
+    @staticmethod
+    def compare_potcar_stats(
+        potcar_stats_1: dict,
+        potcar_stats_2: dict,
+        tolerance: float = 1.0e-6,
+        check_potcar_fields: Sequence[str] = ["header", "data"],
+    ) -> bool:
+        """
+        Compare PotcarSingle._summary_stats to assess if they are the same within a tolerance.
+
+        Args:
+            potcar_stats_1 : dict
+                Dict of potcar summary stats from the first PotcarSingle, from the PotcarSingle._summary stats attr
+            potcar_stats_2 : dict
+                Second dict of summary stats
+            tolerance : float = 1.e-6
+                Tolerance to assess equality of numeric statistical values
+            check_potcar_fields : Sequence[str] = ["header", "data"]
+                The specific fields of the POTCAR to check, whether just the "header", just the "data", or both
+
+        Returns:
+            bool
+                Whether the POTCARs are identical according to their summary stats.
+        """
+
+        key_match = all(
+            set(potcar_stats_1["keywords"].get(key)) == set(potcar_stats_2["keywords"].get(key))
+            for key in check_potcar_fields
+        )
+
+        data_match = False
+        if key_match:
+            data_diff = [
+                abs(potcar_stats_1["stats"].get(key, {}).get(stat) - potcar_stats_2["stats"].get(key, {}).get(stat))
+                for stat in ["MEAN", "ABSMEAN", "VAR", "MIN", "MAX"]
+                for key in check_potcar_fields
+            ]
+            data_match = all(np.array(data_diff) < tolerance)
+
+        return key_match and data_match
+
     def identify_potcar(
         self,
         mode: Literal["data", "file"] = "data",
@@ -2554,19 +2601,9 @@ def identify_potcar(
                 if self.VRHFIN.replace(" ", "") != ref_psp["VRHFIN"]:
                     continue
 
-                key_match = all(
-                    set(ref_psp["keywords"][key]) == set(self._summary_stats["keywords"][key]) for key in check_modes
-                )
-
-                data_diff = [
-                    abs(ref_psp["stats"][key][stat] - self._summary_stats["stats"][key][stat])
-                    for stat in ["MEAN", "ABSMEAN", "VAR", "MIN", "MAX"]
-                    for key in check_modes
-                ]
-
-                data_match = all(np.array(data_diff) < data_tol)
-
-                if key_match and data_match:
+                if self.compare_potcar_stats(
+                    ref_psp, self._summary_stats, tolerance=data_tol, check_potcar_fields=check_modes
+                ):
                     identity["potcar_functionals"].append(func)
                     identity["potcar_symbols"].append(ref_psp["symbol"])
 
@@ -2816,8 +2853,28 @@ def symbols(self, symbols: Sequence[str]) -> None:
 
     @property
     def spec(self) -> list[dict]:
-        """The atomic symbols and hash of all the atoms in the POTCAR file."""
-        return [{"symbol": psingle.symbol, "hash": psingle.md5_computed_file_hash} for psingle in self]
+        """
+        POTCAR spec for all POTCARs in this instance.
+
+        Args:
+            extra_spec : Sequence[str] or None (default)
+                A list of extra POTCAR fields to include in the spec.
+                If None, defaults to ["symbol"] (needed for compatibility with LOBSTER).
+
+        Return:
+            list[dict], a list of PotcarSingle.spec dicts
+        """
+        return [psingle.spec(extra_spec=["symbol"]) for psingle in self]
+
+    def write_potcar_spec(self, filename: str = "POTCAR.spec.json.gz") -> None:
+        """
+        Write POTCAR spec to file.
+
+        Args:
+            filename : str = "POTCAR.spec.json.gz"
+                The name of a file to write the POTCAR spec to.
+        """
+        dumpfn(self.spec, filename)
 
     def as_dict(self) -> dict:
         """MSONable dict representation."""
@@ -2840,22 +2897,19 @@ def from_dict(cls, dct: dict) -> Self:
         return Potcar(symbols=dct["symbols"], functional=dct["functional"])
 
     @classmethod
-    def from_file(cls, filename: PathLike) -> Self:
+    def from_str(cls, data: str):
         """
-        Reads Potcar from file.
+        Read Potcar from a string.
 
-        Args:
-            filename: Filename
+        :param data: Potcar as a string.
 
         Returns:
             Potcar
         """
-        with zopen(filename, mode="rt", encoding="utf-8") as file:
-            fdata = file.read()
-
         potcar = cls()
-        functionals: list[str | None] = []
-        for psingle_str in fdata.split("End of Dataset"):
+
+        functionals = []
+        for psingle_str in data.split("End of Dataset"):
             if p_strip := psingle_str.strip():
                 psingle = PotcarSingle(f"{p_strip}\nEnd of Dataset\n")
                 potcar.append(psingle)
@@ -2867,6 +2921,20 @@ def from_file(cls, filename: PathLike) -> Self:
         potcar.functional = functionals[0]
         return potcar
 
+    @classmethod
+    def from_file(cls, filename: str):
+        """
+        Reads Potcar from file.
+
+        :param filename: Filename
+
+        Returns:
+            Potcar
+        """
+        with zopen(filename, mode="rt", encoding="utf-8") as file:
+            fdata = file.read()
+        return cls.from_str(fdata)
+
     def write_file(self, filename: PathLike) -> None:
         """Write Potcar to a file.
 
@@ -2903,6 +2971,42 @@ def set_symbols(
         else:
             self.extend(PotcarSingle.from_symbol_and_functional(el, functional) for el in symbols)
 
+    @classmethod
+    def from_spec(cls, potcar_spec: list[dict], functionals: list[str] | None = None) -> Potcar:
+        """
+        Generate a POTCAR from a list of POTCAR spec dicts.
+
+        If a set of POTCARs *for the same functional* cannot be found, raises a ValueError.
+        Args:
+            potcar_spec: list[dict]
+                List of POTCAR specs, from Potcar.spec or [PotcarSingle.spec]
+            functionals : list[str] or None (default)
+                If a list of strings, the functionals to restrict the search to.
+
+        Returns:
+            Potcar, a POTCAR using a single functional that matches the input spec.
+        """
+
+        functionals = functionals or list(PotcarSingle._potcar_summary_stats)
+        for functional in functionals:
+            potcar = Potcar()
+            matched = [False for _ in range(len(potcar_spec))]
+            for ispec, spec in enumerate(potcar_spec):
+                titel = spec.get("titel", "")
+                titel_no_spc = titel.replace(" ", "")
+                symbol = titel.split(" ")[1].strip()
+
+                for stats in PotcarSingle._potcar_summary_stats[functional].get(titel_no_spc, []):
+                    if PotcarSingle.compare_potcar_stats(spec["summary_stats"], stats):
+                        potcar.append(PotcarSingle.from_symbol_and_functional(symbol=symbol, functional=functional))
+                        matched[ispec] = True
+                        break
+
+                if all(matched):
+                    return potcar
+
+        raise ValueError("Cannot match the give POTCAR spec to a set of POTCARs generated with the same functional.")
+
 
 class UnknownPotcarWarning(UserWarning):
     """Warning raised when POTCAR hashes do not pass validation."""