emmet-builders 0.78.3__py3-none-any.whl → 0.86.0__py3-none-any.whl

emmet/builders/materials/elasticity.py

@@ -17,8 +17,9 @@ The build proceeds in the below steps:
 7. Fit the elastic tensor.
 """
 
+from __future__ import annotations
+
 from datetime import datetime
-from typing import Any, Dict, Generator, List, Optional, Tuple, Union
 
 import numpy as np
 from maggma.core import Builder, Store
@@ -29,10 +30,18 @@ from pymatgen.core import Structure
 from pymatgen.core.tensors import TensorMapping
 
 from emmet.core.elasticity import ElasticityDoc
-from emmet.core.mpid import MPID
+from emmet.core.mpid import AlphaID
 from emmet.core.utils import jsanitize
 from emmet.core.vasp.calc_types import CalcType
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+    from typing import Any
+
+    from emmet.core.types.typing import IdentifierType
+
 
 class ElasticityBuilder(Builder):
     def __init__(
@@ -40,7 +49,7 @@ class ElasticityBuilder(Builder):
         tasks: Store,
         materials: Store,
         elasticity: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         fitting_method: str = "finite_difference",
         **kwargs,
     ):
@@ -78,7 +87,7 @@ class ElasticityBuilder(Builder):
 
     def get_items(
         self,
-    ) -> Generator[Tuple[str, Dict[str, str], List[Dict]], None, None]:
+    ) -> Generator[tuple[str, dict[str, str], list[dict]], None, None]:
         """
         Gets all items to process into elasticity docs.
 
@@ -129,8 +138,8 @@ class ElasticityBuilder(Builder):
             yield material_id, calc_types, tasks
 
     def process_item(
-        self, item: Tuple[MPID, Dict[str, str], List[Dict]]
-    ) -> Union[Dict, None]:
+        self, item: tuple[IdentifierType, dict[str, str], list[dict]]
+    ) -> dict | None:
         """
         Process all tasks belong to the same material into an elasticity doc.
 
@@ -191,7 +200,7 @@ class ElasticityBuilder(Builder):
 
         # convert to elasticity doc
         deforms = []
-        stresses = []
+        stresses: list[Stress] = []  # TODO: mypy misfires on `Stress`
         deform_task_ids = []
         deform_dir_names = []
         for doc in final_deform:
@@ -200,7 +209,7 @@ class ElasticityBuilder(Builder):
             )
             # 0.1 to convert to GPa from kBar, and the minus sign to flip the stress
             # direction from compressive as positive (in vasp) to tensile as positive
-            stresses.append(-0.1 * Stress(doc["output"]["stress"]))
+            stresses.append(-0.1 * Stress(doc["output"]["stress"]))  # type: ignore[arg-type]
             deform_task_ids.append(doc["task_id"])
             deform_dir_names.append(doc["dir_name"])
 
@@ -208,10 +217,10 @@ class ElasticityBuilder(Builder):
             structure=Structure.from_dict(final_opt["output"]["structure"]),
             material_id=material_id,
             deformations=deforms,
-            stresses=stresses,
+            stresses=stresses,  # type: ignore[arg-type]
             deformation_task_ids=deform_task_ids,
             deformation_dir_names=deform_dir_names,
-            equilibrium_stress=-0.1 * Stress(final_opt["output"]["stress"]),
+            equilibrium_stress=-0.1 * Stress(final_opt["output"]["stress"]),  # type: ignore[arg-type]
             optimization_task_id=final_opt["task_id"],
             optimization_dir_name=final_opt["dir_name"],
             fitting_method="finite_difference",
@@ -220,7 +229,7 @@ class ElasticityBuilder(Builder):
 
         return elasticity_doc
 
-    def update_targets(self, items: List[Dict]):
+    def update_targets(self, items: list[dict]):
         """
         Insert the new elasticity docs into the elasticity collection.
 
