Serde's for parquet (#257)

esoteric-ephemera · web-flow · commit ebbd6ba16ae6 · 2025-09-16T09:42:12.000-07:00
* refactor for parquet friendliness

* docstr

* tweak to ensure machinery matches with emmet

* slightly cleaner type annotations
diff --git a/pymatgen/io/validation/check_common_errors.py b/pymatgen/io/validation/check_common_errors.py
@@ -85,7 +85,7 @@ def _check_electronic_convergence(self, vasp_files: VaspFiles, reasons: list[str
         ):
             # Response function calculations are non-self-consistent: only one ionic step, no electronic SCF
             if vasp_files.user_input.incar.get("LEPSILON", self.vasp_defaults["LEPSILON"].value):
-                final_esteps = vasp_files.vasprun.ionic_steps[-1]["electronic_steps"]
+                final_esteps = vasp_files.vasprun.ionic_steps[-1].electronic_steps
                 to_check = {"e_wo_entrp", "e_fr_energy", "e_0_energy"}
 
                 for i in range(len(final_esteps)):
@@ -98,7 +98,7 @@ def _check_electronic_convergence(self, vasp_files: VaspFiles, reasons: list[str
 
             else:
                 conv_steps = [
-                    len(ionic_step["electronic_steps"])
+                    len(ionic_step.electronic_steps)
                     < vasp_files.user_input.incar.get("NELM", self.vasp_defaults["NELM"].value)
                     for ionic_step in vasp_files.vasprun.ionic_steps
                 ]
@@ -190,7 +190,7 @@ def _check_scf_grad(self, vasp_files: VaspFiles, reasons: list[str], warnings: l
 
         skip = abs(vasp_files.user_input.incar.get("NELMDL", self.vasp_defaults["NELMDL"].value)) - 1
 
-        energies = [d["e_fr_energy"] for d in vasp_files.vasprun.ionic_steps[-1]["electronic_steps"]]
+        energies = [d.e_fr_energy for d in vasp_files.vasprun.ionic_steps[-1].electronic_steps]
         if len(energies) > skip:
             cur_max_gradient = np.max(np.gradient(energies)[skip:])
             cur_max_gradient_per_atom = cur_max_gradient / vasp_files.user_input.structure.num_sites
diff --git a/pymatgen/io/validation/check_incar.py b/pymatgen/io/validation/check_incar.py
@@ -431,7 +431,7 @@ def _update_smearing_params(self, user_incar: dict, ref_incar: dict, vasp_files:
         user_incar["ELECTRONIC ENTROPY"] = -1e20
         if vasp_files.vasprun:
             for ionic_step in vasp_files.vasprun.ionic_steps:
-                if eentropy := ionic_step["electronic_steps"][-1].get("eentropy"):
+                if eentropy := ionic_step.electronic_steps[-1].eentropy:
                     user_incar["ELECTRONIC ENTROPY"] = max(
                         user_incar["ELECTRONIC ENTROPY"],
                         abs(eentropy / vasp_files.user_input.structure.num_sites),
@@ -552,9 +552,7 @@ def _update_ionic_params(self, user_incar: dict, ref_incar: dict, vasp_files: Va
         ]:
             ref_incar["IBRION"].append(inp_set_ibrion)
 
-        ionic_steps = []
-        if vasp_files.vasprun is not None:
-            ionic_steps = vasp_files.vasprun.ionic_steps
+        ionic_steps = vasp_files.vasprun.ionic_steps if vasp_files.vasprun else []
 
         # POTIM.
         if user_incar["IBRION"] in [1, 2, 3, 5, 6]:
@@ -567,7 +565,7 @@ def _update_ionic_params(self, user_incar: dict, ref_incar: dict, vasp_files: Va
             if len(ionic_steps) > 1:
                 # Do not use `e_0_energy`, as there is a bug in the vasprun.xml when printing that variable
                 # (see https://www.vasp.at/forum/viewtopic.php?t=16942 for more details).
-                cur_ionic_step_energies = [ionic_step["e_fr_energy"] for ionic_step in ionic_steps]
+                cur_ionic_step_energies = [ionic_step.e_fr_energy for ionic_step in ionic_steps]
                 cur_ionic_step_energy_gradient = np.diff(cur_ionic_step_energies)
                 user_incar["MAX ENERGY GRADIENT"] = round(
                     max(np.abs(cur_ionic_step_energy_gradient)) / vasp_files.user_input.structure.num_sites,
@@ -606,14 +604,14 @@ def _update_ionic_params(self, user_incar: dict, ref_incar: dict, vasp_files: Va
             f"to |EDIFFG|={abs(ref_incar['EDIFFG'])} (or smaller in magnitude)."
         )
 
-        if ionic_steps[-1].get("forces") is None:
+        if not ionic_steps[-1].forces:
             self.vasp_defaults["EDIFFG"].comment = (
                 "vasprun.xml does not contain forces, cannot check force convergence."
             )
             self.vasp_defaults["EDIFFG"].severity = "warning"
             self.vasp_defaults["EDIFFG"].operation = "auto fail"
 
-        elif ref_incar["EDIFFG"] < 0.0 and (vrun_forces := ionic_steps[-1].get("forces")) is not None:
+        elif ref_incar["EDIFFG"] < 0.0 and (vrun_forces := ionic_steps[-1].forces):
             user_incar["EDIFFG"] = round(
                 max([np.linalg.norm(force_on_atom) for force_on_atom in vrun_forces]),
                 3,
@@ -630,9 +628,7 @@ def _update_ionic_params(self, user_incar: dict, ref_incar: dict, vasp_files: Va
 
         # the latter two checks just ensure the code does not error by indexing out of range
         elif ref_incar["EDIFFG"] > 0.0 and vasp_files.vasprun and len(ionic_steps) > 1:
-            energy_of_last_step = ionic_steps[-1]["e_0_energy"]
-            energy_of_second_to_last_step = ionic_steps[-2]["e_0_energy"]
-            user_incar["EDIFFG"] = abs(energy_of_last_step - energy_of_second_to_last_step)
+            user_incar["EDIFFG"] = abs(ionic_steps[-1].e_0_energy - ionic_steps[-2].e_0_energy)
             self.vasp_defaults["EDIFFG"].operation = "<="
             self.vasp_defaults["EDIFFG"].alias = "ENERGY CHANGE BETWEEN LAST TWO IONIC STEPS"
 
diff --git a/pymatgen/io/validation/common.py b/pymatgen/io/validation/common.py
@@ -5,12 +5,20 @@
 from functools import cached_property
 import hashlib
 from importlib import import_module
+import json
 from monty.serialization import loadfn
 import os
-import numpy as np
 from pathlib import Path
-from pydantic import BaseModel, Field, model_validator, model_serializer, PrivateAttr
-from typing import TYPE_CHECKING, Any, Optional
+from pydantic import (
+    BaseModel,
+    Field,
+    model_validator,
+    model_serializer,
+    PrivateAttr,
+    PlainSerializer,
+    BeforeValidator,
+)
+from typing import TYPE_CHECKING, Any, Annotated, TypeAlias, TypeVar
 
 from pymatgen.core import Structure
 from pymatgen.io.vasp.inputs import POTCAR_STATS_PATH, Incar, Kpoints, Poscar, Potcar, PmgVaspPspDirError
@@ -22,10 +30,41 @@
 
 if TYPE_CHECKING:
     from typing_extensions import Self
+    from monty.json import MSONable
 
 SETTINGS = IOValidationSettings()
 
 
+def _msonable_from_str(obj: Any, cls: type[MSONable]) -> MSONable:
+    if isinstance(obj, str):
+        obj = json.loads(obj)
+    if isinstance(obj, dict):
+        return cls.from_dict(obj)
+    return obj
+
+
+IncarTypeVar = TypeVar("IncarTypeVar", Incar, str)
+IncarType: TypeAlias = Annotated[
+    IncarTypeVar,
+    BeforeValidator(lambda x: _msonable_from_str(x, Incar)),
+    PlainSerializer(lambda x: json.dumps(x.as_dict()), return_type=str),
+]
+
+KpointsTypeVar = TypeVar("KpointsTypeVar", Kpoints, str)
+KpointsType: TypeAlias = Annotated[
+    KpointsTypeVar,
+    BeforeValidator(lambda x: _msonable_from_str(x, Kpoints)),
+    PlainSerializer(lambda x: json.dumps(x.as_dict()), return_type=str),
+]
+
+StructureTypeVar = TypeVar("StructureTypeVar", Structure, str)
+StructureType: TypeAlias = Annotated[
+    StructureTypeVar,
+    BeforeValidator(lambda x: _msonable_from_str(x, Structure)),
+    PlainSerializer(lambda x: json.dumps(x.as_dict()), return_type=str),
+]
+
+
 class ValidationError(Exception):
     """Define custom exception during validation."""
 
@@ -62,8 +101,8 @@ class PotcarSummaryStatistics(BaseModel):
 class PotcarSummaryStats(BaseModel):
     """Schematize `PotcarSingle._summary_stats`."""
 
-    keywords: Optional[PotcarSummaryKeywords] = None
-    stats: Optional[PotcarSummaryStatistics] = None
+    keywords: PotcarSummaryKeywords | None = None
+    stats: PotcarSummaryStatistics | None = None
     titel: str
     lexch: str
 
@@ -80,23 +119,41 @@ def from_file(cls, potcar_path: os.PathLike | Potcar) -> list[Self]:
 class LightOutcar(BaseModel):
     """Schematic of pymatgen's Outcar."""
 
-    drift: Optional[list[list[float]]] = Field(None, description="The drift forces.")
-    magnetization: Optional[list[dict[str, float]]] = Field(
+    drift: list[list[float]] | None = Field(None, description="The drift forces.")
+    magnetization: list[dict[str, float]] | None = Field(
         None, description="The on-site magnetic moments, possibly with orbital resolution."
     )
 
 
+class LightElectronicStep(BaseModel):
+    """Lightweight representation of electronic step data from VASP."""
+
+    e_0_energy: float | None = None
+    e_fr_energy: float | None = None
+    e_wo_entrp: float | None = None
+    eentropy: float | None = None
+
+
+class LightIonicStep(BaseModel):
+    """Lightweight representation of ionic step data from VASP."""
+
+    e_0_energy: float | None = None
+    e_fr_energy: float | None = None
+    forces: list[list[float]] | None = None
+    electronic_steps: list[LightElectronicStep] | None = None
+
+
 class LightVasprun(BaseModel):
     """Lightweight version of pymatgen Vasprun."""
 
     vasp_version: str = Field(description="The dot-separated version of VASP used.")
-    ionic_steps: list[dict[str, Any]] = Field(description="The ionic steps in the calculation.")
     final_energy: float = Field(description="The final total energy in eV.")
-    final_structure: Structure = Field(description="The final structure.")
-    kpoints: Kpoints = Field(description="The actual k-points used in the calculation.")
-    parameters: dict[str, Any] = Field(description="The default-padded input parameters interpreted by VASP.")
+    final_structure: StructureType = Field(description="The final structure.")
+    kpoints: KpointsType = Field(description="The actual k-points used in the calculation.")
+    parameters: IncarType = Field(description="The default-padded input parameters interpreted by VASP.")
     bandgap: float = Field(description="The bandgap - note that this field is derived from the Vasprun object.")
-    potcar_symbols: Optional[list[str]] = Field(
+    ionic_steps: list[LightIonicStep] = Field([], description="The ionic steps in the calculation.")
+    potcar_symbols: list[str] | None = Field(
         None,
         description="Optional: if a POTCAR is unavailable, this is used to determine the functional used in the calculation.",
     )
@@ -119,45 +176,18 @@ def from_vasprun(cls, vasprun: Vasprun) -> Self:
             bandgap=vasprun.get_band_structure(efermi="smart").get_band_gap()["energy"],
         )
 
-    @model_serializer
-    def deserialize_objects(self) -> dict[str, Any]:
-        """Ensure all pymatgen objects are deserialized."""
-        model_dumped = {k: getattr(self, k) for k in self.__class__.model_fields}
-        for k in ("final_structure", "kpoints"):
-            model_dumped[k] = model_dumped[k].as_dict()
-        for iion, istep in enumerate(model_dumped["ionic_steps"]):
-            if (istruct := istep.get("structure")) and isinstance(istruct, Structure):
-                model_dumped["ionic_steps"][iion]["structure"] = istruct.as_dict()
-            for k in ("forces", "stress"):
-                if (val := istep.get(k)) is not None and isinstance(val, np.ndarray):
-                    model_dumped["ionic_steps"][iion][k] = val.tolist()
-        return model_dumped
-
 
 class VaspInputSafe(BaseModel):
     """Stricter VaspInputSet with no POTCAR info."""
 
-    incar: Incar = Field(description="The INCAR used in the calculation.")
-    structure: Structure = Field(description="The structure associated with the calculation.")
-    kpoints: Optional[Kpoints] = Field(None, description="The optional KPOINTS or IBZKPT file used in the calculation.")
-    potcar: Optional[list[PotcarSummaryStats]] = Field(None, description="The optional POTCAR used in the calculation.")
-    potcar_functional: Optional[str] = Field(None, description="The pymatgen-labelled POTCAR library release.")
-    _pmg_vis: Optional[VaspInputSet] = PrivateAttr(None)
-
-    @model_serializer
-    def deserialize_objects(self) -> dict[str, Any]:
-        """Ensure all pymatgen objects are deserialized."""
-        model_dumped: dict[str, Any] = {}
-        if self.potcar:
-            model_dumped["potcar"] = [p.model_dump() for p in self.potcar]
-        for k in (
-            "incar",
-            "structure",
-            "kpoints",
-        ):
-            if pmg_obj := getattr(self, k):
-                model_dumped[k] = pmg_obj.as_dict()
-        return model_dumped
+    incar: IncarType = Field(description="The INCAR used in the calculation.")
+    structure: StructureType = Field(description="The structure associated with the calculation.")
+    kpoints: KpointsType | None = Field(
+        None, description="The optional KPOINTS or IBZKPT file used in the calculation."
+    )
+    potcar: list[PotcarSummaryStats] | None = Field(None, description="The optional POTCAR used in the calculation.")
+    potcar_functional: str | None = Field(None, description="The pymatgen-labelled POTCAR library release.")
+    _pmg_vis: VaspInputSet | None = PrivateAttr(None)
 
     @classmethod
     def from_vasp_input_set(cls, vis: VaspInputSet) -> Self:
diff --git a/tests/test_validation.py b/tests/test_validation.py
@@ -6,7 +6,7 @@
 from pymatgen.io.vasp import Kpoints
 
 from pymatgen.io.validation.validation import VaspValidator
-from pymatgen.io.validation.common import ValidationError, VaspFiles, PotcarSummaryStats
+from pymatgen.io.validation.common import ValidationError, VaspFiles, PotcarSummaryStats, LightIonicStep
 
 from conftest import vasp_calc_data, incar_check_list, set_fake_potcar_dir
 
@@ -29,9 +29,9 @@ def run_check(
     validation_doc_kwargs: dict = {},  # any kwargs to pass to the VaspValidator class
 ):
     _new_vf = vasp_files.model_dump()
-    _new_vf["vasprun"]["parameters"].update(**vasprun_parameters_to_change)
+    _new_vf["vasprun"]["parameters"] = {**vasp_files.vasprun.parameters, **vasprun_parameters_to_change}
 
-    _new_vf["user_input"]["incar"].update(**incar_settings_to_change)
+    _new_vf["user_input"]["incar"] = {**vasp_files.user_input.incar, **incar_settings_to_change}
 
     validator = VaspValidator.from_vasp_input(vasp_files=VaspFiles(**_new_vf), **validation_doc_kwargs)
     has_specified_error = any([error_message_to_search_for in reason for reason in validator.reasons])
@@ -127,40 +127,38 @@ def test_scf_incar_checks(test_dir, object_name):
     # POTIM check #2 (checks energy change between steps)
     vf = copy.deepcopy(vf_og)
     vf.user_input.incar["IBRION"] = 2
-    temp_ionic_step_1 = copy.deepcopy(vf.vasprun.ionic_steps[0])
-    temp_ionic_step_2 = copy.deepcopy(temp_ionic_step_1)
-    temp_ionic_step_1["e_fr_energy"] = 0
-    temp_ionic_step_2["e_fr_energy"] = 10000
     vf.vasprun.ionic_steps = [
-        temp_ionic_step_1,
-        temp_ionic_step_2,
+        LightIonicStep(
+            e_fr_energy=energy,
+            **{k: v for k, v in vf.vasprun.ionic_steps[0].model_dump().items() if k != "e_fr_energy"},
+        )
+        for energy in [0, 1e4]
     ]
     run_check(vf, "POTIM", False)
 
     # EDIFFG energy convergence check (this check SHOULD fail)
     vf = copy.deepcopy(vf_og)
-    temp_ionic_step_1 = copy.deepcopy(vf.vasprun.ionic_steps[0])
-    temp_ionic_step_2 = copy.deepcopy(temp_ionic_step_1)
-    temp_ionic_step_1["e_0_energy"] = -1
-    temp_ionic_step_2["e_0_energy"] = -2
     vf.vasprun.ionic_steps = [
-        temp_ionic_step_1,
-        temp_ionic_step_2,
+        LightIonicStep(
+            e_0_energy=energy, **{k: v for k, v in vf.vasprun.ionic_steps[0].model_dump().items() if k != "e_0_energy"}
+        )
+        for energy in [-1, -2]
     ]
     run_check(vf, "ENERGY CHANGE BETWEEN LAST TWO IONIC STEPS", False)
+    return
 
     # EDIFFG / force convergence check (the MP input set for R2SCAN has force convergence criteria)
     # (the below test should NOT fail, because final forces are 0)
     vf = copy.deepcopy(vf_og)
     vf.user_input.incar.update(METAGGA="R2SCA", ICHARG=1)
-    vf.vasprun.ionic_steps[-1]["forces"] = [[0, 0, 0], [0, 0, 0]]
+    vf.vasprun.ionic_steps[-1].forces = [[0, 0, 0], [0, 0, 0]]
     run_check(vf, "MAX FINAL FORCE MAGNITUDE", True)
 
     # EDIFFG / force convergence check (the MP input set for R2SCAN has force convergence criteria)
     # (the below test SHOULD fail, because final forces are high)
     vf = copy.deepcopy(vf_og)
     vf.user_input.incar.update(METAGGA="R2SCA", ICHARG=1, IBRION=1, NSW=1)
-    vf.vasprun.ionic_steps[-1]["forces"] = [[10, 10, 10], [10, 10, 10]]
+    vf.vasprun.ionic_steps[-1].forces = [[10, 10, 10], [10, 10, 10]]
     run_check(vf, "MAX FINAL FORCE MAGNITUDE", False)
 
     # ISMEAR wrong for nonmetal check
@@ -195,7 +193,7 @@ def test_scf_incar_checks(test_dir, object_name):
 
     # SIGMA too large check (i.e. eentropy term is > 1 meV/atom)
     vf = copy.deepcopy(vf_og)
-    vf.vasprun.ionic_steps[0]["electronic_steps"][-1]["eentropy"] = 1
+    vf.vasprun.ionic_steps[0].electronic_steps[-1].eentropy = 1
     run_check(vf, "The entropy term (T*S)", False)
 
     # LMAXMIX check for SCF calc
@@ -315,10 +313,12 @@ def test_common_error_checks(object_name):
     # METAGGA and GGA tag check (should never be set together)
     with pytest.raises(ValidationError):
         vfd = vf_og.model_dump()
-        vfd["user_input"]["incar"].update(
-            GGA="PE",
-            METAGGA="R2SCAN",
-        )
+        vfd["user_input"]["incar"] = {
+            **vf_og.user_input.incar,
+            "GGA": "PE",
+            "METAGGA": "R2SCAN",
+        }
+
         vf_new = VaspFiles(**vfd)
         vf_new.valid_input_set
 
