From 8775a52ae7a32d796c0634c97e296584e63842a3 Mon Sep 17 00:00:00 2001
From: Jason Munro <jmunro@lbl.gov>
Date: Thu, 2 Sep 2021 15:03:59 -0700
Subject: [PATCH] Provenance to summary (#256)

* Add provenance to summary builder

* Task IDS added from materials to summary

* Fix summary model

* Linting

* Fix summary tests
---
 .../emmet/builders/materials/summary.py       | 30 +++++++------------
 emmet-core/emmet/core/summary.py              |  6 ++++
 tests/emmet-builders/test_summary.py          |  8 +++++
 3 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/emmet-builders/emmet/builders/materials/summary.py b/emmet-builders/emmet/builders/materials/summary.py
index fbee247efb..44a5d00bd5 100644
--- a/emmet-builders/emmet/builders/materials/summary.py
+++ b/emmet-builders/emmet/builders/materials/summary.py
@@ -25,6 +25,7 @@ def __init__(
         substrates,
         surfaces,
         eos,
+        provenance,
         summary,
         chunk_size=100,
         query=None,
@@ -45,6 +46,7 @@ def __init__(
         self.substrates = substrates
         self.surfaces = surfaces
         self.eos = eos
+        self.provenance = provenance
         self.summary = summary
         self.chunk_size = chunk_size
         self.query = query if query else {}
@@ -65,6 +67,7 @@ def __init__(
                 surfaces,
                 substrates,
                 eos,
+                provenance,
             ],
             targets=[summary],
             chunk_size=chunk_size,
@@ -98,32 +101,22 @@ def get_items(self):
                 "materials": self.materials.query_one({self.materials.key: entry}),
                 "thermo": self.thermo.query_one({self.thermo.key: entry}),
                 "xas": list(self.xas.query({self.xas.key: entry})),
-                "grain_boundaries": list(
-                    self.grain_boundaries.query({self.grain_boundaries.key: entry})
-                ),
-                "electronic_structure": self.electronic_structure.query_one(
-                    {self.electronic_structure.key: entry}
-                ),
+                "grain_boundaries": list(self.grain_boundaries.query({self.grain_boundaries.key: entry})),
+                "electronic_structure": self.electronic_structure.query_one({self.electronic_structure.key: entry}),
                 "magnetism": self.magnetism.query_one({self.magnetism.key: entry}),
                 "elasticity": self.elasticity.query_one({self.elasticity.key: entry}),
                 "dielectric": self.dielectric.query_one({self.dielectric.key: entry}),
-                "piezoelectric": self.piezoelectric.query_one(
-                    {self.piezoelectric.key: entry}
-                ),
-                "phonon": self.phonon.query_one(
-                    {self.phonon.key: entry}, [self.phonon.key]
-                ),
+                "piezoelectric": self.piezoelectric.query_one({self.piezoelectric.key: entry}),
+                "phonon": self.phonon.query_one({self.phonon.key: entry}, [self.phonon.key]),
                 "insertion_electrodes": list(
                     self.insertion_electrodes.query(
-                        {self.insertion_electrodes.key: entry},
-                        [self.insertion_electrodes.key],
+                        {self.insertion_electrodes.key: entry}, [self.insertion_electrodes.key],
                     )
                 ),
-                "surface_properties": self.surfaces.query_one(
-                    {self.surfaces.key: entry}
-                ),
+                "surface_properties": self.surfaces.query_one({self.surfaces.key: entry}),
                 "substrates": list(self.surfaces.query({self.substrates.key: entry})),
                 "eos": self.eos.query_one({self.eos.key: entry}, [self.eos.key]),
+                "provenance": self.provenance.query_one({self.provenance.key: entry}),
             }
 
             sub_fields = {
@@ -137,8 +130,7 @@ def get_items(self):
                 if data[collection] is not None:
                     data[collection] = (
                         data[collection][sub_field]
-                        if (sub_field in data[collection])
-                        and (data[collection][sub_field] != {})
+                        if (sub_field in data[collection]) and (data[collection][sub_field] != {})
                         else None
                     )
 
diff --git a/emmet-core/emmet/core/summary.py b/emmet-core/emmet/core/summary.py
index 3d21f5113f..7a4a6678fa 100644
--- a/emmet-core/emmet/core/summary.py
+++ b/emmet-core/emmet/core/summary.py
@@ -85,6 +85,10 @@ class SummaryDoc(PropertyDoc):
 
     structure: Structure = Field(..., description="The lowest energy structure for this material")
 
+    task_ids: List[MPID] = Field(
+        [], title="Calculation IDs", description="List of Calculations IDs associated with this material.",
+    )
+
     # Thermo
 
     uncorrected_energy_per_atom: float = Field(
@@ -268,6 +272,7 @@ def from_docs(cls, material_id: MPID, **docs: Dict[str, Dict]):
         "symmetry",
         "structure",
         "deprecated",
+        "task_ids",
     ],
     "thermo": [
         "uncorrected_energy_per_atom",
@@ -320,6 +325,7 @@ def from_docs(cls, material_id: MPID, **docs: Dict[str, Dict]):
         "weighted_work_function",
         "has_reconstructed",
     ],
+    "provenance": ["theoretical"],
     "eos": [],
     "phonon": [],
     "insertion_electrodes": [],
diff --git a/tests/emmet-builders/test_summary.py b/tests/emmet-builders/test_summary.py
index 10cddaeebc..f3d329b20d 100644
--- a/tests/emmet-builders/test_summary.py
+++ b/tests/emmet-builders/test_summary.py
@@ -86,6 +86,11 @@ def xas():
     return MemoryStore()
 
 
+@pytest.fixture
+def provenance():
+    return MemoryStore()
+
+
 @pytest.fixture
 def summary():
     return MemoryStore(key="material_id")
@@ -106,6 +111,7 @@ def test_summary_builder(
     substrates,
     surfaces,
     eos,
+    provenance,
     summary,
 ):
 
@@ -124,6 +130,7 @@ def test_summary_builder(
         xas=xas,
         grain_boundaries=grain_boundaries,
         eos=eos,
+        provenance=provenance,
         summary=summary,
     )
 
@@ -148,6 +155,7 @@ def test_serialization(tmpdir):
         MemoryStore(),
         MemoryStore(),
         MemoryStore(),
+        MemoryStore(),
     )
 
     dumpfn(builder.as_dict(), Path(tmpdir) / "test.json")