@@ -233,24 +242,26 @@ class ElasticityBuilder(Builder):
 
 
 def filter_opt_tasks(
-    tasks: List[Dict],
-    calc_types: Dict[str, str],
-    target_calc_type: str = CalcType.GGA_Structure_Optimization,
-) -> List[Dict]:
+    tasks: list[dict],
+    calc_types: dict[str, str],
+    target_calc_type: str | CalcType = CalcType.GGA_Structure_Optimization,
+) -> list[dict]:
     """
     Filter optimization tasks, by
         - calculation type
     """
-    opt_tasks = [t for t in tasks if calc_types[str(t["task_id"])] == target_calc_type]
+    opt_tasks = [
+        t for t in tasks if calc_types[str(AlphaID(t["task_id"]))] == target_calc_type
+    ]
 
     return opt_tasks
 
 
 def filter_deform_tasks(
-    tasks: List[Dict],
-    calc_types: Dict[str, str],
-    target_calc_type: str = CalcType.GGA_Deformation,
-) -> List[Dict]:
+    tasks: list[dict],
+    calc_types: dict[str, str],
+    target_calc_type: str | CalcType = CalcType.GGA_Deformation,
+) -> list[dict]:
     """
     Filter deformation tasks, by
         - calculation type
@@ -271,8 +282,8 @@ def filter_deform_tasks(
 
 
 def filter_by_incar_settings(
-    tasks: List[Dict], incar_settings: Optional[Dict[str, Any]] = None
-) -> List[Dict]:
+    tasks: list[dict], incar_settings: dict[str, Any] | None = None
+) -> list[dict]:
     """
     Filter tasks by incar parameters.
     """
@@ -315,7 +326,7 @@ def filter_by_incar_settings(
     return selected
 
 
-def filter_opt_tasks_by_time(tasks: List[Dict], logger) -> Dict:
+def filter_opt_tasks_by_time(tasks: list[dict], logger) -> dict:
     """
     Filter a set of tasks to select the latest completed one.
 
@@ -330,8 +341,8 @@ def filter_opt_tasks_by_time(tasks: List[Dict], logger) -> Dict:
 
 
 def filter_deform_tasks_by_time(
-    tasks: List[Dict], deform_comp_tol: float = 1e-5, logger=None
-) -> List[Dict]:
+    tasks: list[dict], deform_comp_tol: float = 1e-5, logger=None
+) -> list[dict]:
     """
     For deformation tasks with the same deformation, select the latest completed one.
 
@@ -364,7 +375,7 @@ def filter_deform_tasks_by_time(
     return selected
 
 
-def _filter_tasks_by_time(tasks: List[Dict], mode: str, logger) -> Dict:
+def _filter_tasks_by_time(tasks: list[dict], mode: str, logger) -> dict:
     """
     Helper function to filter a set of tasks to select the latest completed one.
     """
@@ -388,11 +399,11 @@ def _filter_tasks_by_time(tasks: List[Dict], mode: str, logger) -> Dict:
 
 
 def select_final_opt_deform_tasks(
-    opt_tasks: List[Tuple[np.ndarray, Dict]],
-    deform_tasks: List[Tuple[np.ndarray, List[Dict]]],
+    opt_tasks: list[tuple[np.ndarray, dict]],
+    deform_tasks: list[tuple[np.ndarray, list[dict]]],
     logger,
     lattice_comp_tol: float = 1e-5,
-) -> Tuple[Union[Dict, None], Union[List[Dict], None]]:
+) -> tuple[dict | None, list[dict] | None]:
     """
     Select the final opt task and deform tasks for fitting.
 
@@ -445,8 +456,8 @@ def select_final_opt_deform_tasks(
 
 
 def group_by_parent_lattice(
-    tasks: List[Dict], mode: str, lattice_comp_tol: float = 1e-5
-) -> List[Tuple[np.ndarray, List[Dict]]]:
+    tasks: list[dict], mode: str, lattice_comp_tol: float = 1e-5
+) -> list[tuple[np.ndarray, list[dict]]]:
     """
     Groups a set of task docs by parent lattice equivalence.
 
@@ -459,11 +470,11 @@ def group_by_parent_lattice(
         lattice_comp_tol: tolerance for comparing lattice equivalence.
 
     Returns:
-        [(lattice, List[tasks])]: each tuple gives the common parent lattice of a
+        [(lattice, list[tasks])]: each tuple gives the common parent lattice of a
             list of the structures before deformation (if any), and the list tasks
             from which the structures are taken.
     """
