use ComputedStructureEntries

emmet-core/emmet/core/thermo.py

@@ -1,15 +1,15 @@
 """ 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
 from pymatgen.core.periodic_table import Element
 
 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)
 

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

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)