Fix thermo type to work with s3 folders (#490)

materialsproject · Aug 1, 2022 · 62414dc · 62414dc
1 parent 7d2d72c
commit 62414dc
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 8 deletions.
diff --git a/emmet-builders/emmet/builders/vasp/thermo.py b/emmet-builders/emmet/builders/vasp/thermo.py
@@ -56,19 +56,25 @@ def __init__(
         self._entries_cache: Dict[str, List[ComputedStructureEntry]] = defaultdict(list)
 
         if self.thermo.key != "thermo_id":
-            warnings.warn(f"Key for the thermo store is incorrect and has been changed from {self.thermo.key} to thermo_id!")
+            warnings.warn(
+                f"Key for the thermo store is incorrect and has been changed from {self.thermo.key} to thermo_id!"
+            )
             self.thermo.key = "thermo_id"
 
         if self.materials.key != "material_id":
-            warnings.warn(f"Key for the materials store is incorrect and has been changed from {self.materials.key} to material_id!")
+            warnings.warn(
+                f"Key for the materials store is incorrect and has been changed from {self.materials.key} to material_id!"
+            )
             self.materials.key = "material_id"
 
         sources = [materials]
 
         if self.oxidation_states is not None:
 
             if self.oxidation_states.key != "material_id":
-                warnings.warn(f"Key for the oxidation states store is incorrect and has been changed from {self.oxidation_states.key} to material_id!")
+                warnings.warn(
+                    f"Key for the oxidation states store is incorrect and has been changed from {self.oxidation_states.key} to material_id!"
+                )
                 self.oxidation_states.key = "material_id"
 
             sources.append(oxidation_states)  # type: ignore
@@ -78,7 +84,9 @@ def __init__(
         if self.phase_diagram is not None:
 
             if self.phase_diagram.key != "phase_diagram_id":
-                warnings.warn(f"Key for the phase diagram store is incorrect and has been changed from {self.thphase_diagramermo.key} to phase_diagram_id!")
+                warnings.warn(
+                    f"Key for the phase diagram store is incorrect and has been changed from {self.thphase_diagramermo.key} to phase_diagram_id!"
+                )
                 self.phase_diagram.key = "phase_diagram_id"
 
             targets.append(phase_diagram)  # type: ignore
@@ -103,8 +111,8 @@ def ensure_indexes(self):
 
         # Search index for phase_diagram
         if self.phase_diagram:
-            self.phase_diagram.ensure_index("chemsys")
-            self.phase_diagram.ensure_index("phase_diagram_id")
+            self.phase_diagram.index.ensure_index("chemsys")
+            self.phase_diagram.index.ensure_index("phase_diagram_id")
 
     def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
         updated_chemsys = self.get_updated_chemsys()
@@ -227,6 +235,7 @@ def process_item(self, item: List[Dict]):
                             phase_diagram=pd,
                             thermo_type=thermo_type,
                         )
+
                         pd_data = jsanitize(pd_doc.dict(), allow_bson=True)
 
                 docs_pd_pair = (

diff --git a/emmet-core/emmet/core/thermo.py b/emmet-core/emmet/core/thermo.py
@@ -34,8 +34,8 @@ class DecompositionProduct(BaseModel):
 
 
 class ThermoType(ValueEnum):
-    GGA_GGA_U = "GGA/GGA+U"
-    GGA_GGA_U_R2SCAN = "GGA/GGA+U/R2SCAN"
+    GGA_GGA_U = "GGA_GGA+U"
+    GGA_GGA_U_R2SCAN = "GGA_GGA+U_R2SCAN"
     UNKNOWN = "UNKNOWN"