emmet-builders 0.84.7rc3__py3-none-any.whl → 0.84.8__py3-none-any.whl

emmet/builders/materials/electronic_structure.py

@@ -1,28 +1,28 @@
-from collections import defaultdict
-from math import ceil
 import itertools
 import re
+from collections import defaultdict
+from math import ceil
+
 import boto3
 import numpy as np
 from botocore.handlers import disable_signing
 from maggma.builders import Builder
 from maggma.utils import grouper
 from pymatgen.analysis.magnetism.analyzer import CollinearMagneticStructureAnalyzer
+from pymatgen.analysis.structure_matcher import StructureMatcher
 from pymatgen.core import Structure
-from pymatgen.electronic_structure.core import Spin
 from pymatgen.electronic_structure.bandstructure import BandStructureSymmLine
+from pymatgen.electronic_structure.core import Spin
 from pymatgen.electronic_structure.dos import CompleteDos
-from pymatgen.symmetry.bandstructure import HighSymmKpath
-from pymatgen.analysis.structure_matcher import StructureMatcher
-from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 from pymatgen.io.vasp.sets import MPStaticSet
+from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
+from pymatgen.symmetry.bandstructure import HighSymmKpath
 
-from emmet.core.settings import EmmetSettings
+from emmet.builders.utils import query_open_data
 from emmet.core.electronic_structure import ElectronicStructureDoc
+from emmet.core.settings import EmmetSettings
 from emmet.core.utils import jsanitize
 
-from emmet.builders.utils import query_open_data
-
 SETTINGS = EmmetSettings()
 
 
@@ -297,7 +297,7 @@ class ElectronicStructureBuilder(Builder):
         Inserts electronic structure documents into the electronic_structure collection
 
         Args:
-            items ([Dict]): A list of ElectronicStructureDoc dictionaries to update
+            items ([dict]): A list of ElectronicStructureDoc dictionaries to update
         """
 
         items = list(filter(None, items))
@@ -577,9 +577,9 @@ class ElectronicStructureBuilder(Builder):
 
                 other_calcs.append(
                     {
-                        "is_static": True
-                        if "Static" in mat["task_types"][task_id]
-                        else False,
+                        "is_static": (
+                            True if "Static" in mat["task_types"][task_id] else False
+                        ),
                         "task_id": task_id,
                         "is_hubbard": int(is_hubbard),
                         "nkpoints": int(nkpoints),
@@ -640,9 +640,9 @@ class ElectronicStructureBuilder(Builder):
                     bs_obj["data"] if bs_obj is not None else None
                 )
 
-                materials_doc["bandstructure"][bs_type][
-                    "output_structure"
-                ] = sorted_bs_data[0]["output_structure"]
+                materials_doc["bandstructure"][bs_type]["output_structure"] = (
+                    sorted_bs_data[0]["output_structure"]
+                )
 
                 materials_doc["origins"].append(
                     {

emmet/builders/materials/magnetism.py

@@ -1,5 +1,5 @@
 from math import ceil
-from typing import Dict, Iterator, Optional
+from typing import Iterator
 
 from maggma.builders import Builder
 from maggma.stores import Store
@@ -18,7 +18,7 @@ class MagneticBuilder(Builder):
         materials: Store,
         magnetism: Store,
         tasks: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs,
     ):
         """
@@ -42,7 +42,7 @@ class MagneticBuilder(Builder):
 
         super().__init__(sources=[materials, tasks], targets=[magnetism], **kwargs)
 
