diff --git a/emmet-api/requirements/ubuntu-latest_py3.10_extras.txt b/emmet-api/requirements/ubuntu-latest_py3.10_extras.txt
index 01589a655d..2501ae65f7 100644
--- a/emmet-api/requirements/ubuntu-latest_py3.10_extras.txt
+++ b/emmet-api/requirements/ubuntu-latest_py3.10_extras.txt
@@ -336,7 +336,7 @@ pygments==2.18.0
     # via
     #   mkdocs-material
     #   rich
-pymatgen==2024.5.1
+pymatgen==2024.4.13
     # via
     #   emmet-core
     #   pymatgen-analysis-alloys
diff --git a/emmet-api/requirements/ubuntu-latest_py3.11_extras.txt b/emmet-api/requirements/ubuntu-latest_py3.11_extras.txt
index ab6a9da5fa..d2a0402a61 100644
--- a/emmet-api/requirements/ubuntu-latest_py3.11_extras.txt
+++ b/emmet-api/requirements/ubuntu-latest_py3.11_extras.txt
@@ -331,7 +331,7 @@ pygments==2.18.0
     # via
     #   mkdocs-material
     #   rich
-pymatgen==2024.5.1
+pymatgen==2024.4.13
     # via
     #   emmet-core
     #   pymatgen-analysis-alloys
diff --git a/emmet-api/requirements/ubuntu-latest_py3.9_extras.txt b/emmet-api/requirements/ubuntu-latest_py3.9_extras.txt
index 6680b8f845..f07b2dbdd5 100644
--- a/emmet-api/requirements/ubuntu-latest_py3.9_extras.txt
+++ b/emmet-api/requirements/ubuntu-latest_py3.9_extras.txt
@@ -346,7 +346,7 @@ pygments==2.18.0
     # via
     #   mkdocs-material
     #   rich
-pymatgen==2024.5.1
+pymatgen==2024.4.13
     # via
     #   emmet-core
     #   pymatgen-analysis-alloys
diff --git a/emmet-builders/emmet/builders/vasp/materials.py b/emmet-builders/emmet/builders/vasp/materials.py
index a9db1aeb6a..a518faa502 100644
--- a/emmet-builders/emmet/builders/vasp/materials.py
+++ b/emmet-builders/emmet/builders/vasp/materials.py
@@ -11,7 +11,7 @@
 from emmet.core.utils import group_structures, jsanitize, undeform_structure
 from emmet.core.vasp.calc_types import TaskType
 from emmet.core.vasp.material import MaterialsDoc
-from emmet.core.vasp.task_valid import TaskDocument
+from emmet.core.tasks import TaskDoc
 
 __author__ = "Shyam Dwaraknath <shyamd@lbl.gov>"
 
