revert to mp_master

emmet-builders/emmet/builders/materials/electrodes.py

@@ -1,24 +1,24 @@
-import math
 import operator
+import math
 from collections import namedtuple
 from datetime import datetime
 from functools import lru_cache
-from itertools import chain, groupby
+from itertools import groupby, chain
 from pprint import pprint
-from typing import Any, Dict, Iterable, List
+from typing import Iterable, Dict, List, Any
 
+from emmet.core.electrode import InsertionElectrodeDoc
+from emmet.core.structure_group import StructureGroupDoc
+from emmet.core.utils import jsanitize
 from maggma.builders import Builder, MapBuilder
 from maggma.stores import MongoStore
 from monty.json import MontyEncoder
 from numpy import unique
-from pymatgen.analysis.structure_matcher import ElementComparator, StructureMatcher
+from pymatgen import Composition
+from pymatgen.analysis.structure_matcher import StructureMatcher, ElementComparator
 from pymatgen.apps.battery.insertion_battery import InsertionElectrode
-from pymatgen.core import Composition, Structure
-from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
-
-from emmet.core.electrode import InsertionElectrodeDoc
-from emmet.core.structure_group import StructureGroupDoc
-from emmet.core.utils import jsanitize
+from pymatgen.core import Structure
+from pymatgen.entries.computed_entries import ComputedStructureEntry
 
 __author__ = "Jimmy Shen"
 __email__ = "jmmshn@lbl.gov"
@@ -89,19 +89,18 @@ def generic_groupby(list_in, comp=operator.eq):
     return list_out
 
 
-
 class StructureGroupBuilder(Builder):
     def __init__(
-            self,
-            materials: MongoStore,
-            sgroups: MongoStore,
-            working_ion: str,
-            query: dict = None,
-            ltol: float = 0.2,
-            stol: float = 0.3,
-            angle_tol: float = 5.0,
-            check_newer: bool = True,
-            **kwargs,
+        self,
+        materials: MongoStore,
+        sgroups: MongoStore,
+        working_ion: str,
+        query: dict = None,
+        ltol: float = 0.2,
+        stol: float = 0.3,
+        angle_tol: float = 5.0,
+        check_newer: bool = True,
+        **kwargs,
     ):
         """
         Aggregate materials entries into sgroups that are topotactically similar to each other.
@@ -278,18 +277,21 @@ def process_item(self, item: Any) -> Any:
     def _remove_targets(self, rm_ids):
         self.sgroups.remove_docs({"material_id": {"$in": rm_ids}})
 
+
 class InsertionElectrodeBuilder(MapBuilder):
     def __init__(
-            self,
-            grouped_materials: MongoStore,
-            insertion_electrode: MongoStore,
-            thermo: MongoStore,
-            query: dict = None,
-            **kwargs,
+        self,
+        grouped_materials: MongoStore,
+        insertion_electrode: MongoStore,
+        thermo: MongoStore,
+        material: MongoStore,
+        query: dict = None,
+        **kwargs,
     ):
         self.grouped_materials = grouped_materials
         self.insertion_electrode = insertion_electrode
         self.thermo = thermo
+        self.material = material
         qq_ = {} if query is None else query
         qq_.update({"structure_matched": True, "has_distinct_compositions": True})
         super().__init__(
@@ -302,12 +304,12 @@ def __init__(
     def get_items(self):
         """"""
 
-        @lru_cache()
+        @lru_cache(None)
         def get_working_ion_entry(working_ion):
             with self.thermo as store:
                 working_ion_docs = [*store.query({"chemsys": working_ion})]
             best_wion = min(
-                working_ion_docs, key=lambda x: x["energy_per_atom"]
+                working_ion_docs, key=lambda x: x["thermo"]["energy_per_atom"]
             )
             return best_wion
 
@@ -323,23 +325,35 @@ def modify_item(item):
                                 {"material_id": {"$in": item["grouped_ids"]}},
                             ]
                         },
-                        properties=["material_id", "_sbxn", "thermo", "entries", "energy_type", "energy_above_hull"],
+                        properties=["material_id", "_sbxn", "thermo"],
                     )
                 ]
 
-            self.logger.debug(f"Found for {len(thermo_docs)} Thermo Documents.")
+            with self.material as store:
+                material_docs = [
+                    *store.query(
+                        {
+                            "$and": [
+                                {"material_id": {"$in": item["grouped_ids"]}},
+                                {"_sbxn": {"$in": ["core"]}},
+                            ]
+                        },
+                        properties=["material_id", "structure"],
+                    )
+                ]
 
+            self.logger.debug(f"Found for {len(thermo_docs)} Thermo Documents.")
             if len(item["ignored_species"]) != 1:
                 raise ValueError(
                     "Insertion electrode can only be defined for one working ion species"
                 )
-
             working_ion_doc = get_working_ion_entry(item["ignored_species"][0])
             return {
                 "material_id": item["material_id"],
                 "working_ion_doc": working_ion_doc,
                 "working_ion": item["ignored_species"][0],
                 "thermo_docs": thermo_docs,
+                "material_docs": material_docs,
             }
 
         yield from map(modify_item, super().get_items())
@@ -349,26 +363,40 @@ def unary_function(self, item):
         - Add volume information to each entry to create the insertion electrode document
         - Add the host structure
         """