-    docs_by_lattice: List[Tuple[np.ndarray, List[Dict]]] = []
+    docs_by_lattice: list[tuple[np.ndarray, list[dict]]] = []
 
     for doc in tasks:
         sim_lattice = get(doc, "output.structure.lattice.matrix")

emmet/builders/materials/electrodes.py

@@ -1,22 +1,27 @@
 import operator
+from collections import defaultdict
 from datetime import datetime
 from functools import lru_cache
 from itertools import chain
 from math import ceil
-from typing import Any, Iterator, Dict, List, Optional
-from collections import defaultdict
 
 from maggma.builders import Builder
 from maggma.stores import MongoStore
 from maggma.utils import grouper
-from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
+from pymatgen.analysis.phase_diagram import Composition, PhaseDiagram
 from pymatgen.entries.compatibility import MaterialsProject2020Compatibility
-from pymatgen.analysis.phase_diagram import PhaseDiagram, Composition
+from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
 
-from emmet.core.electrode import InsertionElectrodeDoc, ConversionElectrodeDoc
+from emmet.builders.settings import EmmetBuildSettings
+from emmet.core.electrode import ConversionElectrodeDoc, InsertionElectrodeDoc
 from emmet.core.structure_group import StructureGroupDoc, _get_id_lexi
 from emmet.core.utils import jsanitize
-from emmet.builders.settings import EmmetBuildSettings
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+    from collections.abc import Iterator
 
 
 def s_hash(el):
@@ -85,7 +90,7 @@ class StructureGroupBuilder(Builder):
         materials: MongoStore,
         sgroups: MongoStore,
         working_ion: str,
-        query: Optional[dict] = None,
+        query: dict | None = None,
         ltol: float = default_build_settings.LTOL,
         stol: float = default_build_settings.STOL,
         angle_tol: float = default_build_settings.ANGLE_TOL,
@@ -115,15 +120,15 @@ class StructureGroupBuilder(Builder):
         self.check_newer = check_newer
         self.chunk_size = chunk_size
 
-        self.query[
-            "deprecated"
-        ] = False  # Ensure only non-deprecated materials are chosen
+        self.query["deprecated"] = (
+            False  # Ensure only non-deprecated materials are chosen
+        )
 
         super().__init__(
             sources=[materials], targets=[sgroups], chunk_size=chunk_size, **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
         """
@@ -268,7 +273,7 @@ class StructureGroupBuilder(Builder):
             else:
                 yield {"chemsys": chemsys, "materials": all_mats_in_chemsys}
 
-    def update_targets(self, items: List):
+    def update_targets(self, items: list):
         items = list(filter(None, chain.from_iterable(items)))
         if len(items) > 0:
             self.logger.info("Updating {} sgroups documents".format(len(items)))
@@ -319,7 +324,7 @@ class InsertionElectrodeBuilder(Builder):
         grouped_materials: MongoStore,
         thermo: MongoStore,
         insertion_electrode: MongoStore,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         strip_structures: bool = False,
         **kwargs,
     ):
@@ -335,7 +340,7 @@ class InsertionElectrodeBuilder(Builder):
             **kwargs,
         )
 
-    def prechunk(self, number_splits: int) -> Iterator[Dict]:
+    def prechunk(self, number_splits: int) -> Iterator[dict]:
         """
         Prechunk method to perform chunking by the key field
         """
@@ -418,7 +423,7 @@ class InsertionElectrodeBuilder(Builder):
             else:
                 yield None
 
-    def process_item(self, item) -> Dict:
+    def process_item(self, item) -> dict:
         """
         - Add volume information to each entry to create the insertion electrode document
         - Add the host structure
@@ -460,14 +465,14 @@ class InsertionElectrodeBuilder(Builder):
             # {"failed_reason": "unable to create InsertionElectrode document"}
         return jsanitize(ie.model_dump())
 
-    def update_targets(self, items: List):
+    def update_targets(self, items: list):
         items = list(filter(None, items))
         if len(items) > 0:
             self.logger.info("Updating {} battery documents".format(len(items)))
             for struct_group_dict in items:
-                struct_group_dict[
-                    self.grouped_materials.last_updated_field
-                ] = datetime.utcnow()
+                struct_group_dict[self.grouped_materials.last_updated_field] = (
+                    datetime.utcnow()
+                )
             self.insertion_electrode.update(docs=items, key=["battery_id"])
         else:
             self.logger.info("No items to update")
@@ -480,7 +485,7 @@ class ConversionElectrodeBuilder(Builder):
         conversion_electrode_store: MongoStore,
         working_ion: str,
         thermo_type: str,
-        query: Optional[dict] = None,
+        query: dict | None = None,
         **kwargs,
     ):
         self.phase_diagram_store = phase_diagram_store
