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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (44) hide show
  1. emmet/builders/abinit/phonon.py +47 -47
  2. emmet/builders/abinit/sound_velocity.py +15 -11
  3. emmet/builders/feff/xas.py +1 -2
  4. emmet/builders/materials/absorption_spectrum.py +25 -14
  5. emmet/builders/materials/alloys.py +10 -11
  6. emmet/builders/materials/chemenv.py +2 -3
  7. emmet/builders/materials/corrected_entries.py +21 -15
  8. emmet/builders/materials/dielectric.py +19 -11
  9. emmet/builders/materials/elasticity.py +44 -33
  10. emmet/builders/materials/electrodes.py +35 -28
  11. emmet/builders/materials/electronic_structure.py +17 -17
  12. emmet/builders/materials/magnetism.py +11 -4
  13. emmet/builders/materials/optimade.py +7 -3
  14. emmet/builders/materials/piezoelectric.py +24 -21
  15. emmet/builders/materials/provenance.py +16 -13
  16. emmet/builders/materials/robocrys.py +2 -3
  17. emmet/builders/materials/substrates.py +9 -8
  18. emmet/builders/materials/summary.py +3 -3
  19. emmet/builders/materials/thermo.py +17 -11
  20. emmet/builders/matscholar/missing_compositions.py +12 -8
  21. emmet/builders/mobility/migration_graph.py +5 -5
  22. emmet/builders/settings.py +21 -17
  23. emmet/builders/utils.py +101 -12
  24. emmet/builders/vasp/materials.py +40 -51
  25. emmet/builders/vasp/mp_potcar_stats.json.gz +0 -0
  26. emmet/builders/vasp/task_validator.py +25 -36
  27. emmet_builders-0.86.0.dist-info/METADATA +37 -0
  28. emmet_builders-0.86.0.dist-info/RECORD +41 -0
  29. {emmet_builders-0.78.3.dist-info → emmet_builders-0.86.0.dist-info}/WHEEL +1 -1
  30. emmet/builders/materials/ml.py +0 -87
  31. emmet/builders/molecules/atomic.py +0 -589
  32. emmet/builders/molecules/bonds.py +0 -324
  33. emmet/builders/molecules/metal_binding.py +0 -526
  34. emmet/builders/molecules/orbitals.py +0 -288
  35. emmet/builders/molecules/redox.py +0 -496
  36. emmet/builders/molecules/summary.py +0 -383
  37. emmet/builders/molecules/thermo.py +0 -500
  38. emmet/builders/molecules/vibration.py +0 -278
  39. emmet/builders/qchem/__init__.py +0 -0
  40. emmet/builders/qchem/molecules.py +0 -734
  41. emmet_builders-0.78.3.dist-info/METADATA +0 -47
  42. emmet_builders-0.78.3.dist-info/RECORD +0 -51
  43. /emmet/builders/{molecules/__init__.py → py.typed} +0 -0
  44. {emmet_builders-0.78.3.dist-info → emmet_builders-0.86.0.dist-info}/top_level.txt +0 -0
@@ -1,43 +1,45 @@
1
- import tempfile
2
1
  import os
2
+ import tempfile
3
3
  from math import ceil
4
- from emmet.builders.settings import EmmetBuildSettings
4
+
5
5
  import numpy as np
6
- from typing import Optional, Dict, List, Iterator, Tuple
6
+ from abipy.abio.inputs import AnaddbInput
7
+ from abipy.core.abinit_units import eV_to_THz
8
+ from abipy.dfpt.anaddbnc import AnaddbNcFile
9
+ from abipy.dfpt.ddb import AnaddbError, DdbFile, DielectricTensorGenerator
10
+ from abipy.dfpt.phonons import PhononBands
11
+ from abipy.flowtk.tasks import AnaddbTask, TaskManager
12
+ from maggma.builders import Builder
13
+ from maggma.core import Store
7
14
  from maggma.utils import grouper
8
-
15
+ from pymatgen.core.structure import Structure
16
+ from pymatgen.io.abinit.abiobjects import KSampling
9
17
  from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine
10
18
  from pymatgen.phonon.dos import CompletePhononDos
11
19
  from pymatgen.phonon.ir_spectra import IRDielectricTensor
12
- from pymatgen.core.structure import Structure
13
- from pymatgen.io.abinit.abiobjects import KSampling
14
- from pymatgen.symmetry.bandstructure import HighSymmKpath
15
20
  from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
16
- from abipy.dfpt.anaddbnc import AnaddbNcFile
17
- from abipy.abio.inputs import AnaddbInput
18
- from abipy.flowtk.tasks import AnaddbTask, TaskManager
19
- from abipy.dfpt.ddb import AnaddbError, DielectricTensorGenerator, DdbFile
20
- from abipy.dfpt.phonons import PhononBands
21
- from abipy.core.abinit_units import eV_to_THz
22
- from maggma.builders import Builder
23
- from maggma.core import Store
21
+ from pymatgen.symmetry.bandstructure import HighSymmKpath
24
22
 
23
+ from emmet.builders.settings import EmmetBuildSettings
25
24
  from emmet.core.phonon import (
26
- PhononWarnings,
27
- ThermodynamicProperties,
28
25
  AbinitPhonon,
29
- VibrationalEnergy,
30
- )
31
- from emmet.core.phonon import (
32
- PhononDos,
33
- PhononBandStructure,
34
- PhononWebsiteBS,
35
26
  Ddb,
27
+ PhononBS,
28
+ PhononDOS,
29
+ PhononWarnings,
30
+ PhononWebsiteBS,
36
31
  ThermalDisplacement,
32
+ ThermodynamicProperties,
33
+ VibrationalEnergy,
37
34
  )
38
- from emmet.core.polar import DielectricDoc, BornEffectiveCharges, IRDielectric
35
+ from emmet.core.polar import BornEffectiveCharges, DielectricDoc, IRDielectric
39
36
  from emmet.core.utils import jsanitize
40
37
 
38
+ from typing import TYPE_CHECKING
39
+
40
+ if TYPE_CHECKING:
41
+ from collections.abc import Iterator
42
+
41
43
  SETTINGS = EmmetBuildSettings()
42
44
 
43
45
 
@@ -52,8 +54,8 @@ class PhononBuilder(Builder):
52
54
  ddb_files: Store,
53
55
  th_disp: Store,
54
56
  phonon_website: Store,
55
- query: Optional[Dict] = None,
56
- manager: Optional[TaskManager] = None,
57
+ query: dict | None = None,
58
+ manager: TaskManager | None = None,
57
59
  symprec: float = SETTINGS.SYMPREC,
58
60
  angle_tolerance: float = SETTINGS.ANGLE_TOL,
59
61
  chunk_size=100,
@@ -136,7 +138,7 @@ class PhononBuilder(Builder):
136
138
  for mpid_chunk in grouper(mats, N):
137
139
  yield {"query": {self.phonon_materials.key: {"$in": list(mpid_chunk)}}}
138
140
 
139
- def get_items(self) -> Iterator[Dict]:
141
+ def get_items(self) -> Iterator[dict]:
140
142
  """
141
143
  Gets all materials that need phonons
142
144
 
@@ -190,7 +192,7 @@ class PhononBuilder(Builder):
190
192
 
191
193
  yield item
192
194
 
193
- def process_item(self, item: Dict) -> Optional[Dict]:
195
+ def process_item(self, item: dict) -> dict | None:
194
196
  """
195
197
  Generates the full phonon document from an item
196
198
 
@@ -221,7 +223,6 @@ class PhononBuilder(Builder):
221
223
  becs = None
222
224
  if phonon_properties["becs"] is not None:
223
225
  becs = BornEffectiveCharges(
224
- material_id=item["mp_id"],
225
226
  symmetrized_value=phonon_properties["becs"],
226
227
  value=sr_break["becs_nosymm"],
227
228
  cnsr_break=sr_break["cnsr"],
@@ -243,25 +244,24 @@ class PhononBuilder(Builder):
243
244
  abinit_input_vars=abinit_input_vars,
244
245
  )
245
246
 