-    def prechunk(self, number_splits: int) -> Iterator[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterator[dict]:  # pragma: no cover
         """
         Prechunk method to perform chunking by the key field
         """

emmet/builders/materials/ml.py

@@ -1,5 +1,4 @@
 from importlib.metadata import version
-from typing import TYPE_CHECKING, Optional, Union
 
 from maggma.builders.map_builder import MapBuilder
 from maggma.core import Store
@@ -16,13 +15,12 @@ from pymatgen.core import Structure
 from emmet.core.ml import MLDoc
 from emmet.core.utils import jsanitize
 
-if TYPE_CHECKING:
-    try:
-        from ase.calculators.calculator import Calculator
+try:
+    from ase.calculators.calculator import Calculator
 
-        ase_installed = True
-    except ImportError:
-        ase_installed = False
+    ase_installed = True
+except ImportError:
+    ase_installed = False
 
 
 class MLBuilder(MapBuilder):
@@ -30,10 +28,10 @@ class MLBuilder(MapBuilder):
         self,
         materials: Store,
         ml_potential: Store,
-        model: Union[str, "Calculator"],
-        model_kwargs: Optional[dict] = None,
-        prop_kwargs: Optional[dict] = None,
-        provenance: Optional[dict] = None,
+        model: str | Calculator,
+        model_kwargs: dict | None = None,
+        prop_kwargs: dict | None = None,
+        provenance: dict | None = None,
         **kwargs,
     ):
         """Machine learning interatomic potential builder.

emmet/builders/materials/optimade.py

@@ -1,5 +1,5 @@
 from math import ceil
-from typing import Dict, Iterator, Optional
+from typing import Iterator
 
 from maggma.builders import Builder
 from maggma.core import Store
@@ -16,7 +16,7 @@ class OptimadeMaterialsBuilder(Builder):
         materials: Store,
         thermo: Store,
         optimade: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs,
     ):
         """
@@ -42,7 +42,7 @@ class OptimadeMaterialsBuilder(Builder):
 
         super().__init__(sources=[materials, thermo], targets=optimade, **kwargs)
 
-    def prechunk(self, number_splits: int) -> Iterator[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterator[dict]:  # pragma: no cover
         """
         Prechunk method to perform chunking by the key field
         """

emmet/builders/materials/piezoelectric.py

@@ -1,5 +1,4 @@
 from math import ceil
-from typing import Dict, Optional
 
 import numpy as np
 from maggma.builders import Builder
@@ -17,7 +16,7 @@ class PiezoelectricBuilder(Builder):
         materials: Store,
         tasks: Store,
         piezoelectric: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs,
     ):
         self.materials = materials

emmet/builders/materials/provenance.py

@@ -1,7 +1,7 @@
 from collections import defaultdict
 from datetime import datetime
 from math import ceil
-from typing import Dict, Iterable, List, Optional, Tuple
+from typing import Iterable
 
 from maggma.core import Builder, Store
 from maggma.utils import grouper
@@ -18,9 +18,9 @@ class ProvenanceBuilder(Builder):
         self,
         materials: Store,
         provenance: Store,
-        source_snls: List[Store],
-        settings: Optional[EmmetBuildSettings] = None,
-        query: Optional[Dict] = None,
+        source_snls: list[Store],
+        settings: EmmetBuildSettings | None = None,
+        query: dict | None = None,
         **kwargs,
     ):
         """
@@ -59,7 +59,7 @@ class ProvenanceBuilder(Builder):
             s.ensure_index("snl_id")
             s.ensure_index("formula_pretty")
 
-    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
         self.ensure_indicies()
 
         # Find all formulas for materials that have been updated since this
@@ -101,7 +101,7 @@ class ProvenanceBuilder(Builder):
         for chunk in grouper(mat_ids, N):
             yield {"query": {"material_id": {"$in": chunk}}}
 
-    def get_items(self) -> Tuple[List[Dict], List[Dict]]:  # type: ignore
+    def get_items(self) -> tuple[list[dict], list[dict]]:  # type: ignore
         """
         Gets all materials to assocaite with SNLs
         Returns:
@@ -178,7 +178,7 @@ class ProvenanceBuilder(Builder):
             self.logger.debug(f"Found {len(snl_structs)} potential snls for {mat_id}")
             yield mat, snl_structs
 