@@ -499,7 +504,7 @@ class ConversionElectrodeBuilder(Builder):
             **kwargs,
         )
 
-    def prechunk(self, number_splits: int) -> Iterator[Dict]:
+    def prechunk(self, number_splits: int) -> Iterator[dict]:
         """
         Prechunk method to perform chunking by the key field
         """
@@ -532,7 +537,7 @@ class ConversionElectrodeBuilder(Builder):
         for phase_diagram_doc in self.phase_diagram_store.query(criteria=q):
             yield phase_diagram_doc
 
-    def process_item(self, item) -> Dict:
+    def process_item(self, item) -> dict:
         """
         - For each phase diagram doc, find all possible conversion electrodes and create conversion electrode docs
         """
@@ -564,14 +569,16 @@ class ConversionElectrodeBuilder(Builder):
                 # Get lowest material_id with matching composition
                 material_ids = [
                     (
-                        lambda x: x.data["material_id"]
-                        if x.composition.reduced_formula == v[1].reduced_formula
-                        else None
+                        lambda x: (
+                            x.data["material_id"]  # type: ignore[attr-defined]
+                            if x.composition.reduced_formula == v[1].reduced_formula
+                            else None
+                        )
                     )(e)
                     for e in pd.entries
                 ]
                 material_ids = list(filter(None, material_ids))