246
- phbs = PhononBandStructure(
247
- material_id=item["mp_id"],
248
- band_structure=phonon_properties["ph_bs"].as_dict(),
247
+ phbs = PhononBS(
248
+ identifier=item["mp_id"],
249
+ **phonon_properties["ph_bs"].as_dict(),
249
250
  )
250
251
 
251
- phws = PhononWebsiteBS(
252
+ phws = PhononWebsiteBS( # type: ignore[call-arg]
252
253
  material_id=item["mp_id"],
253
254
  phononwebsite=phonon_properties["ph_bs"].as_phononwebsite(),
254
255
  )
255
256
 
256
- phdos = PhononDos(
257
- material_id=item["mp_id"],
258
- dos=phonon_properties["ph_dos"].as_dict(),
259
- dos_method=phonon_properties["ph_dos_method"],
257
+ phdos = PhononDOS(
258
+ identifier=item["mp_id"],
259
+ **phonon_properties["ph_dos"].as_dict(),
260
260
  )
261
261
 
262
- ddb = Ddb(material_id=item["mp_id"], ddb=item["ddb_str"])
262
+ ddb = Ddb(material_id=item["mp_id"], ddb=item["ddb_str"]) # type: ignore[call-arg]
263
263
 
264
- th_disp = ThermalDisplacement(
264
+ th_disp = ThermalDisplacement( # type: ignore[call-arg]
265
265
  material_id=item["mp_id"],
266
266
  structure=structure,
267
267
  nsites=len(structure),
@@ -295,7 +295,7 @@ class PhononBuilder(Builder):
295
295
  )
296
296
  return None
297
297
 
298
- def get_phonon_properties(self, item: Dict) -> Dict:
298
+ def get_phonon_properties(self, item: dict) -> dict:
299
299
  """
300
300
  Extracts the phonon properties from the item
301
301
  """
@@ -513,7 +513,7 @@ class PhononBuilder(Builder):
513
513
  dos: str = "tetra",
514
514
  lo_to_splitting: bool = True,
515
515
  use_dieflag: bool = True,
516
- ) -> Tuple[AnaddbInput, Optional[List]]:
516
+ ) -> tuple[AnaddbInput, list | None]:
517
517
  """
518
518
  creates the AnaddbInput object to calculate the phonon properties.
519
519
  It also returns the list of qpoints labels for generating the PhononBandStructureSymmLine.
@@ -639,7 +639,7 @@ class PhononBuilder(Builder):
639
639
 
640
640
  @staticmethod
641
641
  def get_pmg_bs(
642
- phbands: PhononBands, labels_list: List
642
+ phbands: PhononBands, labels_list: list
643
643
  ) -> PhononBandStructureSymmLine:
644
644
  """
645
645
  Generates a PhononBandStructureSymmLine starting from a abipy PhononBands object
@@ -679,7 +679,7 @@ class PhononBuilder(Builder):
679
679
  )
680
680
 
681
681
  ph_bs_sl = PhononBandStructureSymmLine(
682
- qpoints=qpts,
682
+ qpoints=qpts, # type: ignore[arg-type]
683
683
  frequencies=ph_freqs,
684
684
  lattice=structure.reciprocal_lattice,
685
685
  has_nac=phbands.non_anal_ph is not None,
@@ -728,7 +728,7 @@ class PhononBuilder(Builder):
728
728
 
729
729
  return data
730
730
 
731
- def update_targets(self, items: List[Dict]):
731
+ def update_targets(self, items: list[dict]):
732
732
  """
733
733
  Inserts the new task_types into the task_types collection
734
734
 
@@ -771,7 +771,7 @@ class PhononBuilder(Builder):
771
771
 
772
772
  def get_warnings(
773
773
  asr_break: float, cnsr_break: float, ph_bs: PhononBandStructureSymmLine
774
- ) -> List[PhononWarnings]:
774
+ ) -> list[PhononWarnings]:
775
775
  """
776
776
 
777
777
  Args:
@@ -812,7 +812,7 @@ def get_warnings(
812
812
 
813
813
  def get_thermodynamic_properties(
814
814
  ph_dos: CompletePhononDos,
815
- ) -> Tuple[ThermodynamicProperties, VibrationalEnergy]:
815
+ ) -> tuple[ThermodynamicProperties, VibrationalEnergy]:
816
816
  """
817
817
  Calculates the thermodynamic properties from a phonon DOS
818
818
 
@@ -841,7 +841,7 @@ def get_thermodynamic_properties(
841
841
  zpe = ph_dos.zero_point_energy(ph_dos.structure)
842
842
 
843
843
  temperatures = temp.tolist()
844
- tp = ThermodynamicProperties(temperatures=temperatures, cv=cv, entropy=entropy)
844
+ tp = ThermodynamicProperties(temperatures=temperatures, cv=cv, entropy=entropy) # type: ignore[call-arg]
845
845
 
846
846
  ve = VibrationalEnergy(
847
847
  temperatures=temperatures,
@@ -1,18 +1,22 @@
1
1
  import tempfile
2
2
  import traceback
3
3
  from math import ceil
4
- from maggma.utils import grouper
5
- from typing import Optional, Dict, List, Iterator
6
4
 
7
- from abipy.dfpt.vsound import SoundVelocity as AbiSoundVelocity
8
5
  from abipy.dfpt.ddb import DdbFile
6
+ from abipy.dfpt.vsound import SoundVelocity as AbiSoundVelocity
7
+ from abipy.flowtk.tasks import TaskManager
9
8
  from maggma.builders import Builder
10
9
  from maggma.core import Store
11
- from abipy.flowtk.tasks import TaskManager
10
+ from maggma.utils import grouper
12
11
 
13
12
  from emmet.core.phonon import SoundVelocity
14
13
  from emmet.core.utils import jsanitize
15
14
 
15
+ from typing import TYPE_CHECKING
16
+
17
+ if TYPE_CHECKING:
18
+ from collections.abc import Iterator
19
+
16
20
 
17
21
  class SoundVelocityBuilder(Builder):
18
22
  def __init__(
@@ -20,8 +24,8 @@ class SoundVelocityBuilder(Builder):
20
24
  phonon_materials: Store,
21
25
  ddb_source: Store,
22
26
  sound_vel: Store,
23
- query: Optional[dict] = None,
24
- manager: Optional[TaskManager] = None,
27
+ query: dict | None = None,
28
+ manager: TaskManager | None = None,
25
29
  **kwargs
26
30
  ):
27
31
  """
@@ -72,7 +76,7 @@ class SoundVelocityBuilder(Builder):
72
76
  for mpid_chunk in grouper(mats, N):
73
77
  yield {"query": {self.phonon_materials.key: {"$in": list(mpid_chunk)}}}
74
78
 
75
- def get_items(self) -> Iterator[Dict]:
79
+ def get_items(self) -> Iterator[dict]:
76
80
  """
77
81
  Gets all materials that need sound velocity.
78
82
 
@@ -117,7 +121,7 @@ class SoundVelocityBuilder(Builder):
117
121
 
118
122
  yield item
119
123
 
120
- def process_item(self, item: Dict) -> Optional[Dict]:
124
+ def process_item(self, item: dict) -> dict | None:
121
125
  """
122
126
  Generates the sound velocity document from an item
123
127
 
@@ -132,7 +136,7 @@ class SoundVelocityBuilder(Builder):
132
136
  try:
133
137
  sound_vel_data = self.get_sound_vel(item)
134
138
 
135
- sv = SoundVelocity(
139
+ sv = SoundVelocity( # type: ignore[call-arg]
136
140
  material_id=item["mp_id"],
137
141
  structure=sound_vel_data["structure"],
138
142
  directions=sound_vel_data["directions"],
@@ -153,7 +157,7 @@ class SoundVelocityBuilder(Builder):
153
157
  return None
154
158
 
155
159
  @staticmethod
156
- def get_sound_vel(item: Dict) -> Dict:
160
+ def get_sound_vel(item: dict) -> dict:
157
161
  """
158
162
  Runs anaddb and return the extracted data for the speed of sound.
159
163
 
@@ -187,7 +191,7 @@ class SoundVelocityBuilder(Builder):
187
191
 
188
192
  return sv_data
189
193
 
190
- def update_targets(self, items: List[Dict]):
194
+ def update_targets(self, items: list[dict]):
191
195
  """
192
196
  Inserts the new task_types into the task_types collection
193
197
 
@@ -1,7 +1,6 @@
1
1
  import traceback
2
2
  from datetime import datetime
3
3
  from itertools import chain
4
- from typing import Dict, List
5
4
 
6
5
  from maggma.builders import GroupBuilder
7
6
  from maggma.core import Store
@@ -27,7 +26,7 @@ class XASBuilder(GroupBuilder):
27
26
  super().__init__(source=tasks, target=xas, grouping_keys=["mp_id"])
28
27
  self._target_keys_field = "xas_ids"
29
28
 
30
- def process_item(self, spectra: List[Dict]) -> Dict:
29
+ def process_item(self, spectra: list[dict]) -> dict:
31
30
  # TODO: Change this to do structure matching against materials collection
32
31
  mpid = spectra[0]["mp_id"]
33
32
 
@@ -1,5 +1,5 @@
1
+ from __future__ import annotations
1
2
  from math import ceil
2
- from typing import Dict, Iterator, List, Optional
3
3
 
4
4
  import numpy as np
5
5
  from maggma.builders import Builder
@@ -10,6 +10,11 @@ from pymatgen.core.structure import Structure
10
10
  from emmet.core.absorption import AbsorptionDoc
11
11
  from emmet.core.utils import jsanitize
12
12
 
13
+ from typing import TYPE_CHECKING
14
+
15
+ if TYPE_CHECKING:
16
+ from collections.abc import Iterator
17
+
13
18
 
14
19
  class AbsorptionBuilder(Builder):
15
20
  def __init__(
@@ -17,7 +22,7 @@ class AbsorptionBuilder(Builder):
17
22
  materials: Store,
18
23
  tasks: Store,
19
24
  absorption: Store,
20
- query: Optional[Dict] = None,
25
+ query: dict | None = None,
21
26
  **kwargs,
22
27
  ):
23
28
  self.materials = materials
@@ -32,7 +37,7 @@ class AbsorptionBuilder(Builder):
32
37
 
33
38
  super().__init__(sources=[materials, tasks], targets=[absorption], **kwargs)
34
39
 
35
- def prechunk(self, number_splits: int) -> Iterator[Dict]: # pragma: no cover
40
+ def prechunk(self, number_splits: int) -> Iterator[dict]: # pragma: no cover
36
41
  """
37
42
  Prechunk method to perform chunking by the key field
38
43
  """
@@ -44,7 +49,7 @@ class AbsorptionBuilder(Builder):
44
49
  for split in grouper(keys, N):
45
50
  yield {"query": {self.materials.key: {"$in": list(split)}}}
46
51
 
47
- def get_items(self) -> Iterator[List[Dict]]:
52
+ def get_items(self) -> Iterator[list[dict]]:
48
53
  """
49
54
  Gets all items to process
50
55
 
@@ -139,23 +144,25 @@ class AbsorptionBuilder(Builder):
139
144
  task_query = self.tasks.query_one(
140
145
  properties=[
141
146
  "orig_inputs.kpoints",
142
- "orig_inputs.poscar.structure",
147
+ "orig_inputs.structure",
143
148
  "input.parameters",
144
149
  "input.structure",
145
150
  "output.dielectric.energy",
146
151
  "output.dielectric.real",
147
152
  "output.dielectric.imag",
148
- "output.optical_absorption_coeff",
153
+ "calcs_reversed",
149
154
  "output.bandgap",
150
155
  ],
151
156
  criteria={self.tasks.key: task_id},
152
157
  )
153
158
 
154
- if task_query["output"]["optical_absorption_coeff"] is not None:
159
+ if (cr := task_query.get("calcs_reversed", [])) and (
160
+ oac := cr[0]["output"]["optical_absorption_coeff"]
161
+ ):
155
162
  try:
156
- structure = task_query["orig_inputs"]["poscar"]["structure"]
157
- except KeyError:
158
163
  structure = task_query["input"]["structure"]
164
+ except KeyError:
165
+ structure = task_query["orig_inputs"]["structure"]
159
166
 
160
167
  if (
161
168
  task_query["orig_inputs"]["kpoints"]["generation_style"]
@@ -176,12 +183,16 @@ class AbsorptionBuilder(Builder):
176
183
  {
177
184
  "task_id": task_id,
178
185
  "nkpoints": int(nkpoints),
179
- "energies": task_query["output"]["dielectric"]["energy"],
180
- "real_dielectric": task_query["output"]["dielectric"]["real"],
181
- "imag_dielectric": task_query["output"]["dielectric"]["imag"],
182
- "optical_absorption_coeff": task_query["output"][
183
- "optical_absorption_coeff"
186
+ "energies": cr[0]["output"]["frequency_dependent_dielectric"][
187
+ "energy"
184
188
  ],
189
+ "real_dielectric": cr[0]["output"][
190
+ "frequency_dependent_dielectric"
191
+ ]["real"],
192
+ "imag_dielectric": cr[0]["output"][
193
+ "frequency_dependent_dielectric"
194
+ ]["imaginary"],
195
+ "optical_absorption_coeff": oac,
185
196
  "bandgap": task_query["output"]["bandgap"],
186
197
  "structure": structure,
187
198
  "updated_on": lu_dt,
@@ -1,19 +1,18 @@
1
- from itertools import combinations, chain
2
- from typing import Tuple, List, Dict, Union
1
+ from itertools import chain, combinations
3
2
 
4
- from tqdm import tqdm
5
3
  from maggma.builders import Builder
6
- from pymatgen.core.structure import Structure
7
4
  from matminer.datasets import load_dataset
8
- from emmet.core.thermo import ThermoType
9
-
10
5
  from pymatgen.analysis.alloys.core import (
11
- AlloyPair,
12
- InvalidAlloy,
13
6
  KNOWN_ANON_FORMULAS,
14
7
  AlloyMember,
8
+ AlloyPair,
15
9
  AlloySystem,
10
+ InvalidAlloy,
16
11
  )
12
+ from pymatgen.core.structure import Structure
13
+ from tqdm import tqdm
14
+
15
+ from emmet.core.types.enums import ThermoType
17
16
 
18
17
  # rough sort of ANON_FORMULAS by "complexity"
19
18
  ANON_FORMULAS = sorted(KNOWN_ANON_FORMULAS, key=lambda af: len(af))
@@ -40,7 +39,7 @@ class AlloyPairBuilder(Builder):
40
39
  provenance,
41
40
  oxi_states,
42
41
  alloy_pairs,
43
- thermo_type: Union[ThermoType, str] = ThermoType.GGA_GGA_U_R2SCAN,
42
+ thermo_type: ThermoType | str = ThermoType.GGA_GGA_U_R2SCAN,
44
43
  ):
45
44
  self.materials = materials
46
45
  self.thermo = thermo
@@ -274,7 +273,7 @@ class AlloyPairMemberBuilder(Builder):
274
273
  if structures:
275
274
  yield (pairs, structures)
276
275
 
277
- def process_item(self, item: Tuple[List[AlloyPair], Dict[str, Structure]]):
276
+ def process_item(self, item: tuple[list[AlloyPair], dict[str, Structure]]):
278
277
  pairs, structures = item
279
278
 
280
279
  all_pair_members = []
@@ -291,7 +290,7 @@ class AlloyPairMemberBuilder(Builder):
291
290
  is_ordered=structure.is_ordered,
292
291
  x=pair.get_x(structure.composition),
293
292
  )
294
- pair_members["members"].append(member.as_dict())
293
+ pair_members["members"].append(member.as_dict()) # type: ignore[attr-defined]
295
294
  except Exception as exc:
296
295
  print(f"Exception for {db_id}: {exc}")
297
296
  if pair_members["members"]:
@@ -1,8 +1,7 @@
1
- from typing import Dict, Optional
2
1
  from maggma.builders.map_builder import MapBuilder
3
2
  from maggma.core import Store
4
-
5
3
  from pymatgen.core.structure import Structure
4
+
6
5
  from emmet.core.chemenv import ChemEnvDoc
7
6
  from emmet.core.utils import jsanitize
8
7
 
@@ -12,7 +11,7 @@ class ChemEnvBuilder(MapBuilder):
12
11
  self,
13
12
  oxidation_states: Store,
14
13
  chemenv: Store,
15
- query: Optional[Dict] = None,
14
+ query: dict | None = None,
16
15
  **kwargs
17
16
  ):
18
17
  self.oxidation_states = oxidation_states
@@ -1,20 +1,26 @@
1
+ from __future__ import annotations
2
+
3
+ import copy
1
4
  import warnings
2
5
  from collections import defaultdict
6
+ from datetime import datetime
3
7
  from itertools import chain
4
- from typing import Dict, Iterable, Iterator, List, Optional
5
8
  from math import ceil
6
- import copy
7
- from datetime import datetime
8
9
 
9
10
  from maggma.core import Builder, Store
10
11
  from maggma.utils import grouper
11
- from pymatgen.entries.computed_entries import ComputedStructureEntry
12
12
  from pymatgen.entries.compatibility import Compatibility
13
+ from pymatgen.entries.computed_entries import ComputedStructureEntry
13
14
 
14
- from emmet.core.utils import jsanitize
15
- from emmet.builders.utils import chemsys_permutations, HiddenPrints
16
- from emmet.core.thermo import ThermoType
15
+ from emmet.builders.utils import HiddenPrints, chemsys_permutations
17
16
  from emmet.core.corrected_entries import CorrectedEntriesDoc
17
+ from emmet.core.types.enums import ThermoType
18
+ from emmet.core.utils import jsanitize
19
+
20
+ from typing import TYPE_CHECKING
21
+
22
+ if TYPE_CHECKING:
23
+ from collections.abc import Iterable, Iterator
18
24
 
19
25
 
20
26
  class CorrectedEntriesBuilder(Builder):
@@ -22,9 +28,9 @@ class CorrectedEntriesBuilder(Builder):
22
28
  self,
23
29
  materials: Store,
24
30
  corrected_entries: Store,
25
- oxidation_states: Optional[Store] = None,
26
- query: Optional[Dict] = None,
27
- compatibility: Optional[List[Compatibility]] = None,
31
+ oxidation_states: Store | None = None,
32
+ query: dict | None = None,
33
+ compatibility: list[Compatibility] | list[None] | None = [None],
28
34
  chunk_size: int = 1000,
29
35
  **kwargs,
30
36
  ):
@@ -45,10 +51,10 @@ class CorrectedEntriesBuilder(Builder):
45
51
  self.materials = materials
46
52
  self.query = query if query else {}
47
53
  self.corrected_entries = corrected_entries
48
- self.compatibility = compatibility or [None]
54
+ self.compatibility = compatibility
49
55
  self.oxidation_states = oxidation_states
50
56
  self.chunk_size = chunk_size
51
- self._entries_cache: Dict[str, List[ComputedStructureEntry]] = defaultdict(list)
57
+ self._entries_cache: dict[str, list[dict]] = defaultdict(list)
52
58
 
53
59
  if self.corrected_entries.key != "chemsys":
54
60
  warnings.warn(
@@ -93,7 +99,7 @@ class CorrectedEntriesBuilder(Builder):
93
99
  # Search index for corrected_entries
94
100
  self.corrected_entries.ensure_index("chemsys")
95
101
 
96
- def prechunk(self, number_splits: int) -> Iterable[Dict]: # pragma: no cover
102
+ def prechunk(self, number_splits: int) -> Iterable[dict]: # pragma: no cover
97
103
  to_process_chemsys = self._get_chemsys_to_process()
98
104
 
99
105
  N = ceil(len(to_process_chemsys) / number_splits)
@@ -101,7 +107,7 @@ class CorrectedEntriesBuilder(Builder):
101
107
  for chemsys_chunk in grouper(to_process_chemsys, N):
102
108
  yield {"query": {"chemsys": {"$in": list(chemsys_chunk)}}}
103
109
 
104
- def get_items(self) -> Iterator[List[Dict]]:
110
+ def get_items(self) -> Iterator[list[dict]]:
105
111
  """
106
112
  Gets whole chemical systems of entries to process
107
113
  """
@@ -210,7 +216,7 @@ class CorrectedEntriesBuilder(Builder):
210
216
  else:
211
217
  self.logger.info("No corrected entry items to update")
212
218
 
213
- def get_entries(self, chemsys: str) -> List[Dict]:
219
+ def get_entries(self, chemsys: str) -> list[dict]:
214
220
  """
215
221
  Gets entries from the materials collection for the corresponding chemical systems
216
222
  Args:
@@ -1,6 +1,6 @@
1
- from math import ceil
2
- from typing import Dict, Optional, Iterator
1
+ from __future__ import annotations
3
2
 
3
+ from math import ceil
4
4
 
5
5
  import numpy as np
6
6
  from maggma.builders import Builder
@@ -11,6 +11,11 @@ from pymatgen.core.structure import Structure
11
11
  from emmet.core.polar import DielectricDoc
12
12
  from emmet.core.utils import jsanitize
13
13
 
14
+ from typing import TYPE_CHECKING
15
+
16
+ if TYPE_CHECKING:
17
+ from collections.abc import Iterator
18
+
14
19
 
15
20
  class DielectricBuilder(Builder):
16
21
  def __init__(
@@ -18,7 +23,7 @@ class DielectricBuilder(Builder):
18
23
  materials: Store,
19
24
  tasks: Store,
20
25
  dielectric: Store,
21
- query: Optional[Dict] = None,
26
+ query: dict | None = None,
22
27
  **kwargs,
23
28
  ):
24
29
  self.materials = materials
@@ -33,7 +38,7 @@ class DielectricBuilder(Builder):
33
38
 
34
39
  super().__init__(sources=[materials, tasks], targets=[dielectric], **kwargs)
35
40
 
36
- def prechunk(self, number_splits: int) -> Iterator[Dict]: # pragma: no cover
41
+ def prechunk(self, number_splits: int) -> Iterator[dict]: # pragma: no cover
37
42
  """
38
43
  Prechunk method to perform chunking by the key field
39
44
  """
@@ -143,11 +148,10 @@ class DielectricBuilder(Builder):
143
148
  "last_updated",
144
149
  "input.is_hubbard",
145
150
  "orig_inputs.kpoints",
146
- "orig_inputs.poscar.structure",
151
+ "orig_inputs.structure",
147
152
  "input.parameters",
148
153
  "input.structure",
149
- "output.epsilon_static",
150
- "output.epsilon_ionic",
154
+ "calcs_reversed",
151
155
  "output.bandgap",
152
156
  ],
153
157
  criteria={self.tasks.key: str(task_id)},
@@ -155,9 +159,9 @@ class DielectricBuilder(Builder):
155
159
 
156
160
  if task_query["output"]["bandgap"] > 0:
157
161
  try:
158
- structure = task_query["orig_inputs"]["poscar"]["structure"]
159
- except KeyError:
160
162
  structure = task_query["input"]["structure"]
163
+ except KeyError:
164
+ structure = task_query["orig_inputs"]["structure"]
161
165
 
162
166
  is_hubbard = task_query["input"]["is_hubbard"]
163
167
 
@@ -182,8 +186,12 @@ class DielectricBuilder(Builder):
182
186
  "task_id": task_id,
183
187
  "is_hubbard": int(is_hubbard),
184
188
  "nkpoints": int(nkpoints),
185
- "epsilon_static": task_query["output"]["epsilon_static"],
186
- "epsilon_ionic": task_query["output"]["epsilon_ionic"],
189
+ "epsilon_static": task_query["calcs_reversed"][0]["output"][
190
+ "epsilon_static"
191
+ ],
192
+ "epsilon_ionic": task_query["calcs_reversed"][0]["output"][
193
+ "epsilon_ionic"
194
+ ],
187
195
  "structure": structure,
188
196
  "updated_on": lu_dt,
189
197
  "task_updated": task_updated,