emmet-builders 0.57.2__tar.gz → 0.58.0__tar.gz

{emmet-builders-0.57.2 → emmet-builders-0.58.0}/PKG-INFO

@@ -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

{emmet-builders-0.57.2 → emmet-builders-0.58.0}/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-0.57.2 → emmet-builders-0.58.0}/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 → emmet-builders-0.58.0}/emmet_builders.egg-info/PKG-INFO

@@ -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

{emmet-builders-0.57.2 → emmet-builders-0.58.0}/emmet_builders.egg-info/requires.txt

@@ -1,5 +1,5 @@
 emmet-core[all]
-maggma>=0.49.17
+maggma>=0.51.9
 matminer>=0.7.3
 
 [docs]

{emmet-builders-0.57.2 → emmet-builders-0.58.0}/requirements/deployment.txt

@@ -14,9 +14,9 @@ bcrypt==4.0.1
     # via paramiko
 blinker==1.6.2
     # via flask
-boto3==1.26.158
+boto3==1.26.160
     # via maggma
-botocore==1.29.158
+botocore==1.29.160
     # via
     #   boto3
     #   s3transfer
@@ -42,18 +42,20 @@ dnspython==2.3.0
     # via
     #   maggma
     #   pymongo
-emmet-core[all]==0.57.1
+emmet-core[all]==0.57.2
     # via
     #   emmet-builders (emmet/emmet-builders/setup.py)
     #   mp-api
 exceptiongroup==1.1.1
     # via anyio
-fastapi==0.97.0
+fastapi==0.98.0
     # via maggma
 flask==2.3.2
     # via mongogrant
 fonttools==4.40.0
     # via matplotlib
+frozendict==2.3.8
+    # via pymatgen
 future==0.18.3
     # via
     #   matminer
@@ -82,7 +84,7 @@ kiwisolver==1.4.4
     # via matplotlib
 latexcodec==2.0.1
     # via pybtex
-maggma==0.51.8
+maggma==0.51.9
     # via emmet-builders (emmet/emmet-builders/setup.py)
 markupsafe==2.1.3
     # via
@@ -172,7 +174,7 @@ pydantic==1.10.9
     #   maggma
 pydash==7.0.4
     # via maggma
-pymatgen==2023.5.31
+pymatgen==2023.6.23
     # via
     #   emmet-core
     #   matminer
@@ -225,7 +227,7 @@ s3transfer==0.6.1
     # via boto3
 scikit-learn==1.2.2
     # via matminer
-scipy==1.10.1
+scipy==1.11.0
     # via
     #   -r python/requirements.txt
     #   pymatgen

{emmet-builders-0.57.2 → emmet-builders-0.58.0}/requirements/macos-latest_py3.10.txt

@@ -14,9 +14,9 @@ bcrypt==4.0.1
     # via paramiko
 blinker==1.6.2
     # via flask
-boto3==1.26.158
+boto3==1.26.160
     # via maggma
-botocore==1.29.158
+botocore==1.29.160
     # via
     #   boto3
     #   s3transfer
@@ -42,18 +42,20 @@ dnspython==2.3.0
     # via
     #   maggma
     #   pymongo
-emmet-core[all]==0.57.1
+emmet-core[all]==0.57.2
     # via
     #   emmet-builders (setup.py)
     #   mp-api
 exceptiongroup==1.1.1
     # via anyio
-fastapi==0.97.0
+fastapi==0.98.0
     # via maggma
 flask==2.3.2
     # via mongogrant
 fonttools==4.40.0
     # via matplotlib
+frozendict==2.3.8
+    # via pymatgen
 future==0.18.3
     # via
     #   matminer
@@ -82,7 +84,7 @@ kiwisolver==1.4.4
     # via matplotlib
 latexcodec==2.0.1
     # via pybtex
-maggma==0.51.8
+maggma==0.51.9
     # via emmet-builders (setup.py)
 markupsafe==2.1.3
     # via
@@ -168,7 +170,7 @@ pydantic==1.10.9
     #   maggma
 pydash==7.0.4
     # via maggma
-pymatgen==2023.5.31
+pymatgen==2023.6.23
     # via
     #   emmet-core
     #   matminer
@@ -221,7 +223,7 @@ s3transfer==0.6.1
     # via boto3
 scikit-learn==1.2.2
     # via matminer
-scipy==1.10.1
+scipy==1.11.0
     # via
     #   pymatgen
     #   robocrys