-                lowest_id = min(material_ids, key=_get_id_lexi)
+                lowest_id = min(material_ids, key=_get_id_lexi)  # type: ignore[arg-type]
                 conversion_electrode_doc = (
                     ConversionElectrodeDoc.from_composition_and_pd(
                         comp=v[1],
@@ -597,7 +604,7 @@ class ConversionElectrodeBuilder(Builder):
             relevant_entry_data = []
             for e in pd.entries:
                 if e.composition == Composition(c):
-                    relevant_entry_data.append((e.energy_per_atom, e.entry_id))
+                    relevant_entry_data.append((e.energy_per_atom, e.entry_id))  # type: ignore[attr-defined]
             relevant_entry_data.sort(key=lambda x: x[0])
             entry_id_mapping[c] = relevant_entry_data[0][1]
 
@@ -609,7 +616,7 @@ class ConversionElectrodeBuilder(Builder):
 
         return new_docs  # type: ignore
 
-    def update_targets(self, items: List):
+    def update_targets(self, items: list):
         combined_items = []
         for _items in items:
             _items = list(filter(None, _items))

emmet/builders/materials/electronic_structure.py

@@ -1,29 +1,29 @@
-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()
+SETTINGS = EmmetSettings()  # type: ignore[call-arg]
 
 
 class ElectronicStructureBuilder(Builder):
@@ -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,6 @@
+from __future__ import annotations
+
 from math import ceil
-from typing import Dict, Iterator, Optional
 
 from maggma.builders import Builder
 from maggma.stores import Store
@@ -7,8 +8,14 @@ from maggma.utils import grouper
 from pymatgen.core.structure import Structure
 
 from emmet.core.magnetism import MagnetismDoc
+from emmet.core.mpid import AlphaID
 from emmet.core.utils import jsanitize
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 __author__ = "Shyam Dwaraknath <shyamd@lbl.gov>, Matthew Horton <mkhorton@lbl.gov>"
 
 
@@ -18,7 +25,7 @@ class MagneticBuilder(Builder):
         materials: Store,
         magnetism: Store,
         tasks: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs,
     ):
         """
@@ -42,7 +49,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
         """
@@ -131,7 +138,7 @@ class MagneticBuilder(Builder):
 
         for origin in mat_doc["origins"]:
             if origin["name"] == "structure":
-                task_id = origin["task_id"]
+                task_id = str(AlphaID(origin["task_id"]))
 
         task_query = self.tasks.query_one(
             properties=["last_updated", "calcs_reversed"],

emmet/builders/materials/optimade.py

@@ -1,5 +1,4 @@
 from math import ceil
-from typing import Dict, Iterator, Optional
 
 from maggma.builders import Builder
 from maggma.core import Store
@@ -9,6 +8,11 @@ from pymatgen.core.structure import Structure
 from emmet.core.optimade import OptimadeMaterialsDoc
 from emmet.core.utils import jsanitize
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 
 class OptimadeMaterialsBuilder(Builder):
     def __init__(
@@ -16,7 +20,7 @@ class OptimadeMaterialsBuilder(Builder):
         materials: Store,
         thermo: Store,
         optimade: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs,
     ):
         """
@@ -42,7 +46,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
@@ -8,6 +7,7 @@ from maggma.utils import grouper
 from pymatgen.core.structure import Structure
 
 from emmet.core.polar import PiezoelectricDoc
+from emmet.core.mpid import AlphaID
 from emmet.core.utils import jsanitize
 
 
@@ -17,7 +17,7 @@ class PiezoelectricBuilder(Builder):
         materials: Store,
         tasks: Store,
         piezoelectric: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs,
     ):
         self.materials = materials
@@ -190,20 +190,19 @@ class PiezoelectricBuilder(Builder):
                     "last_updated",
                     "input.is_hubbard",
                     "orig_inputs.kpoints",
-                    "orig_inputs.poscar.structure",
+                    "orig_inputs.structure",
                     "input.parameters",
                     "input.structure",
-                    "output.piezo_tensor",
-                    "output.piezo_ionic_tensor",
+                    "calcs_reversed",
                     "output.bandgap",
                 ],
-                criteria={self.tasks.key: str(task_id)},
+                criteria={self.tasks.key: str(AlphaID(task_id))},
             )
             if task_query["output"]["bandgap"] > 0:
                 try:
-                    structure = task_query["orig_inputs"]["poscar"]["structure"]
-                except KeyError:
                     structure = task_query["input"]["structure"]
+                except KeyError:
+                    structure = task_query["orig_inputs"]["structure"]
 
                 is_hubbard = task_query["input"]["is_hubbard"]
 
@@ -223,19 +222,23 @@ class PiezoelectricBuilder(Builder):
                 lu_dt = mat_doc["last_updated"]
                 task_updated = task_query["last_updated"]
 
-                final_docs.append(
-                    {
-                        "task_id": task_id,
-                        "is_hubbard": int(is_hubbard),
-                        "nkpoints": int(nkpoints),
-                        "piezo_static": task_query["output"]["piezo_tensor"],
-                        "piezo_ionic": task_query["output"]["piezo_ionic_tensor"],
-                        "structure": structure,
-                        "updated_on": lu_dt,
-                        "task_updated": task_updated,
-                        self.materials.key: mat_doc[self.materials.key],
-                    }
-                )
+                if (cr := task_query.get("calcs_reversed", [])) and (
+                    outcar := cr[0].get("output", {}).get("outcar", {})
+                ):
+
+                    final_docs.append(
+                        {
+                            "task_id": task_id,
+                            "is_hubbard": int(is_hubbard),
+                            "nkpoints": int(nkpoints),
+                            "piezo_static": outcar.get("piezo_tensor"),
+                            "piezo_ionic": outcar.get("piezo_ionic_tensor"),
+                            "structure": structure,
+                            "updated_on": lu_dt,
+                            "task_updated": task_updated,
+                            self.materials.key: mat_doc[self.materials.key],
+                        }
+                    )
 
         if len(final_docs) > 0:
             sorted_final_docs = sorted(