-    def process_item(self, item) -> Dict:
+    def process_item(self, item) -> dict:
         """
         Matches SNLS and Materials
         Args:

emmet/builders/materials/robocrys.py

@@ -1,8 +1,7 @@
-from typing import Dict, Optional
 from maggma.builders.map_builder import MapBuilder
 from maggma.core import Store
-
 from pymatgen.core.structure import Structure
+
 from emmet.core.robocrys import RobocrystallogapherDoc
 from emmet.core.utils import jsanitize
 
@@ -12,7 +11,7 @@ class RobocrystallographerBuilder(MapBuilder):
         self,
         oxidation_states: Store,
         robocrys: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs
     ):
         self.oxidation_states = oxidation_states

emmet/builders/materials/substrates.py

@@ -1,14 +1,15 @@
-from typing import Optional, Dict, Iterable
-from emmet.core.mpid import MPID
-from maggma.core.store import Store
+from typing import Iterable
+
 from maggma.core.builder import Builder
-from pymatgen.core.structure import Structure
+from maggma.core.store import Store
+from maggma.utils import grouper
 from pymatgen.analysis.elasticity.elastic import ElasticTensor
+from pymatgen.core.structure import Structure
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 
+from emmet.core.mpid import MPID
 from emmet.core.substrates import SubstratesDoc
 from emmet.core.utils import jsanitize
-from maggma.utils import grouper
 
 
 class SubstratesBuilder(Builder):
@@ -17,7 +18,7 @@ class SubstratesBuilder(Builder):
         materials: Store,
         substrates: Store,
         elasticity: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs,
     ):
         """
@@ -47,7 +48,7 @@ class SubstratesBuilder(Builder):
             **kwargs,
         )
 
-    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
         to_process_mat_ids = self._find_to_process()
 
         return [

emmet/builders/materials/thermo.py

@@ -1,8 +1,8 @@
-from math import ceil
 import warnings
-from itertools import chain
-from typing import Dict, Iterator, List, Optional, Set
 from datetime import datetime
+from itertools import chain
+from math import ceil
+from typing import Iterator, Set
 
 from maggma.core import Builder, Store
 from maggma.stores import S3Store
@@ -12,7 +12,7 @@ from pymatgen.analysis.phase_diagram import PhaseDiagramError
 from pymatgen.entries.computed_entries import ComputedStructureEntry
 
 from emmet.builders.utils import HiddenPrints
-from emmet.core.thermo import ThermoDoc, PhaseDiagramDoc
+from emmet.core.thermo import PhaseDiagramDoc, ThermoDoc
 from emmet.core.utils import jsanitize
 
 
@@ -21,9 +21,9 @@ class ThermoBuilder(Builder):
         self,
         thermo: Store,
         corrected_entries: Store,
-        phase_diagram: Optional[Store] = None,
-        query: Optional[Dict] = None,
-        num_phase_diagram_eles: Optional[int] = None,
+        phase_diagram: Store | None = None,
+        query: dict | None = None,
+        num_phase_diagram_eles: int | None = None,
         chunk_size: int = 1000,
         **kwargs,
     ):
@@ -111,7 +111,7 @@ class ThermoBuilder(Builder):
             coll.ensure_index("chemsys")
             coll.ensure_index("phase_diagram_id")
 
