emmet-builders 0.57.2__py3-none-any.whl → 0.58.0__py3-none-any.whl

emmet/builders/materials/electronic_structure.py

@@ -1,6 +1,8 @@
 from collections import defaultdict
 from math import ceil
 import itertools
+import re
+import boto3
 import numpy as np
 from maggma.builders import Builder
 from maggma.utils import grouper
@@ -18,6 +20,8 @@ from emmet.core.settings import EmmetSettings
 from emmet.core.electronic_structure import ElectronicStructureDoc
 from emmet.core.utils import jsanitize
 
+from emmet.builders.utils import query_open_data
+
 SETTINGS = EmmetSettings()
 
 
@@ -43,8 +47,10 @@ class ElectronicStructureBuilder(Builder):
         tasks (Store): Store of task documents
         materials (Store): Store of materials documents
         electronic_structure (Store): Store of electronic structure summary data documents
-        bandstructure_fs (Store): Store of bandstructures
-        dos_fs (Store): Store of DOS
+        bandstructure_fs (Store, str): Store of bandstructures, or S3 URL string with prefix
+            (e.g. s3://materialsproject-parsed/bandstructures).
+        dos_fs (Store, str): Store of DOS, or S3 URL string with bucket and prefix
+            (e.g. s3://materialsproject-parsed/dos).
         chunk_size (int): Chunk size to use for processing. Defaults to 10.
         query (dict): Dictionary to limit materials to be analyzed
         """
@@ -57,8 +63,27 @@ class ElectronicStructureBuilder(Builder):
         self.chunk_size = chunk_size
         self.query = query if query else {}
 
+        self._s3_resource = None
+
+        sources = [tasks, materials]
+
+        fs_stores = [bandstructure_fs, dos_fs]
+
+        for store in fs_stores:
+            if isinstance(store, str):
+                if not re.match("^s3://.*", store):
+                    raise ValueError(
+                        "Please provide an S3 URL "
+                        "in the format s3://{bucket_name}/{prefix}"
+                    )
+
+                if self._s3_resource is None:
+                    self._s3_resource = boto3.resource("s3")
+            else:
+                sources.append(store)
+
         super().__init__(
-            sources=[tasks, materials, bandstructure_fs, dos_fs],
+            sources=sources,
             targets=[electronic_structure],
             chunk_size=chunk_size,
             **kwargs,
@@ -385,6 +410,7 @@ class ElectronicStructureBuilder(Builder):
                     structure = Structure.from_dict(task_query["output"]["structure"])
 
                     kpoints = task_query["orig_inputs"]["kpoints"]
+
                     labels_dict = {
                         label: point
                         for label, point in zip(kpoints["labels"], kpoints["kpoints"])
@@ -397,16 +423,32 @@ class ElectronicStructureBuilder(Builder):
                         bs_type = None
 
                     if bs_type is None:
-                        bs_dict = self.bandstructure_fs.query_one(
-                            {self.bandstructure_fs.key: str(task_id)}
-                        )
+                        if isinstance(self.bandstructure_fs, str):
+                            _, _, bucket, prefix = self.bandstructure_fs.strip(
+                                "/"
+                            ).split("/")
+
+                            bs_dict = query_open_data(
+                                bucket,
+                                prefix,
+                                task_id,
+                                monty_decode=False,
+                                s3_resource=self._s3_resource,
+                            )
+                        else:
+                            bs_dict = self.bandstructure_fs.query_one(
+                                {self.bandstructure_fs.key: str(task_id)}
+                            )
 
                         if bs_dict is not None:
                             bs = BandStructureSymmLine.from_dict(bs_dict["data"])
 
-                            bs_type = self._obtain_path_type(
-                                bs.labels_dict, bs.structure
-                            )
+                            labels_dict = {
+                                label: kpoint.frac_coords
+                                for label, kpoint in bs.labels_dict.items()
+                            }
+
+                            bs_type = self._obtain_path_type(labels_dict, bs.structure)
 
                     is_hubbard = task_query["input"]["is_hubbard"]
                     lmaxmix = task_query["input"]["incar"].get(
@@ -425,6 +467,7 @@ class ElectronicStructureBuilder(Builder):
                                 "nkpoints": int(nkpoints),
                                 "updated_on": lu_dt,
                                 "output_structure": structure,
+                                "labels_dict": labels_dict,
                             }
                         )
 
@@ -565,10 +608,19 @@ class ElectronicStructureBuilder(Builder):
                 materials_doc["bandstructure"][bs_type]["lmaxmix"] = sorted_bs_data[0][
                     "lmaxmix"
                 ]
-
-                bs_obj = self.bandstructure_fs.query_one(
-                    criteria={"fs_id": sorted_bs_data[0]["fs_id"]}
-                )
+                if isinstance(self.bandstructure_fs, str):
+                    _, _, bucket, prefix = self.bandstructure_fs.strip("/").split("/")
+                    bs_obj = query_open_data(
+                        bucket,
+                        prefix,
+                        sorted_bs_data[0]["task_id"],
+                        monty_decode=False,
+                        s3_resource=self._s3_resource,
+                    )
+                else:
+                    bs_obj = self.bandstructure_fs.query_one(
+                        criteria={"fs_id": sorted_bs_data[0]["fs_id"]}
+                    )
 
                 materials_doc["bandstructure"][bs_type]["object"] = (
                     bs_obj["data"] if bs_obj is not None else None
@@ -602,9 +654,20 @@ class ElectronicStructureBuilder(Builder):
 
             materials_doc["dos"]["lmaxmix"] = sorted_dos_data[0]["lmaxmix"]
 
-            dos_obj = self.dos_fs.query_one(
-                criteria={"fs_id": sorted_dos_data[0]["fs_id"]}
-            )
+            if isinstance(self.bandstructure_fs, str):
+                _, _, bucket, prefix = self.dos_fs.strip("/").split("/")
+                dos_obj = query_open_data(
+                    bucket,
+                    prefix,
+                    sorted_dos_data[0]["task_id"],
+                    monty_decode=False,
+                    s3_resource=self._s3_resource,
+                )
+            else:
+                dos_obj = self.dos_fs.query_one(
+                    criteria={"fs_id": sorted_dos_data[0]["fs_id"]}
+                )
+
             materials_doc["dos"]["object"] = (
                 dos_obj["data"] if dos_obj is not None else None
             )

emmet/builders/utils.py

@@ -1,6 +1,12 @@
-from typing import Set, Union
+from typing import Set, Union, Any
 import sys
 import os
+from gzip import GzipFile
+import orjson
+import json
+from io import BytesIO
+from monty.serialization import MontyDecoder
+from botocore.exceptions import ClientError
 from itertools import chain, combinations
 from pymatgen.core import Structure
 from pymatgen.analysis.diffusion.neb.full_path_mapper import MigrationGraph
@@ -146,6 +152,56 @@ def get_hop_cutoff(
         return None
 
 
+def query_open_data(
+    bucket: str,
+    prefix: str,
+    key: str,
+    monty_decode: bool = True,
+    s3_resource: Any = None,
+) -> Union[dict, None]:
+    """Query a Materials Project AWS S3 Open Data bucket directly with boto3
+
+    Args:
+        bucket (str): Materials project bucket name
+        prefix (str): Full set of file prefixes
+        key (str): Key for file
+        monty_decode (bool): Whether to monty decode or keep as dictionary. Defaults to True.
+        s3_resource (Optional[Any]): S3 resource. One will be instantiated if none are provided
+
+    Returns:
+        dict: MontyDecoded data or None
+    """
+
+    def decode(content, monty_decode):
+        if monty_decode:
+            result = MontyDecoder().decode(content)
+        else:
+            result = orjson.loads(content)
+        return result
+
+    try:
+        file_key = f"{prefix}/{key}.json.gz"
+        ref = s3_resource.Object(bucket, file_key)  # type: ignore
+        bytes = ref.get()["Body"]  # type: ignore
+
+        with GzipFile(fileobj=bytes, mode="r") as gzipfile:
+            content = gzipfile.read()
+
+        try:
+            result = decode(content, monty_decode)
+        except (orjson.JSONDecodeError, json.JSONDecodeError):
+            try:
+                with GzipFile(fileobj=BytesIO(content), mode="r") as gzipfile_nested:
+                    result = decode(gzipfile_nested.read(), monty_decode)
+            except Exception:
+                print(f"Issue decoding {file_key} from bucket {bucket}")
+                return None
+    except ClientError:
+        return None
+
+    return result
+
+
 # From: https://stackoverflow.com/a/45669280
 class HiddenPrints:
     def __enter__(self):

{emmet_builders-0.57.2.dist-info → emmet_builders-0.58.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: emmet-builders
-Version: 0.57.2
+Version: 0.58.0
 Summary: Builders for the Emmet Library
 Home-page: https://github.com/materialsproject/emmet
 Author: The Materials Project
@@ -9,7 +9,7 @@ License: modified BSD
 Platform: UNKNOWN
 Requires-Python: >=3.8
 Requires-Dist: emmet-core[all]
-Requires-Dist: maggma (>=0.49.17)
+Requires-Dist: maggma (>=0.51.9)
 Requires-Dist: matminer (>=0.7.3)
 Provides-Extra: docs
 Requires-Dist: mkdocs ; extra == 'docs'

{emmet_builders-0.57.2.dist-info → emmet_builders-0.58.0.dist-info}/RECORD

@@ -1,6 +1,6 @@
 emmet/builders/__init__.py,sha256=y-ZREtieuFK3MaYvCBDMPf3YLxnsZG1VNho9lMjnRDU,221
 emmet/builders/settings.py,sha256=xVIqpQjSxn--m6mbKJIr7g-nqoqZr1Vcxi3YNa3etOg,2766
-emmet/builders/utils.py,sha256=VtA_77cBqzKBaA747FBbcPX3M0QmQE-l2nMLzlQwijw,5775
+emmet/builders/utils.py,sha256=iaCWMGjttInSW21a_zEk4CqN4HPlSx7BX3a48QMZw3c,7537
 emmet/builders/abinit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 emmet/builders/abinit/phonon.py,sha256=dWQ3pGjxTES7jkwZVoqeNgTMhXORXcg8rkxAA9lLd5Q,32026
 emmet/builders/abinit/sound_velocity.py,sha256=3S3EaAaLwKKZ9-4Lnv3nUX6dHkIahRkpr3ecM_8nWx8,6999
@@ -16,7 +16,7 @@ emmet/builders/materials/corrected_entries.py,sha256=GctwGyE3JYBAnKcCN2PJjOa1stm
 emmet/builders/materials/dielectric.py,sha256=-AagE2LwsB6U5A1cG-wPEKK-4zd3n57XzWEOQ0qCLps,6522
 emmet/builders/materials/elasticity.py,sha256=O6AdNlPwUCz9bQmgRVoqtlwMTkgK31e-kKK94HmUVi8,16297
 emmet/builders/materials/electrodes.py,sha256=7bkCCObs-XMcJ3wv6BjX6vWzdK2jdwdeAMNGQkvro6Y,23258
-emmet/builders/materials/electronic_structure.py,sha256=wadGgWGuW7zsZMDODVZVPaBjR7EQSwZiVs1xm-tPzN8,26395
+emmet/builders/materials/electronic_structure.py,sha256=gn7krFpbNs4q82TWYsycQmRrhdD1tjx4JdOi3qGspRY,28892
 emmet/builders/materials/magnetism.py,sha256=h3sMPracp1KaYKxBEQLt9AlLmbUl_OAE2g5YZdQykZA,4599
 emmet/builders/materials/optimade.py,sha256=fdY9sM2eJt0g_ka-QkGQSSSHJ2xfiWN-G7LIvHuctyo,5158
 emmet/builders/materials/oxidation_states.py,sha256=bgHFpVzrHyyifTE1jdGTMHi39GRiZ9wCcmQKbCULXlY,1540
@@ -43,7 +43,7 @@ emmet/builders/qchem/molecules.py,sha256=TpCysoSQzWGZ_P4sDfz76wq6tcodyPFDJv0v9xt
 emmet/builders/vasp/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 emmet/builders/vasp/materials.py,sha256=RGJawnvUwwR1PIJBEIJSbIuGzwsF2lj04P9rl37RIZ0,13332
 emmet/builders/vasp/task_validator.py,sha256=sV-yxBq-ZTKjI73jdEzWdaFTC0Ng18xMGoHglu5Fh9Q,3600
-emmet_builders-0.57.2.dist-info/METADATA,sha256=vxPoIE9wQIfsnKKekj6TJx_JB-IEb8rvmq6tQ0NjTRw,1449
-emmet_builders-0.57.2.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
-emmet_builders-0.57.2.dist-info/top_level.txt,sha256=6GcpbmWPeFhNCTfDFilb8GQ4T1UQu4z9c5jpobjwE-Q,6
-emmet_builders-0.57.2.dist-info/RECORD,,
+emmet_builders-0.58.0.dist-info/METADATA,sha256=c_4UxlNCjeSZEeqGuIzGxCzKowSRyPYrAVcbaZGuOHs,1448
+emmet_builders-0.58.0.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
+emmet_builders-0.58.0.dist-info/top_level.txt,sha256=6GcpbmWPeFhNCTfDFilb8GQ4T1UQu4z9c5jpobjwE-Q,6
+emmet_builders-0.58.0.dist-info/RECORD,,