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

emmet/builders/abinit/phonon.py

@@ -1,7 +1,6 @@
 import os
 import tempfile
 from math import ceil
-from typing import Dict, Iterator, List, Optional, Tuple
 
 import numpy as np
 from abipy.abio.inputs import AnaddbInput
@@ -25,8 +24,8 @@ from emmet.builders.settings import EmmetBuildSettings
 from emmet.core.phonon import (
     AbinitPhonon,
     Ddb,
-    PhononBandStructure,
-    PhononDos,
+    PhononBS,
+    PhononDOS,
     PhononWarnings,
     PhononWebsiteBS,
     ThermalDisplacement,
@@ -36,6 +35,11 @@ from emmet.core.phonon import (
 from emmet.core.polar import BornEffectiveCharges, DielectricDoc, IRDielectric
 from emmet.core.utils import jsanitize
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 SETTINGS = EmmetBuildSettings()
 
 
@@ -50,8 +54,8 @@ class PhononBuilder(Builder):
         ddb_files: Store,
         th_disp: Store,
         phonon_website: Store,
-        query: Optional[Dict] = None,
-        manager: Optional[TaskManager] = None,
+        query: dict | None = None,
+        manager: TaskManager | None = None,
         symprec: float = SETTINGS.SYMPREC,
         angle_tolerance: float = SETTINGS.ANGLE_TOL,
         chunk_size=100,
@@ -134,7 +138,7 @@ class PhononBuilder(Builder):
         for mpid_chunk in grouper(mats, N):
             yield {"query": {self.phonon_materials.key: {"$in": list(mpid_chunk)}}}
 
-    def get_items(self) -> Iterator[Dict]:
+    def get_items(self) -> Iterator[dict]:
         """
         Gets all materials that need phonons
 
@@ -188,7 +192,7 @@ class PhononBuilder(Builder):
 
             yield item
 
-    def process_item(self, item: Dict) -> Optional[Dict]:
+    def process_item(self, item: dict) -> dict | None:
         """
         Generates the full phonon document from an item
 
@@ -219,7 +223,6 @@ class PhononBuilder(Builder):
             becs = None
             if phonon_properties["becs"] is not None:
                 becs = BornEffectiveCharges(
-                    material_id=item["mp_id"],
                     symmetrized_value=phonon_properties["becs"],
                     value=sr_break["becs_nosymm"],
                     cnsr_break=sr_break["cnsr"],
@@ -241,25 +244,24 @@ class PhononBuilder(Builder):
                 abinit_input_vars=abinit_input_vars,
             )
 
-            phbs = PhononBandStructure(
-                material_id=item["mp_id"],
-                band_structure=phonon_properties["ph_bs"].as_dict(),
+            phbs = PhononBS(
+                identifier=item["mp_id"],
+                **phonon_properties["ph_bs"].as_dict(),
             )
 
-            phws = PhononWebsiteBS(
+            phws = PhononWebsiteBS(  # type: ignore[call-arg]
                 material_id=item["mp_id"],
                 phononwebsite=phonon_properties["ph_bs"].as_phononwebsite(),
             )
 
-            phdos = PhononDos(
-                material_id=item["mp_id"],
-                dos=phonon_properties["ph_dos"].as_dict(),
-                dos_method=phonon_properties["ph_dos_method"],
+            phdos = PhononDOS(
+                identifier=item["mp_id"],
+                **phonon_properties["ph_dos"].as_dict(),
             )
 
-            ddb = Ddb(material_id=item["mp_id"], ddb=item["ddb_str"])
+            ddb = Ddb(material_id=item["mp_id"], ddb=item["ddb_str"])  # type: ignore[call-arg]
 
-            th_disp = ThermalDisplacement(
+            th_disp = ThermalDisplacement(  # type: ignore[call-arg]
                 material_id=item["mp_id"],
                 structure=structure,
                 nsites=len(structure),
@@ -293,7 +295,7 @@ class PhononBuilder(Builder):
             )
             return None
 
-    def get_phonon_properties(self, item: Dict) -> Dict:
+    def get_phonon_properties(self, item: dict) -> dict:
         """
         Extracts the phonon properties from the item
         """
@@ -511,7 +513,7 @@ class PhononBuilder(Builder):
         dos: str = "tetra",
         lo_to_splitting: bool = True,
         use_dieflag: bool = True,
-    ) -> Tuple[AnaddbInput, Optional[List]]:
+    ) -> tuple[AnaddbInput, list | None]:
         """
         creates the AnaddbInput object to calculate the phonon properties.
         It also returns the list of qpoints labels for generating the PhononBandStructureSymmLine.
@@ -637,7 +639,7 @@ class PhononBuilder(Builder):
 
     @staticmethod
     def get_pmg_bs(
-        phbands: PhononBands, labels_list: List
+        phbands: PhononBands, labels_list: list
     ) -> PhononBandStructureSymmLine:
         """
         Generates a PhononBandStructureSymmLine starting from a abipy PhononBands object
@@ -726,7 +728,7 @@ class PhononBuilder(Builder):
 
         return data
 
-    def update_targets(self, items: List[Dict]):
+    def update_targets(self, items: list[dict]):
         """
         Inserts the new task_types into the task_types collection
 
@@ -769,7 +771,7 @@ class PhononBuilder(Builder):
 
 def get_warnings(
     asr_break: float, cnsr_break: float, ph_bs: PhononBandStructureSymmLine
-) -> List[PhononWarnings]:
+) -> list[PhononWarnings]:
     """
 
     Args:
@@ -810,7 +812,7 @@ def get_warnings(
 
 def get_thermodynamic_properties(
     ph_dos: CompletePhononDos,
-) -> Tuple[ThermodynamicProperties, VibrationalEnergy]:
+) -> tuple[ThermodynamicProperties, VibrationalEnergy]:
     """
     Calculates the thermodynamic properties from a phonon DOS
 
@@ -839,7 +841,7 @@ def get_thermodynamic_properties(
     zpe = ph_dos.zero_point_energy(ph_dos.structure)
 
     temperatures = temp.tolist()
-    tp = ThermodynamicProperties(temperatures=temperatures, cv=cv, entropy=entropy)
+    tp = ThermodynamicProperties(temperatures=temperatures, cv=cv, entropy=entropy)  # type: ignore[call-arg]
 
     ve = VibrationalEnergy(
         temperatures=temperatures,

emmet/builders/abinit/sound_velocity.py

@@ -1,18 +1,22 @@
 import tempfile
 import traceback
 from math import ceil
-from maggma.utils import grouper
-from typing import Optional, Dict, List, Iterator
 
-from abipy.dfpt.vsound import SoundVelocity as AbiSoundVelocity
 from abipy.dfpt.ddb import DdbFile
+from abipy.dfpt.vsound import SoundVelocity as AbiSoundVelocity
+from abipy.flowtk.tasks import TaskManager
 from maggma.builders import Builder
 from maggma.core import Store
-from abipy.flowtk.tasks import TaskManager
+from maggma.utils import grouper
 
 from emmet.core.phonon import SoundVelocity
 from emmet.core.utils import jsanitize
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 
 class SoundVelocityBuilder(Builder):
     def __init__(
@@ -20,8 +24,8 @@ class SoundVelocityBuilder(Builder):
         phonon_materials: Store,
         ddb_source: Store,
         sound_vel: Store,
-        query: Optional[dict] = None,
-        manager: Optional[TaskManager] = None,
+        query: dict | None = None,
+        manager: TaskManager | None = None,
         **kwargs
     ):
         """
@@ -72,7 +76,7 @@ class SoundVelocityBuilder(Builder):
         for mpid_chunk in grouper(mats, N):
             yield {"query": {self.phonon_materials.key: {"$in": list(mpid_chunk)}}}
 
-    def get_items(self) -> Iterator[Dict]:
+    def get_items(self) -> Iterator[dict]:
         """
         Gets all materials that need sound velocity.
 
@@ -117,7 +121,7 @@ class SoundVelocityBuilder(Builder):
 
             yield item
 
-    def process_item(self, item: Dict) -> Optional[Dict]:
+    def process_item(self, item: dict) -> dict | None:
         """
         Generates the sound velocity document from an item
 
@@ -132,7 +136,7 @@ class SoundVelocityBuilder(Builder):
         try:
             sound_vel_data = self.get_sound_vel(item)
 
-            sv = SoundVelocity(
+            sv = SoundVelocity(  # type: ignore[call-arg]
                 material_id=item["mp_id"],
                 structure=sound_vel_data["structure"],
                 directions=sound_vel_data["directions"],
@@ -153,7 +157,7 @@ class SoundVelocityBuilder(Builder):
             return None
 
     @staticmethod
-    def get_sound_vel(item: Dict) -> Dict:
+    def get_sound_vel(item: dict) -> dict:
         """
         Runs anaddb and return the extracted data for the speed of sound.
 
@@ -187,7 +191,7 @@ class SoundVelocityBuilder(Builder):
 
             return sv_data
 
-    def update_targets(self, items: List[Dict]):
+    def update_targets(self, items: list[dict]):
         """
         Inserts the new task_types into the task_types collection
 

emmet/builders/feff/xas.py

@@ -1,7 +1,6 @@
 import traceback
 from datetime import datetime
 from itertools import chain
-from typing import Dict, List
 
 from maggma.builders import GroupBuilder
 from maggma.core import Store
@@ -27,7 +26,7 @@ class XASBuilder(GroupBuilder):
         super().__init__(source=tasks, target=xas, grouping_keys=["mp_id"])
         self._target_keys_field = "xas_ids"
 
-    def process_item(self, spectra: List[Dict]) -> Dict:
+    def process_item(self, spectra: list[dict]) -> dict:
         # TODO: Change this to do structure matching against materials collection
         mpid = spectra[0]["mp_id"]
 

emmet/builders/materials/absorption_spectrum.py

@@ -1,5 +1,5 @@
+from __future__ import annotations
 from math import ceil
-from typing import Dict, Iterator, List, Optional
 
 import numpy as np
 from maggma.builders import Builder
@@ -10,6 +10,11 @@ from pymatgen.core.structure import Structure
 from emmet.core.absorption import AbsorptionDoc
 from emmet.core.utils import jsanitize
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 
 class AbsorptionBuilder(Builder):
     def __init__(
@@ -17,7 +22,7 @@ class AbsorptionBuilder(Builder):
         materials: Store,
         tasks: Store,
         absorption: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs,
     ):
         self.materials = materials
@@ -32,7 +37,7 @@ class AbsorptionBuilder(Builder):
 
         super().__init__(sources=[materials, tasks], targets=[absorption], **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
         """
@@ -44,7 +49,7 @@ class AbsorptionBuilder(Builder):
         for split in grouper(keys, N):
             yield {"query": {self.materials.key: {"$in": list(split)}}}
 
-    def get_items(self) -> Iterator[List[Dict]]:
+    def get_items(self) -> Iterator[list[dict]]:
         """
         Gets all items to process
 
@@ -139,23 +144,25 @@ class AbsorptionBuilder(Builder):
             task_query = self.tasks.query_one(
                 properties=[
                     "orig_inputs.kpoints",
-                    "orig_inputs.poscar.structure",
+                    "orig_inputs.structure",
                     "input.parameters",
                     "input.structure",
                     "output.dielectric.energy",
                     "output.dielectric.real",
                     "output.dielectric.imag",
-                    "output.optical_absorption_coeff",
+                    "calcs_reversed",
                     "output.bandgap",
                 ],
                 criteria={self.tasks.key: task_id},
             )
 
-            if task_query["output"]["optical_absorption_coeff"] is not None:
+            if (cr := task_query.get("calcs_reversed", [])) and (
+                oac := cr[0]["output"]["optical_absorption_coeff"]
+            ):
                 try:
-                    structure = task_query["orig_inputs"]["poscar"]["structure"]
-                except KeyError:
                     structure = task_query["input"]["structure"]
+                except KeyError:
+                    structure = task_query["orig_inputs"]["structure"]
 
                 if (
                     task_query["orig_inputs"]["kpoints"]["generation_style"]
@@ -176,12 +183,16 @@ class AbsorptionBuilder(Builder):
                     {
                         "task_id": task_id,
                         "nkpoints": int(nkpoints),
-                        "energies": task_query["output"]["dielectric"]["energy"],
-                        "real_dielectric": task_query["output"]["dielectric"]["real"],
-                        "imag_dielectric": task_query["output"]["dielectric"]["imag"],
-                        "optical_absorption_coeff": task_query["output"][
-                            "optical_absorption_coeff"
+                        "energies": cr[0]["output"]["frequency_dependent_dielectric"][
+                            "energy"
                         ],
+                        "real_dielectric": cr[0]["output"][
+                            "frequency_dependent_dielectric"
+                        ]["real"],
+                        "imag_dielectric": cr[0]["output"][
+                            "frequency_dependent_dielectric"
+                        ]["imaginary"],
+                        "optical_absorption_coeff": oac,
                         "bandgap": task_query["output"]["bandgap"],
                         "structure": structure,
                         "updated_on": lu_dt,

emmet/builders/materials/alloys.py

@@ -1,5 +1,4 @@
 from itertools import chain, combinations
-from typing import Dict, List, Tuple, Union
 
 from maggma.builders import Builder
 from matminer.datasets import load_dataset
@@ -13,7 +12,7 @@ from pymatgen.analysis.alloys.core import (
 from pymatgen.core.structure import Structure
 from tqdm import tqdm
 
-from emmet.core.thermo import ThermoType
+from emmet.core.types.enums import ThermoType
 
 # rough sort of ANON_FORMULAS by "complexity"
 ANON_FORMULAS = sorted(KNOWN_ANON_FORMULAS, key=lambda af: len(af))
@@ -40,7 +39,7 @@ class AlloyPairBuilder(Builder):
         provenance,
         oxi_states,
         alloy_pairs,
-        thermo_type: Union[ThermoType, str] = ThermoType.GGA_GGA_U_R2SCAN,
+        thermo_type: ThermoType | str = ThermoType.GGA_GGA_U_R2SCAN,
     ):
         self.materials = materials
         self.thermo = thermo
@@ -274,7 +273,7 @@ class AlloyPairMemberBuilder(Builder):
             if structures:
                 yield (pairs, structures)
 
-    def process_item(self, item: Tuple[List[AlloyPair], Dict[str, Structure]]):
+    def process_item(self, item: tuple[list[AlloyPair], dict[str, Structure]]):
         pairs, structures = item
 
         all_pair_members = []

emmet/builders/materials/chemenv.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.chemenv import ChemEnvDoc
 from emmet.core.utils import jsanitize
 
@@ -12,7 +11,7 @@ class ChemEnvBuilder(MapBuilder):
         self,
         oxidation_states: Store,
         chemenv: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs
     ):
         self.oxidation_states = oxidation_states

emmet/builders/materials/corrected_entries.py

@@ -1,10 +1,11 @@
+from __future__ import annotations
+
 import copy
 import warnings
 from collections import defaultdict
 from datetime import datetime
 from itertools import chain
 from math import ceil
-from typing import Dict, Iterable, Iterator, List, Optional, Union
 
 from maggma.core import Builder, Store
 from maggma.utils import grouper
@@ -13,18 +14,23 @@ from pymatgen.entries.computed_entries import ComputedStructureEntry
 
 from emmet.builders.utils import HiddenPrints, chemsys_permutations
 from emmet.core.corrected_entries import CorrectedEntriesDoc
-from emmet.core.thermo import ThermoType
+from emmet.core.types.enums import ThermoType
 from emmet.core.utils import jsanitize
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator
+
 
 class CorrectedEntriesBuilder(Builder):
     def __init__(
         self,
         materials: Store,
         corrected_entries: Store,
-        oxidation_states: Optional[Store] = None,
-        query: Optional[Dict] = None,
-        compatibility: Optional[Union[List[Compatibility], List[None]]] = [None],
+        oxidation_states: Store | None = None,
+        query: dict | None = None,
+        compatibility: list[Compatibility] | list[None] | None = [None],
         chunk_size: int = 1000,
         **kwargs,
     ):
@@ -48,7 +54,7 @@ class CorrectedEntriesBuilder(Builder):
         self.compatibility = compatibility
         self.oxidation_states = oxidation_states
         self.chunk_size = chunk_size
-        self._entries_cache: Dict[str, List[Dict]] = defaultdict(list)
+        self._entries_cache: dict[str, list[dict]] = defaultdict(list)
 
         if self.corrected_entries.key != "chemsys":
             warnings.warn(
@@ -93,7 +99,7 @@ class CorrectedEntriesBuilder(Builder):
         # Search index for corrected_entries
         self.corrected_entries.ensure_index("chemsys")
 
-    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
         to_process_chemsys = self._get_chemsys_to_process()
 
         N = ceil(len(to_process_chemsys) / number_splits)
@@ -101,7 +107,7 @@ class CorrectedEntriesBuilder(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
         """
@@ -210,7 +216,7 @@ class CorrectedEntriesBuilder(Builder):
         else:
             self.logger.info("No corrected entry items to update")
 
-    def get_entries(self, chemsys: str) -> List[Dict]:
+    def get_entries(self, chemsys: str) -> list[dict]:
         """
         Gets entries from the materials collection for the corresponding chemical systems
         Args:

emmet/builders/materials/dielectric.py

@@ -1,6 +1,6 @@
-from math import ceil
-from typing import Dict, Optional, Iterator
+from __future__ import annotations
 
+from math import ceil
 
 import numpy as np
 from maggma.builders import Builder
@@ -11,6 +11,11 @@ from pymatgen.core.structure import Structure
 from emmet.core.polar import DielectricDoc
 from emmet.core.utils import jsanitize
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 
 class DielectricBuilder(Builder):
     def __init__(
@@ -18,7 +23,7 @@ class DielectricBuilder(Builder):
         materials: Store,
         tasks: Store,
         dielectric: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs,
     ):
         self.materials = materials
@@ -33,7 +38,7 @@ class DielectricBuilder(Builder):
 
         super().__init__(sources=[materials, tasks], targets=[dielectric], **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
         """
@@ -143,11 +148,10 @@ class DielectricBuilder(Builder):
                     "last_updated",
                     "input.is_hubbard",
                     "orig_inputs.kpoints",
-                    "orig_inputs.poscar.structure",
+                    "orig_inputs.structure",
                     "input.parameters",
                     "input.structure",
-                    "output.epsilon_static",
-                    "output.epsilon_ionic",
+                    "calcs_reversed",
                     "output.bandgap",
                 ],
                 criteria={self.tasks.key: str(task_id)},
@@ -155,9 +159,9 @@ class DielectricBuilder(Builder):
 
             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"]
 
@@ -182,8 +186,12 @@ class DielectricBuilder(Builder):
                         "task_id": task_id,
                         "is_hubbard": int(is_hubbard),
                         "nkpoints": int(nkpoints),
-                        "epsilon_static": task_query["output"]["epsilon_static"],
-                        "epsilon_ionic": task_query["output"]["epsilon_ionic"],
+                        "epsilon_static": task_query["calcs_reversed"][0]["output"][
+                            "epsilon_static"
+                        ],
+                        "epsilon_ionic": task_query["calcs_reversed"][0]["output"][
+                            "epsilon_ionic"
+                        ],
                         "structure": structure,
                         "updated_on": lu_dt,
                         "task_updated": task_updated,