{emmet-builders-0.57.2 → emmet-builders-0.58.0}/requirements/macos-latest_py3.10_extras.txt

@@ -14,9 +14,9 @@ bcrypt==4.0.1
     # via paramiko
 blinker==1.6.2
     # via flask
-boto3==1.26.158
+boto3==1.26.160
     # via maggma
-botocore==1.29.158
+botocore==1.29.160
     # via
     #   boto3
     #   s3transfer
@@ -55,7 +55,7 @@ dnspython==2.3.0
     # via
     #   maggma
     #   pymongo
-emmet-core[all]==0.57.1
+emmet-core[all]==0.57.2
     # via
     #   emmet-builders (setup.py)
     #   mp-api
@@ -63,7 +63,7 @@ exceptiongroup==1.1.1
     # via
     #   anyio
     #   pytest
-fastapi==0.97.0
+fastapi==0.98.0
     # via maggma
 filelock==3.12.2
     # via virtualenv
@@ -73,6 +73,8 @@ flask==2.3.2
     # via mongogrant
 fonttools==4.40.0
     # via matplotlib
+frozendict==2.3.8
+    # via pymatgen
 future==0.18.3
     # via
     #   matminer
@@ -120,7 +122,7 @@ latexcodec==2.0.1
     # via pybtex
 livereload==2.6.3
     # via emmet-builders (setup.py)
-maggma==0.51.8
+maggma==0.51.9
     # via emmet-builders (setup.py)
 markdown==3.3.7
     # via
@@ -193,7 +195,7 @@ msgpack==1.0.5
     # via
     #   maggma
     #   mp-api
-mypy==1.4.0
+mypy==1.4.1
     # via emmet-builders (setup.py)
 mypy-extensions==1.0.0
     # via
@@ -240,7 +242,7 @@ paramiko==3.2.0
     # via sshtunnel
 pillow==9.5.0
     # via matplotlib
-platformdirs==3.7.0
+platformdirs==3.8.0
     # via virtualenv
 plotly==5.15.0
     # via pymatgen
@@ -275,7 +277,7 @@ pyflakes==3.0.1
     # via flake8
 pygments==2.15.1
     # via mkdocs-material
-pymatgen==2023.5.31
+pymatgen==2023.6.23
     # via
     #   emmet-core
     #   matminer
@@ -302,7 +304,7 @@ pyparsing==3.1.0
     # via matplotlib
 pyrsistent==0.19.3
     # via jsonschema
-pytest==7.3.2
+pytest==7.4.0
     # via
     #   emmet-builders (setup.py)
     #   pytest-cov
@@ -347,7 +349,7 @@ s3transfer==0.6.1
     # via boto3
 scikit-learn==1.2.2
     # via matminer
-scipy==1.10.1
+scipy==1.11.0
     # via
     #   pymatgen
     #   robocrys

{emmet-builders-0.57.2 → emmet-builders-0.58.0}/requirements/macos-latest_py3.11.txt

@@ -14,9 +14,9 @@ bcrypt==4.0.1
     # via paramiko
 blinker==1.6.2
     # via flask
-boto3==1.26.158
+boto3==1.26.160
     # via maggma
-botocore==1.29.158
+botocore==1.29.160
     # via
     #   boto3
     #   s3transfer
@@ -42,16 +42,18 @@ dnspython==2.3.0
     # via
     #   maggma
     #   pymongo
-emmet-core[all]==0.57.1
+emmet-core[all]==0.57.2
     # via
     #   emmet-builders (setup.py)
     #   mp-api
-fastapi==0.97.0
+fastapi==0.98.0
     # via maggma
 flask==2.3.2
     # via mongogrant
 fonttools==4.40.0
     # via matplotlib
+frozendict==2.3.8
+    # via pymatgen
 future==0.18.3
     # via
     #   matminer
@@ -80,7 +82,7 @@ kiwisolver==1.4.4
     # via matplotlib
 latexcodec==2.0.1
     # via pybtex
-maggma==0.51.8
+maggma==0.51.9
     # via emmet-builders (setup.py)
 markupsafe==2.1.3
     # via
@@ -166,7 +168,7 @@ pydantic==1.10.9
     #   maggma
 pydash==7.0.4
     # via maggma
-pymatgen==2023.5.31
+pymatgen==2023.6.23
     # via
     #   emmet-core
     #   matminer
@@ -219,7 +221,7 @@ s3transfer==0.6.1
     # via boto3
 scikit-learn==1.2.2
     # via matminer