-    def prechunk(self, number_splits: int) -> Iterator[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterator[dict]:  # pragma: no cover
         to_process_chemsys = self._get_chemsys_to_process()
 
         N = ceil(len(to_process_chemsys) / number_splits)
@@ -119,7 +119,7 @@ class ThermoBuilder(Builder):
         for chemsys_chunk in grouper(to_process_chemsys, N):
             yield {"query": {"chemsys": {"$in": list(chemsys_chunk)}}}
 
-    def get_items(self) -> Iterator[List[Dict]]:
+    def get_items(self) -> Iterator[list[dict]]:
         """
         Gets whole chemical systems of entries to process
         """
@@ -224,7 +224,7 @@ class ThermoBuilder(Builder):
         """
         Inserts the thermo and phase diagram docs into the thermo collection
         Args:
-            items ([[tuple(List[dict],List[dict])]]): a list of a list of thermo and phase diagram dict pairs to update
+            items ([[tuple(list[dict],list[dict])]]): a list of a list of thermo and phase diagram dict pairs to update
         """
 
         thermo_docs = [pair[0] for pair_list in items for pair in pair_list]

emmet/builders/matscholar/missing_compositions.py

@@ -1,10 +1,10 @@
-from itertools import combinations
 import itertools
+from itertools import combinations
 from math import ceil
-from typing import Dict, List, Iterator, Optional
+from typing import Iterator
 
 from maggma.core import Builder
-from maggma.stores import S3Store, MongoURIStore, MongoStore
+from maggma.stores import MongoStore, MongoURIStore, S3Store
 from maggma.utils import grouper
 from pymatgen.core import Composition, Element
 
@@ -21,7 +21,7 @@ class MissingCompositionsBuilder(Builder):
         phase_diagram: S3Store,
         mpcontribs: MongoURIStore,
         missing_compositions: MongoStore,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs,
     ):
         """
@@ -47,7 +47,7 @@ class MissingCompositionsBuilder(Builder):
             **kwargs,
         )
 
-    def prechunk(self, number_splits: int) -> Iterator[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterator[dict]:  # pragma: no cover
         """
         Prechunk method to perform chunking by the key field
         """
@@ -61,7 +61,7 @@ class MissingCompositionsBuilder(Builder):
         for split in grouper(keys, N):
             yield {"query": {self.phase_diagram.key: {"$in": list(split)}}}
 
-    def get_items(self) -> Iterator[Dict]:
+    def get_items(self) -> Iterator[dict]:
         """
         Returns all chemical systems (combinations of elements)
         to process.
@@ -112,7 +112,7 @@ class MissingCompositionsBuilder(Builder):
                 self.logger.error(f"Erro looking for phase diagram for {sys}: {ex}")
                 continue
 
