diff --git a/emmet-core/emmet/core/thermo.py b/emmet-core/emmet/core/thermo.py index f1e55e058b..37d6dce012 100644 --- a/emmet-core/emmet/core/thermo.py +++ b/emmet-core/emmet/core/thermo.py @@ -1,7 +1,7 @@ """ Core definition of a Thermo Document """ from datetime import datetime from enum import Enum -from typing import ClassVar, Dict, List +from typing import ClassVar, Dict, List, Union from pydantic import BaseModel, Field from pymatgen.analysis.phase_diagram import PhaseDiagram, PhaseDiagramError @@ -9,7 +9,7 @@ from emmet.core.material_property import PropertyDoc from emmet.core.structure import StructureMetadata -from emmet.stubs import Composition, ComputedEntry +from emmet.stubs import Composition, ComputedEntry, ComputedStructureEntry class DecompositionProduct(BaseModel): @@ -83,14 +83,14 @@ class ThermoDoc(PropertyDoc): description="List of available energy types computed for this material" ) - entries: Dict[str, ComputedEntry] = Field( - None, + entries: Dict[str, Union[ComputedEntry, ComputedStructureEntry]] = Field( + ..., description="List of all entries that are valid for this material." " The keys for this dictionary are names of various calculation types", ) @classmethod - def from_entries(cls, entries: List[ComputedEntry]): + def from_entries(cls, entries: List[Union[ComputedEntry, ComputedStructureEntry]]): pd = PhaseDiagram(entries) diff --git a/emmet-core/emmet/core/vasp/material.py b/emmet-core/emmet/core/vasp/material.py index a82cc163bf..87feb02ac5 100644 --- a/emmet-core/emmet/core/vasp/material.py +++ b/emmet-core/emmet/core/vasp/material.py @@ -12,7 +12,7 @@ from emmet.core.structure import StructureMetadata from emmet.core.vasp.calc_types import CalcType, RunType, TaskType from emmet.core.vasp.task import TaskDocument -from emmet.stubs import ComputedEntry, Structure +from emmet.stubs import ComputedStructureEntry, Structure class MaterialsDoc(CoreMaterialsDoc, StructureMetadata): @@ -34,7 +34,7 @@ class MaterialsDoc(CoreMaterialsDoc, StructureMetadata): None, description="Mappingionary for tracking the provenance of properties" ) - entries: Mapping[RunType, ComputedEntry] = Field( + entries: Mapping[RunType, ComputedStructureEntry] = Field( None, description="Dictionary for tracking entries for VASP calculations" ) @@ -136,7 +136,7 @@ def _structure_eval(task: TaskDocument): if len(relevant_calcs) > 0: best_task_doc = relevant_calcs[0] - entry = best_task_doc.entry + entry = best_task_doc.structure_entry entry.data["task_id"] = entry.entry_id entry.entry_id = material_id entries[rt] = entry diff --git a/emmet-core/emmet/core/vasp/task.py b/emmet-core/emmet/core/vasp/task.py index 91ffb7fc83..8db3d5a3a5 100644 --- a/emmet-core/emmet/core/vasp/task.py +++ b/emmet-core/emmet/core/vasp/task.py @@ -18,7 +18,13 @@ run_type, task_type, ) -from emmet.stubs import ComputedEntry, Matrix3D, Structure, Vector3D +from emmet.stubs import ( + ComputedEntry, + ComputedStructureEntry, + Matrix3D, + Structure, + Vector3D, +) class Status(ValueEnum): @@ -153,3 +159,11 @@ def entry(self): } return ComputedEntry.from_dict(entry_dict) + + @property + def structure_entry(self): + """ Turns a Task Doc into a ComputedStructureEntry""" + entry_dict = self.entry.as_dict() + entry_dict["structure"] = self.output.structure + + return ComputedStructureEntry.from_dict(entry_dict)