-scipy==1.10.1
+scipy==1.11.0
     # via
     #   pymatgen
     #   robocrys

{emmet-builders-0.57.2 → emmet-builders-0.58.0}/requirements/macos-latest_py3.11_extras.txt

@@ -14,9 +14,9 @@ bcrypt==4.0.1
     # via paramiko
 blinker==1.6.2
     # via flask
-boto3==1.26.158
+boto3==1.26.160
     # via maggma
-botocore==1.29.158
+botocore==1.29.160
     # via
     #   boto3
     #   s3transfer
@@ -55,11 +55,11 @@ dnspython==2.3.0
     # via
     #   maggma
     #   pymongo
-emmet-core[all]==0.57.1
+emmet-core[all]==0.57.2
     # via
     #   emmet-builders (setup.py)
     #   mp-api
-fastapi==0.97.0
+fastapi==0.98.0
     # via maggma
 filelock==3.12.2
     # via virtualenv
@@ -69,6 +69,8 @@ flask==2.3.2
     # via mongogrant
 fonttools==4.40.0
     # via matplotlib
+frozendict==2.3.8
+    # via pymatgen
 future==0.18.3
     # via
     #   matminer
@@ -116,7 +118,7 @@ latexcodec==2.0.1
     # via pybtex
 livereload==2.6.3
     # via emmet-builders (setup.py)
-maggma==0.51.8
+maggma==0.51.9
     # via emmet-builders (setup.py)
 markdown==3.3.7
     # via
@@ -189,7 +191,7 @@ msgpack==1.0.5
     # via
     #   maggma
     #   mp-api
-mypy==1.4.0
+mypy==1.4.1
     # via emmet-builders (setup.py)
 mypy-extensions==1.0.0
     # via
@@ -236,7 +238,7 @@ paramiko==3.2.0
     # via sshtunnel
 pillow==9.5.0
     # via matplotlib
-platformdirs==3.7.0
+platformdirs==3.8.0
     # via virtualenv
 plotly==5.15.0
     # via pymatgen
@@ -271,7 +273,7 @@ pyflakes==3.0.1
     # via flake8
 pygments==2.15.1
     # via mkdocs-material
-pymatgen==2023.5.31
+pymatgen==2023.6.23
     # via
     #   emmet-core
     #   matminer
@@ -298,7 +300,7 @@ pyparsing==3.1.0
     # via matplotlib
 pyrsistent==0.19.3
     # via jsonschema
-pytest==7.3.2
+pytest==7.4.0
     # via
     #   emmet-builders (setup.py)
     #   pytest-cov
@@ -343,7 +345,7 @@ s3transfer==0.6.1
     # via boto3
 scikit-learn==1.2.2
     # via matminer
-scipy==1.10.1
+scipy==1.11.0
     # via
     #   pymatgen
     #   robocrys

{emmet-builders-0.57.2 → emmet-builders-0.58.0}/requirements/macos-latest_py3.8.txt

@@ -14,9 +14,9 @@ bcrypt==4.0.1
     # via paramiko
 blinker==1.6.2
     # via flask
-boto3==1.26.158
+boto3==1.26.160
     # via maggma
-botocore==1.29.158
+botocore==1.29.160
     # via
     #   boto3
     #   s3transfer
@@ -42,18 +42,20 @@ dnspython==2.3.0
     # via
     #   maggma
     #   pymongo
-emmet-core[all]==0.57.1
+emmet-core[all]==0.57.2
     # via
     #   emmet-builders (setup.py)
     #   mp-api
 exceptiongroup==1.1.1
     # via anyio
-fastapi==0.97.0
+fastapi==0.98.0
     # via maggma
 flask==2.3.2
     # via mongogrant
 fonttools==4.40.0
     # via matplotlib
+frozendict==2.3.8
+    # via pymatgen
 future==0.18.3
     # via
     #   matminer
@@ -88,7 +90,7 @@ kiwisolver==1.4.4
     # via matplotlib
 latexcodec==2.0.1
     # via pybtex
-maggma==0.51.8
+maggma==0.51.9
     # via emmet-builders (setup.py)
 markupsafe==2.1.3
     # via
@@ -176,7 +178,7 @@ pydantic==1.10.9
     #   maggma
 pydash==7.0.4
     # via maggma
-pymatgen==2023.5.31
+pymatgen==2023.6.23
     # via
     #   emmet-core
     #   matminer