-    def process_item(self, item: Dict) -> Dict:
+    def process_item(self, item: dict) -> dict:
         """
         Processes a chemical system and finds missing c
         ompositions for that system.
@@ -195,7 +195,7 @@ class MissingCompositionsBuilder(Builder):
         else:
             self.logger.info("No items to update")
 
-    def _get_entries_in_chemsys(self, chemsys) -> List:
+    def _get_entries_in_chemsys(self, chemsys) -> list:
         """Queries the MPContribs Store for entries in a chemical system."""
         # get sub-systems
         chemsys_subsystems = []

emmet/builders/mobility/migration_graph.py

@@ -1,10 +1,10 @@
 from maggma.builders.map_builder import MapBuilder
 from maggma.stores import MongoStore
-from typing import Tuple
-from emmet.core.mobility.migrationgraph import MigrationGraphDoc
-from emmet.builders.utils import get_hop_cutoff
-from pymatgen.apps.battery.insertion_battery import InsertionElectrode
 from pymatgen.analysis.diffusion.neb.full_path_mapper import MigrationGraph
+from pymatgen.apps.battery.insertion_battery import InsertionElectrode
+
+from emmet.builders.utils import get_hop_cutoff
+from emmet.core.mobility.migrationgraph import MigrationGraphDoc
 from emmet.core.utils import jsanitize
 
 
@@ -18,7 +18,7 @@ class MigrationGraphBuilder(MapBuilder):
         max_hop_distance: float = 7,
         populate_sc_fields: bool = True,
         min_length_sc: float = 8,
-        minmax_num_atoms: Tuple[int, int] = (80, 120),
+        minmax_num_atoms: tuple[int, int] = (80, 120),
         ltol: float = 0.2,
         stol: float = 0.3,
         angle_tol: float = 5,

emmet/builders/molecules/atomic.py

@@ -2,23 +2,22 @@ from collections import defaultdict
 from datetime import datetime
 from itertools import chain
 from math import ceil
-from typing import Optional, Iterable, Iterator, List, Dict
+from typing import Iterable, Iterator
 
 from maggma.builders import Builder
 from maggma.core import Store
 from maggma.utils import grouper
 
-from emmet.core.qchem.task import TaskDocument
-from emmet.core.qchem.molecule import MoleculeDoc, evaluate_lot
+from emmet.builders.settings import EmmetBuildSettings
 from emmet.core.molecules.atomic import (
-    PartialChargesDoc,
-    PartialSpinsDoc,
     CHARGES_METHODS,
     SPINS_METHODS,
+    PartialChargesDoc,
+    PartialSpinsDoc,
 )
+from emmet.core.qchem.molecule import MoleculeDoc, evaluate_lot
+from emmet.core.qchem.task import TaskDocument
 from emmet.core.utils import jsanitize
-from emmet.builders.settings import EmmetBuildSettings
-
 
 __author__ = "Evan Spotte-Smith"
 
@@ -55,9 +54,9 @@ class PartialChargesBuilder(Builder):
         tasks: Store,
         molecules: Store,
         charges: Store,
-        query: Optional[Dict] = None,
-        methods: Optional[List] = None,
-        settings: Optional[EmmetBuildSettings] = None,
+        query: dict | None = None,
+        methods: list | None = None,
+        settings: EmmetBuildSettings | None = None,
         **kwargs,
     ):
         self.tasks = tasks
@@ -105,7 +104,7 @@ class PartialChargesBuilder(Builder):
         self.charges.ensure_index("last_updated")
         self.charges.ensure_index("formula_alphabetical")
 
-    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
         """Prechunk the builder for distributed computation"""
 
         temp_query = dict(self.query)
@@ -131,7 +130,7 @@ class PartialChargesBuilder(Builder):
             query["species_hash"] = {"$in": list(hash_chunk)}
             yield {"query": query}
 
-    def get_items(self) -> Iterator[List[Dict]]:
+    def get_items(self) -> Iterator[list[dict]]:
         """
         Gets all items to process into partial charges documents.
         This does no datetime checking; relying on on whether
@@ -178,12 +177,12 @@ class PartialChargesBuilder(Builder):
 
             yield molecules
 
-    def process_item(self, items: List[Dict]) -> List[Dict]:
+    def process_item(self, items: list[dict]) -> list[dict]:
         """
         Process the tasks into PartialChargesDocs
 
         Args:
-            tasks List[Dict] : a list of MoleculeDocs in dict form
+            tasks list[dict] : a list of MoleculeDocs in dict form
 
         Returns:
             [dict] : a list of new partial charges docs
@@ -275,7 +274,7 @@ class PartialChargesBuilder(Builder):
 
         return jsanitize([doc.model_dump() for doc in charges_docs], allow_bson=True)
 
-    def update_targets(self, items: List[List[Dict]]):
+    def update_targets(self, items: list[list[dict]]):
         """
         Inserts the new documents into the charges collection
 
@@ -334,9 +333,9 @@ class PartialSpinsBuilder(Builder):
         tasks: Store,
         molecules: Store,
         spins: Store,
-        query: Optional[Dict] = None,
-        methods: Optional[List] = None,
-        settings: Optional[EmmetBuildSettings] = None,
+        query: dict | None = None,
+        methods: list | None = None,
+        settings: EmmetBuildSettings | None = None,
         **kwargs,
     ):
         self.tasks = tasks
@@ -384,7 +383,7 @@ class PartialSpinsBuilder(Builder):
         self.spins.ensure_index("last_updated")
         self.spins.ensure_index("formula_alphabetical")
 
-    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
         """Prechunk the builder for distributed computation"""
 
         temp_query = dict(self.query)
@@ -410,7 +409,7 @@ class PartialSpinsBuilder(Builder):
             query["species_hash"] = {"$in": list(hash_chunk)}
             yield {"query": query}
 