-        entries = [tdoc_["entries"][tdoc_["energy_type"]] for tdoc_ in item["thermo_docs"]]
-        entries = list(map(ComputedStructureEntry.from_dict, entries))
+        entries = [tdoc_["thermo"]["entry"] for tdoc_ in item["thermo_docs"]]
+        entries = list(map(ComputedEntry.from_dict, entries))
         working_ion_entry = ComputedEntry.from_dict(
-            item["working_ion_doc"]["entries"][item["working_ion_doc"]['energy_type']]
+            item["working_ion_doc"]["thermo"]["entry"]
         )
         working_ion = working_ion_entry.composition.reduced_formula
-
         decomp_energies = {
-            d_["material_id"]: d_["energy_above_hull"]
+            d_["material_id"]: d_["thermo"]["e_above_hull"]
             for d_ in item["thermo_docs"]
         }
+        mat_structures = {
+            mat_d_["material_id"]: Structure.from_dict(mat_d_["structure"])
+            for mat_d_ in item["material_docs"]
+        }
 
         least_wion_ent = min(
             entries, key=lambda x: x.composition.get_atomic_fraction(working_ion)
         )
-        host_structure = least_wion_ent.structure.copy()
+        mdoc_ = next(
+            filter(
+                lambda x: x["material_id"] == least_wion_ent.entry_id,
+                item["material_docs"],
+            )
+        )
+        host_structure = Structure.from_dict(mdoc_["structure"])
         host_structure.remove_species([item["working_ion"]])
 
         for ient in entries:
-            ient.data["volume"] = ient.structure.volume
+            if mat_structures[ient.entry_id].composition != ient.composition:
+                raise RuntimeError(
+                    f"In {item['material_id']}: the compositions for task {ient.entry_id} are matched "
+                    "between the StructureGroup DB and the Thermo DB "
+                )
+            ient.data["volume"] = mat_structures[ient.entry_id].volume
             ient.data["decomposition_energy"] = decomp_energies[ient.entry_id]
 
         ie = InsertionElectrodeDoc.from_entries(

emmet-core/emmet/core/electrode.py

@@ -117,8 +117,6 @@ class InsertionElectrodeDoc(InsertionVoltagePairDoc):
 
     framework: Composition
 
-    electrode_object: Dict
-
     # Make sure that the datetime field is properly formatted
     @validator("last_updated", pre=True)
     def last_updated_dict_ok(cls, v):
@@ -134,18 +132,17 @@ def from_entries(
     ) -> Union["InsertionElectrodeDoc", None]:
         try:
             ie = InsertionElectrode.from_entries(
-                entries=grouped_entries, working_ion_entry=working_ion_entry, strip_structures=True
+                entries=grouped_entries, working_ion_entry=working_ion_entry
             )
         except IndexError:
             return None
         d = ie.get_summary_dict()
         d["num_steps"] = d.pop("nsteps", None)
         d["last_updated"] = datetime.utcnow()
         return cls(
-            battery_id=task_id,
+            task_id=task_id,
             host_structure=host_structure.as_dict(),
             framework=Composition(d["framework_formula"]),
-            electrode_object=ie.as_dict(),
             **d
         )