Merge pull request #217 from materialsproject/fix-compositon-serializ…

…ation

Fix oxidation doc composition

emmet-builders/emmet/builders/materials/oxidation_states.py

@@ -4,6 +4,7 @@
 from pymatgen.core import __version__ as pymatgen_version
 
 from emmet.core.oxidation_states import OxidationStateDoc
+from emmet.core.utils import jsanitize
 
 
 class OxidationStatesBuilder(MapBuilder):
@@ -37,14 +38,11 @@ def __init__(
 
     def unary_function(self, item):
         structure = Structure.from_dict(item["structure"])
-        oxi_doc = OxidationStateDoc.from_structure(structure)
-        doc = oxi_doc.dict()
-
-        doc.update(
-            {
-                "pymatgen_version": pymatgen_version,
-                "successful": True,
-            }
+        mpid = item["material_id"]
+
+        oxi_doc = OxidationStateDoc.from_structure(
+            structure=structure, material_id=mpid
         )
+        doc = jsanitize(oxi_doc.dict(), allow_bson=True)
 
         return doc

emmet-core/emmet/core/structure.py

@@ -85,13 +85,15 @@ def from_composition(
             if fields is None
             else fields
         )
+        composition = composition.remove_charges()
+
         elsyms = sorted(set([e.symbol for e in composition.elements]))
 
         data = {
             "elements": elsyms,
             "nelements": len(elsyms),
             "composition": composition,
-            "composition_reduced": composition.reduced_composition,
+            "composition_reduced": composition.reduced_composition.remove_charges(),
             "formula_pretty": composition.reduced_formula,
             "formula_anonymous": composition.anonymized_formula,
             "chemsys": "-".join(elsyms),
@@ -126,7 +128,7 @@ def from_structure(
             if fields is None
             else fields
         )
-        comp = structure.composition
+        comp = structure.composition.remove_charges()
         elsyms = sorted(set([e.symbol for e in comp.elements]))
         symmetry = SymmetryData.from_structure(structure)
 

tests/emmet-builders/test_oxidation.py

@@ -0,0 +1,30 @@
+import pytest
+from maggma.stores import JSONStore, MemoryStore
+
+from emmet.builders.materials.oxidation_states import OxidationStatesBuilder
+
+
+@pytest.fixture(scope="session")
+def fake_materias(test_dir):
+    entries = JSONStore(test_dir / "LiTiO2_batt.json", key="entry_id")
+    entries.connect()
+
+    materials_store = MemoryStore(key="material_id")
+    materials_store.connect()
+
+    for doc in entries.query():
+        materials_store.update(
+            {"material_id": doc["entry_id"], "structure": doc["structure"]}
+        )
+    return materials_store
+
+
+def test_oxidation_store(fake_materias):
+    oxi_store = MemoryStore()
+    builder = OxidationStatesBuilder(
+        materials=fake_materias, oxidation_states=oxi_store
+    )
+    builder.run()
+
+    assert oxi_store.count() == 6
+    assert all([isinstance(d["composition"], dict) for d in oxi_store.query()])