@@ -176,7 +176,7 @@ def get_items(self) -> Iterator[List[Dict]]:
             invalid_ids = set()
 
         projected_fields = [
-            "last_updated",
+            # "last_updated",
             "completed_at",
             "task_id",
             "formula_pretty",
@@ -190,9 +190,11 @@ def get_items(self) -> Iterator[List[Dict]]:
             "input.structure",
             # needed for entry from task_doc
             "output.energy",
+            "calcs_reversed.output.energy",
             "input.is_hubbard",
             "input.hubbards",
-            "input.potcar_spec",
+            "calcs_reversed.input.potcar_spec",
+            "calcs_reversed.output.structure",
             # needed for transform deformation structure back for grouping
             "transformations",
             # misc info for materials doc
@@ -222,7 +224,9 @@ def process_item(self, items: List[Dict]) -> List[Dict]:
                 were processed
         """
 
-        tasks = [TaskDocument(**task) for task in items]
+        tasks = [
+            TaskDoc(**task) for task in items
+        ]  # [TaskDoc(**task) for task in items]
         formula = tasks[0].formula_pretty
         task_ids = [task.task_id for task in tasks]
 
@@ -290,8 +294,8 @@ def update_targets(self, items: List[List[Dict]]):
             self.logger.info("No items to update")
 
     def filter_and_group_tasks(
-        self, tasks: List[TaskDocument], task_transformations: List[Union[Dict, None]]
-    ) -> Iterator[List[TaskDocument]]:
+        self, tasks: List[TaskDoc], task_transformations: List[Union[Dict, None]]
+    ) -> Iterator[List[TaskDoc]]:
         """
         Groups tasks by structure matching
         """
diff --git a/emmet-builders/emmet/builders/vasp/task_validator.py b/emmet-builders/emmet/builders/vasp/task_validator.py
index 3f7637686d..afa9b27571 100644
--- a/emmet-builders/emmet/builders/vasp/task_validator.py
+++ b/emmet-builders/emmet/builders/vasp/task_validator.py
@@ -5,8 +5,8 @@
 
 from emmet.builders.settings import EmmetBuildSettings
 from emmet.builders.utils import get_potcar_stats
+from emmet.core.tasks import TaskDoc
 from emmet.core.vasp.calc_types.enums import CalcType
-from emmet.core.vasp.task_valid import TaskDocument
 from emmet.core.vasp.validation import DeprecationMessage, ValidationDoc
 
 
@@ -39,7 +39,7 @@ def __init__(
         # Set up potcar cache if appropriate
         if self.settings.VASP_VALIDATE_POTCAR_STATS:
             if not self.potcar_stats:
-                self.potcar_stats = get_potcar_stats()
+                self.potcar_stats = get_potcar_stats(strict=False)
         else:
             self.potcar_stats = None
 
@@ -65,7 +65,7 @@ def unary_function(self, item):
         Args:
             item (dict): a (projection of a) task doc
         """
-        task_doc = TaskDocument(**item)
+        task_doc = TaskDoc(**item)
         validation_doc = ValidationDoc.from_task_doc(
             task_doc=task_doc,
             kpts_tolerance=self.settings.VASP_KPTS_TOLERANCE,
diff --git a/emmet-builders/requirements/ubuntu-latest_py3.10_extras.txt b/emmet-builders/requirements/ubuntu-latest_py3.10_extras.txt
index a93aa1d7bc..6b71fbadb3 100644
--- a/emmet-builders/requirements/ubuntu-latest_py3.10_extras.txt
+++ b/emmet-builders/requirements/ubuntu-latest_py3.10_extras.txt
@@ -440,7 +440,7 @@ pygments==2.18.0
     # via
     #   mkdocs-material
     #   rich
-pymatgen==2024.5.1
+pymatgen==2024.4.13
     # via
     #   chgnet
     #   emmet-core
diff --git a/emmet-builders/requirements/ubuntu-latest_py3.11_extras.txt b/emmet-builders/requirements/ubuntu-latest_py3.11_extras.txt
index 45aab844a1..c1c66618f3 100644
--- a/emmet-builders/requirements/ubuntu-latest_py3.11_extras.txt
+++ b/emmet-builders/requirements/ubuntu-latest_py3.11_extras.txt
@@ -434,7 +434,7 @@ pygments==2.18.0
     # via
     #   mkdocs-material
     #   rich
-pymatgen==2024.5.1
+pymatgen==2024.4.13
     # via
     #   chgnet
     #   emmet-core
diff --git a/emmet-builders/requirements/ubuntu-latest_py3.9_extras.txt b/emmet-builders/requirements/ubuntu-latest_py3.9_extras.txt
index 6407874eec..e977526c03 100644
--- a/emmet-builders/requirements/ubuntu-latest_py3.9_extras.txt
+++ b/emmet-builders/requirements/ubuntu-latest_py3.9_extras.txt
@@ -452,7 +452,7 @@ pygments==2.18.0
     # via
     #   mkdocs-material
     #   rich
-pymatgen==2024.5.1
+pymatgen==2024.4.13
     # via
     #   chgnet
     #   emmet-core
diff --git a/emmet-builders/tests/test_corrected_entries_thermo.py b/emmet-builders/tests/test_corrected_entries_thermo.py
index 620e662125..f7f03af057 100644
--- a/emmet-builders/tests/test_corrected_entries_thermo.py
+++ b/emmet-builders/tests/test_corrected_entries_thermo.py
@@ -34,7 +34,7 @@ def thermo_store():
 
 @pytest.fixture
 def phase_diagram_store():
-    return MemoryStore(key="chemsys")
+    return MemoryStore(key="phase_diagram_id")
 
 
 def test_corrected_entries_builder(corrected_entries_store, materials_store):
diff --git a/emmet-core/emmet/core/base.py b/emmet-core/emmet/core/base.py
index 0492dcb447..9985d80fb0 100644
--- a/emmet-core/emmet/core/base.py
+++ b/emmet-core/emmet/core/base.py
@@ -11,6 +11,7 @@
 
 from emmet.core import __version__
 from emmet.core.common import convert_datetime
+from emmet.core.utils import utcnow
 
 T = TypeVar("T", bound="EmmetBaseModel")
 
@@ -34,7 +35,7 @@ class EmmetMeta(BaseModel):
     )
 
     build_date: Optional[datetime] = Field(  # type: ignore
-        default_factory=datetime.utcnow,
+        default_factory=utcnow,
         description="The build date for this document.",
     )
 
diff --git a/emmet-core/emmet/core/tasks.py b/emmet-core/emmet/core/tasks.py
index dfa0b89f47..158a3c7972 100644
--- a/emmet-core/emmet/core/tasks.py
+++ b/emmet-core/emmet/core/tasks.py
@@ -9,9 +9,19 @@
 
 import numpy as np
 from emmet.core.common import convert_datetime
+from emmet.core.mpid import MPID
 from emmet.core.structure import StructureMetadata
-from emmet.core.vasp.calc_types import CalcType, TaskType
+from emmet.core.utils import utcnow
+from emmet.core.vasp.calc_types import (
+    CalcType,
+    calc_type,
+    TaskType,
+    run_type,
+    RunType,
+    task_type,
+)
 from emmet.core.vasp.calculation import (
+    CalculationBaseModel,
     Calculation,
     PotcarSpec,
     RunStatistics,
@@ -20,7 +30,14 @@
 from emmet.core.vasp.task_valid import TaskState
 from monty.json import MontyDecoder
 from monty.serialization import loadfn
-from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    field_validator,
+    model_validator,
+    PrivateAttr,
+)
 from pymatgen.analysis.structure_analyzer import oxide_type
 from pymatgen.core.structure import Structure
 from pymatgen.core.trajectory import Trajectory
@@ -45,7 +62,7 @@ class Potcar(BaseModel):
     )
 
 