-    def get_items(self) -> Iterator[List[Dict]]:
+    def get_items(self) -> Iterator[list[dict]]:
         """
         Gets all items to process into partial spins documents.
         This does no datetime checking; relying on on whether
@@ -457,12 +456,12 @@ class PartialSpinsBuilder(Builder):
 
             yield molecules
 
-    def process_item(self, items: List[Dict]) -> List[Dict]:
+    def process_item(self, items: list[dict]) -> list[dict]:
         """
         Process the tasks into PartialSpinsDocs
 
         Args:
-            tasks List[Dict] : a list of MoleculeDocs in dict form
+            tasks list[dict] : a list of MoleculeDocs in dict form
 
         Returns:
             [dict] : a list of new partial spins docs
@@ -555,7 +554,7 @@ class PartialSpinsBuilder(Builder):
 
         return jsanitize([doc.model_dump() for doc in spins_docs], allow_bson=True)
 
-    def update_targets(self, items: List[List[Dict]]):
+    def update_targets(self, items: list[list[dict]]):
         """
         Inserts the new documents into the spins collection
 

emmet/builders/molecules/bonds.py

@@ -2,18 +2,17 @@ from collections import defaultdict
 from datetime import datetime
 from itertools import chain
 from math import ceil
-from typing import Optional, Iterable, Iterator, List, Dict
+from typing import Iterable, Iterator
 
 from maggma.builders import Builder
 from maggma.core import Store
 from maggma.utils import grouper
 
-from emmet.core.qchem.task import TaskDocument
+from emmet.builders.settings import EmmetBuildSettings
+from emmet.core.molecules.bonds import BOND_METHODS, MoleculeBondingDoc
 from emmet.core.qchem.molecule import MoleculeDoc, evaluate_lot
-from emmet.core.molecules.bonds import MoleculeBondingDoc, BOND_METHODS
+from emmet.core.qchem.task import TaskDocument
 from emmet.core.utils import jsanitize
-from emmet.builders.settings import EmmetBuildSettings
-
 
 __author__ = "Evan Spotte-Smith"
 
@@ -54,9 +53,9 @@ class BondingBuilder(Builder):
         tasks: Store,
         molecules: Store,
         bonds: Store,
-        query: Optional[Dict] = None,
-        methods: Optional[List] = None,
-        settings: Optional[EmmetBuildSettings] = None,
+        query: dict | None = None,
+        methods: list | None = None,
+        settings: EmmetBuildSettings | None = None,
         **kwargs,
     ):
         self.tasks = tasks
@@ -104,7 +103,7 @@ class BondingBuilder(Builder):
         self.bonds.ensure_index("last_updated")
         self.bonds.ensure_index("formula_alphabetical")
 
-    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
         """Prechunk the builder for distributed computation"""
 
         temp_query = dict(self.query)
@@ -130,7 +129,7 @@ class BondingBuilder(Builder):
             query["species_hash"] = {"$in": list(hash_chunk)}
             yield {"query": query}
 
-    def get_items(self) -> Iterator[List[Dict]]:
+    def get_items(self) -> Iterator[list[dict]]:
         """
         Gets all items to process into bonding documents.
         This does no datetime checking; relying on on whether
@@ -177,12 +176,12 @@ class BondingBuilder(Builder):
 
             yield molecules
 
-    def process_item(self, items: List[Dict]) -> List[Dict]:
+    def process_item(self, items: list[dict]) -> list[dict]:
         """
         Process the tasks into MoleculeBondingDocs
 
         Args:
-            tasks List[Dict] : a list of MoleculeDocs in dict form
+            tasks list[dict] : a list of MoleculeDocs in dict form
 
         Returns:
             [dict] : a list of new bonding docs
@@ -292,7 +291,7 @@ class BondingBuilder(Builder):
 
         return jsanitize([doc.model_dump() for doc in bonding_docs], allow_bson=True)
 
-    def update_targets(self, items: List[List[Dict]]):
+    def update_targets(self, items: list[list[dict]]):
         """
         Inserts the new documents into the bonds collection