-class OrigInputs(BaseModel):
+class OrigInputs(CalculationBaseModel):
     incar: Optional[Union[Incar, Dict]] = Field(
         None,
         description="Pymatgen object representing the INCAR file.",
@@ -84,7 +101,7 @@ class OutputDoc(BaseModel):
     )
 
     density: Optional[float] = Field(None, description="Density of in units of g/cc.")
-    energy: float = Field(..., description="Total Energy in units of eV.")
+    energy: Optional[float] = Field(None, description="Total Energy in units of eV.")
     forces: Optional[List[List[float]]] = Field(
         None, description="The force on each atom in units of eV/A^2."
     )
@@ -337,11 +354,11 @@ class TaskDoc(StructureMetadata, extra="allow"):
         None, description="Final output structure from the task"
     )
 
-    task_type: Optional[Union[CalcType, TaskType]] = Field(
+    task_type: Optional[Union[TaskType, CalcType]] = Field(
         None, description="The type of calculation."
     )
 
-    task_id: Optional[str] = Field(
+    task_id: Optional[Union[MPID, str]] = Field(
         None,
         description="The (task) ID of this calculation, used as a universal reference across property documents."
         "This comes in the form: mp-******.",
@@ -401,10 +418,44 @@ class TaskDoc(StructureMetadata, extra="allow"):
     )
 
     last_updated: Optional[datetime] = Field(
-        None,
+        utcnow(),
         description="Timestamp for the most recent calculation for this task document",
     )
 
+    batch_id: Optional[str] = Field(
+        None,
+        description="Identifier for this calculation; should provide rough information about the calculation origin and purpose.",
+    )
+
+    # Note that these private fields are needed because TaskDoc permits extra info
+    # added to the model, unlike TaskDocument. Because of this, when pydantic looks up
+    # attrs on the model, it searches for them in the model extra dict first, and if it
+    # can't find them, throws an AttributeError. It does this before looking to see if the
+    # class has that attr defined on it.
+
+    # Private field used to store output of `TaskDoc.calc_type`
+    _calc_type: Optional[CalcType] = PrivateAttr(None)
+    # Private field used to store output of `TaskDoc.run_type`
+    _run_type: Optional[RunType] = PrivateAttr(None)
+    # Private field used to store output of `TaskDoc.structure_entry`.
+    _structure_entry: Optional[ComputedStructureEntry] = PrivateAttr(None)
+
+    def model_post_init(self, __context: Any) -> None:
+        # Needed for compatibility with TaskDocument
+        if self.task_type is None:
+            self.task_type = task_type(self.orig_inputs)
+
+        if isinstance(self.task_type, CalcType):
+            # For a while, the TaskDoc.task_type was allowed to be a CalcType or TaskType
+            # For backwards compatibility with TaskDocument, ensure that isinstance(TaskDoc.task_type, TaskType)
+            temp = str(self.task_type).split(" ")
+            self._run_type = RunType(temp[0])
+            self.task_type = TaskType(" ".join(temp[1:]))
+
+        # TODO: remove after imposing TaskDoc schema on older tasks in collection
+        if self.structure is None:
+            self.structure = self.calcs_reversed[0].output.structure
+
     # Make sure that the datetime field is properly formatted
     # (Unclear when this is not the case, please leave comment if observed)
     @field_validator("last_updated", mode="before")
@@ -412,9 +463,26 @@ class TaskDoc(StructureMetadata, extra="allow"):
     def last_updated_dict_ok(cls, v) -> datetime:
         return convert_datetime(cls, v)
 
+    @field_validator("batch_id", mode="before")
+    @classmethod
+    def _validate_batch_id(cls, v) -> str:
+        if v is not None:
+            invalid_chars = set(
+                char for char in v if (not char.isalnum()) or (char not in {"-", "_"})
+            )
+            if len(invalid_chars) > 0:
+                raise ValueError(
+                    f"Invalid characters in batch_id:\n{' '.join(invalid_chars)}"
+                )
+        return v
+
     @model_validator(mode="after")
     def set_entry(self) -> datetime:
-        if not self.entry and self.calcs_reversed:
+        if (
+            not self.entry
+            and self.calcs_reversed
+            and getattr(self.calcs_reversed[0].output, "structure", None)
+        ):
             self.entry = self.get_entry(self.calcs_reversed, self.task_id)
         return self
 
@@ -426,6 +494,7 @@ def from_directory(
         store_additional_json: bool = True,
         additional_fields: Optional[Dict[str, Any]] = None,
         volume_change_warning_tol: float = 0.2,
+        task_names: Optional[list[str]] = None,
         **vasp_calculation_kwargs,
     ) -> _T:
         """
@@ -444,6 +513,9 @@ def from_directory(
         volume_change_warning_tol
             Maximum volume change allowed in VASP relaxations before the calculation is
             tagged with a warning.
+        task_names
+            Naming scheme for multiple calculations in on folder e.g. ["relax1","relax2"].
+            Can be subfolder or extension.
         **vasp_calculation_kwargs
             Additional parsing options that will be passed to the
             :obj:`.Calculation.from_vasp_files` function.
@@ -457,7 +529,9 @@ def from_directory(
 
         additional_fields = {} if additional_fields is None else additional_fields
         dir_name = Path(dir_name)
-        task_files = _find_vasp_files(dir_name, volumetric_files=volumetric_files)
+        task_files = _find_vasp_files(
+            dir_name, volumetric_files=volumetric_files, task_names=task_names
+        )
 
         if len(task_files) == 0:
             raise FileNotFoundError("No VASP files found!")
@@ -600,7 +674,7 @@ def from_vasprun(
 
     @staticmethod
     def get_entry(
-        calcs_reversed: List[Calculation], task_id: Optional[str] = None
+        calcs_reversed: List[Calculation], task_id: Optional[Union[MPID, str]] = None
     ) -> ComputedEntry:
         """
         Get a computed entry from a list of VASP calculation documents.
@@ -624,7 +698,10 @@ def get_entry(
             "energy": calcs_reversed[0].output.energy,
             "parameters": {
                 # Cannot be PotcarSpec document, pymatgen expects a dict
-                "potcar_spec": [dict(d) for d in calcs_reversed[0].input.potcar_spec],
+                # Note that `potcar_spec` is optional
+                "potcar_spec": [dict(d) for d in calcs_reversed[0].input.potcar_spec]
+                if calcs_reversed[0].input.potcar_spec
+                else [],
                 # Required to be compatible with MontyEncoder for the ComputedEntry
                 "run_type": str(calcs_reversed[0].run_type),
                 "is_hubbard": calcs_reversed[0].input.is_hubbard,
@@ -633,11 +710,40 @@ def get_entry(
             "data": {
                 "oxide_type": oxide_type(calcs_reversed[0].output.structure),
                 "aspherical": calcs_reversed[0].input.parameters.get("LASPH", False),
-                "last_updated": str(datetime.utcnow()),
+                "last_updated": str(utcnow()),
             },
         }
         return ComputedEntry.from_dict(entry_dict)
 
+    @property
+    def calc_type(self) -> CalcType:
+        """
+        Get the calc type for this TaskDoc.
+
+        Returns
+        --------
+        CalcType
+            The type of calculation.
+        """
+        inputs = (
+            self.calcs_reversed[0].input.model_dump()
+            if len(self.calcs_reversed) > 0
+            else self.orig_inputs
+        )
+        params = self.calcs_reversed[0].input.parameters
+        incar = self.calcs_reversed[0].input.incar
+
+        self._calc_type = calc_type(inputs, {**params, **incar})
+        return self._calc_type
+
+    @property
+    def run_type(self) -> RunType:
+        params = self.calcs_reversed[0].input.parameters
+        incar = self.calcs_reversed[0].input.incar
+
+        self._run_type = run_type({**params, **incar})
+        return self._run_type
+
     @property
     def structure_entry(self) -> ComputedStructureEntry:
         """
@@ -648,7 +754,7 @@ def structure_entry(self) -> ComputedStructureEntry:
         ComputedStructureEntry
             The TaskDoc.entry with corresponding TaskDoc.structure added.
         """
-        return ComputedStructureEntry(
+        self._structure_entry = ComputedStructureEntry(
             structure=self.structure,
             energy=self.entry.energy,
             correction=self.entry.correction,
@@ -658,6 +764,7 @@ def structure_entry(self) -> ComputedStructureEntry:
             data=self.entry.data,
             entry_id=self.entry.entry_id,
         )
+        return self._structure_entry
 
 
 class TrajectoryDoc(BaseModel):
@@ -912,6 +1019,7 @@ def _get_run_stats(calcs_reversed: List[Calculation]) -> Dict[str, RunStatistics
 def _find_vasp_files(
     path: Union[str, Path],
     volumetric_files: Tuple[str, ...] = _VOLUMETRIC_FILES,
+    task_names: Optional[list[str]] = None,
 ) -> Dict[str, Any]:
     """
     Find VASP files in a directory.
diff --git a/emmet-core/emmet/core/utils.py b/emmet-core/emmet/core/utils.py
index 707653eaf2..c0b275c0e0 100644
--- a/emmet-core/emmet/core/utils.py
+++ b/emmet-core/emmet/core/utils.py
@@ -373,3 +373,8 @@ class {enum_name}(ValueEnum):
     items = [f'    {const} = "{val}"' for const, val in items.items()]
 
     return header + "\n".join(items)
+
+
+def utcnow() -> datetime.datetime:
+    """Get UTC time right now."""
+    return datetime.datetime.now(datetime.timezone.utc)
diff --git a/emmet-core/emmet/core/vasp/calculation.py b/emmet-core/emmet/core/vasp/calculation.py
index b6b1fc2832..c7628bde7a 100644
--- a/emmet-core/emmet/core/vasp/calculation.py
+++ b/emmet-core/emmet/core/vasp/calculation.py
@@ -68,6 +68,13 @@ class StoreTrajectoryOption(ValueEnum):
     NO = "no"
 
 
+class CalculationBaseModel(BaseModel):
+    """Wrapper around pydantic BaseModel with extra functionality."""
+
+    def get(self, key: Any, default_value: Optional[Any] = None) -> Any:
+        return getattr(self, key, default_value)
+
+
 class PotcarSpec(BaseModel):
     """Document defining a VASP POTCAR specification."""
 
@@ -116,7 +123,7 @@ def from_potcar(cls, potcar: Potcar) -> List["PotcarSpec"]:
         return [cls.from_potcar_single(p) for p in potcar]
 
 
-class CalculationInput(BaseModel):
+class CalculationInput(CalculationBaseModel):
     """Document defining VASP calculation inputs."""
 
     incar: Optional[Dict[str, Any]] = Field(
@@ -572,7 +579,7 @@ def from_vasp_outputs(
         )
 
 
-class Calculation(BaseModel):
+class Calculation(CalculationBaseModel):
     """Full VASP calculation inputs and outputs."""
 
     dir_name: Optional[str] = Field(
diff --git a/emmet-core/emmet/core/vasp/material.py b/emmet-core/emmet/core/vasp/material.py
index 171c117475..dc41f3595c 100644
--- a/emmet-core/emmet/core/vasp/material.py
+++ b/emmet-core/emmet/core/vasp/material.py
@@ -12,7 +12,7 @@
 from emmet.core.settings import EmmetSettings
 from emmet.core.structure import StructureMetadata
 from emmet.core.vasp.calc_types import CalcType, RunType, TaskType
-from emmet.core.vasp.task_valid import TaskDocument
+from emmet.core.tasks import TaskDoc
 
 SETTINGS = EmmetSettings()
 
@@ -51,7 +51,7 @@ class MaterialsDoc(CoreMaterialsDoc, StructureMetadata):
     @classmethod
     def from_tasks(
         cls,
-        task_group: List[TaskDocument],
+        task_group: List[TaskDoc],
         structure_quality_scores: Dict[
             str, int
         ] = SETTINGS.VASP_STRUCTURE_QUALITY_SCORES,
@@ -101,7 +101,7 @@ def from_tasks(
         # Always prefer a static over a structure opt
         structure_task_quality_scores = {"Structure Optimization": 1, "Static": 2}
 
-        def _structure_eval(task: TaskDocument):
+        def _structure_eval(task: TaskDoc):
             """
             Helper function to order structures optimization and statics calcs by
             - Functional Type
@@ -155,7 +155,7 @@ def _structure_eval(task: TaskDocument):
         # Always prefer a static over a structure opt
         entry_task_quality_scores = {"Structure Optimization": 1, "Static": 2}
 
-        def _entry_eval(task: TaskDocument):
+        def _entry_eval(task: TaskDoc):
             """
             Helper function to order entries and statics calcs by
             - Spin polarization
@@ -222,7 +222,7 @@ def _entry_eval(task: TaskDocument):
     @classmethod
     def construct_deprecated_material(
         cls,
-        task_group: List[TaskDocument],
+        task_group: List[TaskDoc],
         commercial_license: bool = True,
     ) -> "MaterialsDoc":
         """
diff --git a/emmet-core/emmet/core/vasp/validation.py b/emmet-core/emmet/core/vasp/validation.py
index 39017b2502..2c480d70b8 100644
--- a/emmet-core/emmet/core/vasp/validation.py
+++ b/emmet-core/emmet/core/vasp/validation.py
@@ -4,6 +4,7 @@
 import numpy as np
 from pydantic import ConfigDict, Field, ImportString
 from pymatgen.core.structure import Structure
+from pymatgen.io.vasp.inputs import Kpoints
 from pymatgen.io.vasp.sets import VaspInputSet
 
 from emmet.core.settings import EmmetSettings
@@ -93,10 +94,15 @@ def from_task_doc(
         run_type = task_doc.run_type
         inputs = task_doc.orig_inputs
         chemsys = task_doc.chemsys
-        calcs_reversed = task_doc.calcs_reversed
+        calcs_reversed = [
+            calc if not hasattr(calc, "model_dump") else calc.model_dump()
+            for calc in task_doc.calcs_reversed
+        ]
 
         if calcs_reversed[0].get("input", {}).get("structure", None):
-            structure = Structure.from_dict(calcs_reversed[0]["input"]["structure"])
+            structure = calcs_reversed[0]["input"]["structure"]
+            if isinstance(structure, dict):
+                structure = Structure.from_dict(structure)
         else:
             structure = task_doc.input.structure or task_doc.output.structure
 
@@ -140,7 +146,7 @@ def from_task_doc(
                 # Checking K-Points
                 # Calculations that use KSPACING will not have a .kpoints attr
 
-                if task_type != task_type.NSCF_Line:
+                if task_type != TaskType.NSCF_Line:
                     # Not validating k-point data for line-mode calculations as constructing
                     # the k-path is too costly for the builder and the uniform input set is used.
 
@@ -252,7 +258,7 @@ def _scf_upward_check(calcs_reversed, inputs, data, max_allowed_scf_gradient, wa
 
 def _u_value_checks(task_doc, valid_input_set, warnings):
     # NOTE: Reverting to old method of just using input.hubbards which is wrong in many instances
-    input_hubbards = task_doc.input.hubbards
+    input_hubbards = {} if task_doc.input.hubbards is None else task_doc.input.hubbards
 
     if valid_input_set.incar.get("LDAU", False) or len(input_hubbards) > 0:
         # Assemble required input_set LDAU params into dictionary
@@ -297,9 +303,12 @@ def _kpoint_check(input_set, inputs, calcs_reversed, data, kpts_tolerance):
     else:
         input_dict = inputs
 
-    num_kpts = input_dict.get("kpoints", {}).get("nkpoints", 0) or np.prod(
-        input_dict.get("kpoints", {}).get("kpoints", [1, 1, 1])
-    )
+    kpoints = input_dict.get("kpoints", {})
+    if isinstance(kpoints, Kpoints):
+        kpoints = kpoints.as_dict()
+    elif kpoints is None:
+        kpoints = {}
+    num_kpts = kpoints.get("nkpoints", 0) or np.prod(kpoints.get("kpoints", [1, 1, 1]))
 
     data["kpts_ratio"] = num_kpts / valid_num_kpts
     return data["kpts_ratio"] < kpts_tolerance
@@ -333,14 +342,14 @@ def _potcar_stats_check(task_doc, potcar_stats: dict):
     data_tol = 1.0e-6
 
     try:
-        potcar_details = task_doc.calcs_reversed[0]["input"]["potcar_spec"]
+        potcar_details = task_doc.calcs_reversed[0].model_dump()["input"]["potcar_spec"]
 
     except KeyError:
         # Assume it is an old calculation without potcar_spec data and treat it as passing POTCAR hash check
         return False
 
     use_legacy_hash_check = False
-    if any(len(entry.get("summary_stats", {})) == 0 for entry in potcar_details):
+    if any(entry.get("summary_stats", None) is None for entry in potcar_details):
         # potcar_spec doesn't include summary_stats kwarg needed to check potcars
         # fall back to header hash checking
         use_legacy_hash_check = True
diff --git a/emmet-core/requirements/ubuntu-latest_py3.10_extras.txt b/emmet-core/requirements/ubuntu-latest_py3.10_extras.txt
index 402a669005..32191f10e1 100644
--- a/emmet-core/requirements/ubuntu-latest_py3.10_extras.txt
+++ b/emmet-core/requirements/ubuntu-latest_py3.10_extras.txt
@@ -443,7 +443,7 @@ pygments==2.18.0
     # via
     #   mkdocs-material
     #   rich
-pymatgen==2024.5.1
+pymatgen==2024.4.13
     # via
     #   chgnet
     #   emmet-core
diff --git a/emmet-core/requirements/ubuntu-latest_py3.11_extras.txt b/emmet-core/requirements/ubuntu-latest_py3.11_extras.txt
index 3fc2640102..d132d48018 100644
--- a/emmet-core/requirements/ubuntu-latest_py3.11_extras.txt
+++ b/emmet-core/requirements/ubuntu-latest_py3.11_extras.txt
@@ -437,7 +437,7 @@ pygments==2.18.0
     # via
     #   mkdocs-material
     #   rich
-pymatgen==2024.5.1
+pymatgen==2024.4.13
     # via
     #   chgnet
     #   emmet-core
diff --git a/emmet-core/requirements/ubuntu-latest_py3.9_extras.txt b/emmet-core/requirements/ubuntu-latest_py3.9_extras.txt
index 9430390aa0..d7fa354485 100644
--- a/emmet-core/requirements/ubuntu-latest_py3.9_extras.txt
+++ b/emmet-core/requirements/ubuntu-latest_py3.9_extras.txt
@@ -352,7 +352,7 @@ pyflakes==3.2.0
     # via flake8
 pygments==2.18.0
     # via mkdocs-material
-pymatgen==2024.5.1
+pymatgen==2024.4.13
     # via
     #   chgnet
     #   emmet-core
diff --git a/emmet-core/setup.py b/emmet-core/setup.py
index 7ef8756b0d..04b381aa32 100644
--- a/emmet-core/setup.py
+++ b/emmet-core/setup.py
@@ -27,7 +27,7 @@
     },
     include_package_data=True,
     install_requires=[
-        "pymatgen>=2023.10.11",
+        "pymatgen==2024.4.13",
         "monty>=2024.2.2",
         "pydantic>=2.0",
         "pydantic-settings>=2.0",
diff --git a/emmet-core/tests/test_task.py b/emmet-core/tests/test_task.py
index 4999f11216..de33e092e4 100644
--- a/emmet-core/tests/test_task.py
+++ b/emmet-core/tests/test_task.py
@@ -109,14 +109,16 @@ def test_output_summary(test_dir, object_name, task_name):
         pytest.param("SiNonSCFUniform", id="SiNonSCFUniform"),
     ],
 )
-def test_task_doc(test_dir, object_name):
+def test_task_doc(test_dir, object_name, tmpdir):
     from monty.json import jsanitize
     from monty.serialization import dumpfn
+    import os
     from pymatgen.alchemy.materials import TransformedStructure
     from pymatgen.entries.computed_entries import ComputedEntry
     from pymatgen.transformations.standard_transformations import (
         DeformStructureTransformation,
     )
+    import shutil
 
     from emmet.core.tasks import TaskDoc
 
@@ -156,8 +158,11 @@ def test_task_doc(test_dir, object_name):
         ],
     )
     ts_json = jsanitize(ts.as_dict())
-    dumpfn(ts, f"{dir_name}/transformations.json")
-    test_doc = TaskDoc.from_directory(dir_name)
+    dumpfn(ts, f"{tmpdir}/transformations.json")
+    for f in os.listdir(dir_name):
+        if os.path.isfile(os.path.join(dir_name, f)):
+            shutil.copy(os.path.join(dir_name, f), tmpdir)
+    test_doc = TaskDoc.from_directory(tmpdir)
     # if other_parameters == {}, this is popped from the TaskDoc.transformations field
     # seems like @version is added by monty serialization
     # jsanitize needed because pymatgen.core.Structure.pbc is a tuple
diff --git a/emmet-core/tests/vasp/test_vasp.py b/emmet-core/tests/vasp/test_vasp.py
index 0d24ea6f9f..fef8f27127 100644
--- a/emmet-core/tests/vasp/test_vasp.py
+++ b/emmet-core/tests/vasp/test_vasp.py
@@ -4,7 +4,7 @@
 from monty.io import zopen
 
 from emmet.core.vasp.calc_types import RunType, TaskType, run_type, task_type
-from emmet.core.vasp.task_valid import TaskDocument
+from emmet.core.tasks import TaskDoc
 from emmet.core.vasp.validation import ValidationDoc, _potcar_stats_check
 
 
@@ -45,7 +45,7 @@ def tasks(test_dir):
     with zopen(test_dir / "test_si_tasks.json.gz") as f:
         data = json.load(f)
 
-    return [TaskDocument(**d) for d in data]
+    return [TaskDoc(**d) for d in data]
 
 
 def test_validator(tasks):
@@ -58,7 +58,7 @@ def test_validator(tasks):
 def test_validator_failed_symmetry(test_dir):
     with zopen(test_dir / "failed_elastic_task.json.gz", "r") as f:
         failed_task = json.load(f)
-    taskdoc = TaskDocument(**failed_task)
+    taskdoc = TaskDoc(**failed_task)
     validation = ValidationDoc.from_task_doc(taskdoc)
     assert any("SYMMETRY" in repr(reason) for reason in validation.reasons)
 
@@ -74,20 +74,20 @@ def task_ldau(test_dir):
     with zopen(test_dir / "test_task.json") as f:
         data = json.load(f)
 
-    return TaskDocument(**data)
+    return TaskDoc(**data)
 
 
 def test_ldau(task_ldau):
     task_ldau.input.is_hubbard = True
     assert task_ldau.run_type == RunType.GGA_U
-    assert ValidationDoc.from_task_doc(task_ldau).valid is False
+    assert not ValidationDoc.from_task_doc(task_ldau).valid
 
 
 def test_ldau_validation(test_dir):
     with open(test_dir / "old_aflow_ggau_task.json") as f:
         data = json.load(f)
 
-    task = TaskDocument(**data)
+    task = TaskDoc(**data)
     assert task.run_type == "GGA+U"
 
     valid = ValidationDoc.from_task_doc(task)
@@ -113,21 +113,24 @@ def test_potcar_stats_check(test_dir):
         < filename > )
     I cannot rebuild the TaskDoc without excluding the `orig_inputs` key.
     """
-    task_doc = TaskDocument(**{key: data[key] for key in data if key != "last_updated"})
+    # task_doc = TaskDocument(**{key: data[key] for key in data if key != "last_updated"})
+    task_doc = TaskDoc(**data)
     try:
         # First check: generate hashes from POTCARs in TaskDoc, check should pass
         calc_type = str(task_doc.calc_type)
         expected_hashes = {calc_type: {}}
-        for spec in task_doc.calcs_reversed[0]["input"]["potcar_spec"]:
-            symbol = spec["titel"].split(" ")[1]
+        for spec in task_doc.calcs_reversed[0].input.potcar_spec:
+            symbol = spec.titel.split(" ")[1]
             potcar = PotcarSingle.from_symbol_and_functional(
                 symbol=symbol, functional="PBE"
             )
-            expected_hashes[calc_type][symbol] = {
-                **potcar._summary_stats,
-                "hash": potcar.md5_header_hash,
-                "titel": potcar.TITEL,
-            }
+            expected_hashes[calc_type][symbol] = [
+                {
+                    **potcar._summary_stats,
+                    "hash": potcar.md5_header_hash,
+                    "titel": potcar.TITEL,
+                }
+            ]
 
         assert not _potcar_stats_check(task_doc, expected_hashes)
 
@@ -141,8 +144,8 @@ def test_potcar_stats_check(test_dir):
         # Third check: change data in expected hashes, check should fail
 
         wrong_hashes = {calc_type: {**expected_hashes[calc_type]}}
-        for key in wrong_hashes[calc_type][first_element]["stats"]["data"]:
-            wrong_hashes[calc_type][first_element]["stats"]["data"][key] *= 1.1
+        for key in wrong_hashes[calc_type][first_element][0]["stats"]["data"]:
+            wrong_hashes[calc_type][first_element][0]["stats"]["data"][key] *= 1.1
 
         assert _potcar_stats_check(task_doc, wrong_hashes)
 
@@ -159,7 +162,7 @@ def test_potcar_stats_check(test_dir):
             }
             for potcar in legacy_data["calcs_reversed"][0]["input"]["potcar_spec"]
         ]
-        legacy_task_doc = TaskDocument(
+        legacy_task_doc = TaskDoc(
             **{key: legacy_data[key] for key in legacy_data if key != "last_updated"}
         )
         assert not _potcar_stats_check(legacy_task_doc, expected_hashes)
@@ -180,7 +183,7 @@ def test_potcar_stats_check(test_dir):
         legacy_data["calcs_reversed"][0]["input"]["potcar_spec"][0][
             "hash"
         ] = legacy_data["calcs_reversed"][0]["input"]["potcar_spec"][0]["hash"][:-1]
-        legacy_task_doc = TaskDocument(
+        legacy_task_doc = TaskDoc(
             **{key: legacy_data[key] for key in legacy_data if key != "last_updated"}
         )
         assert _potcar_stats_check(legacy_task_doc, expected_hashes)