Skip to content

Commit 421b07c

Browse files
elinsc-botclaude
elinsc-bot
and
claude
committed
Gather the chain Codes TypedDicts into one module
The TypedDicts lived in their chain modules, so reading a chain's code requirements meant importing the chain — and the purpose strings named internal route names no input-file author would recognize. - Move every per-chain Codes TypedDict into aiida_koopmans.workgraphs.codes, which imports no chain code, so the dispatcher can introspect requirements cheaply. - Phrase each purpose string in input-file vocabulary: the keyword that triggers the need ("Needed for init_orbitals: mlwfs or projwfs", "Needed when block_wannierization_threshold is set", "Needed for eps_inf: auto"). - Skip `from __future__ import annotations` in the new module: with stringified annotations TypedDict.__required_keys__ reports every member as required (python/cpython#97727, still true on 3.13). - Test the introspection surface the dispatcher will use: __required_keys__/__optional_keys__ through NotRequired, and get_type_hints(include_extras=True) surfacing the SocketMeta helps. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent d919be9 commit 421b07c

12 files changed

Lines changed: 212 additions & 134 deletions

File tree

‎src/aiida_koopmans/workgraphs/auto_wannierize.py‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@
5757
WannierizeBlockOutputs,
5858
WannierizeOverrides,
5959
)
60+
from aiida_koopmans.workgraphs.codes import SplitBlockCodes
6061
from aiida_koopmans.workgraphs.pw import assemble_pw_base_step
6162
from aiida_koopmans.workgraphs.utils.wannier_merge import (
6263
merge_wannier_centres_file_contents,
@@ -340,15 +341,6 @@ def _subblock_w90_parameters(
340341
# ----------------------------------------------------------------------
341342

342343

343-
class SplitBlockCodes(TypedDict):
344-
"""Codes for the wannierize-and-split path (:func:`WannierizeAndSplitBlock`)."""
345-
346-
pw: orm.AbstractCode
347-
pw2wannier90: orm.AbstractCode
348-
wannier90: orm.AbstractCode
349-
wannierjl: orm.AbstractCode
350-
351-
352344
class RewannierizeSplitOutputs(TypedDict):
353345
"""Outputs of :func:`RewannierizeSplitBlocks`.
354346

‎src/aiida_koopmans/workgraphs/block_wannierize.py‎

Lines changed: 1 addition & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,6 @@
6161
)
6262
from aiida_wannier90_workflows.workflows import Wannier90WorkChain
6363
from aiida_workgraph import dynamic, task
64-
from aiida_workgraph.socket_spec import SocketMeta
6564
from aiida_workgraph.utils import get_dict_from_builder
6665

6766
from aiida_koopmans.parallelization import (
@@ -80,6 +79,7 @@
8079
from aiida_koopmans.spin import SpinChannel
8180
from aiida_koopmans.variational_orbitals import VariationalOrbital
8281
from aiida_koopmans.workgraphs import unwrap_enum
82+
from aiida_koopmans.workgraphs.codes import WannierizeBlockCodes, WannierizeBlocksCodes
8383
from aiida_koopmans.workgraphs.pw import PwOutputs, RunScfNscf
8484
from aiida_koopmans.workgraphs.variational_orbitals import (
8585
initial_orbital_partition,
@@ -107,33 +107,6 @@
107107
)
108108

109109

110-
class WannierizeBlockCodes(TypedDict):
111-
"""Codes for one block's wannierization (:func:`WannierizeBlock`).
112-
113-
``pw`` never runs here (the block skips scf / nscf) but the upstream
114-
``Wannier90WorkChain`` builder still requires it to assemble the pw
115-
namespaces this graph discards.
116-
"""
117-
118-
pw: orm.AbstractCode
119-
pw2wannier90: orm.AbstractCode
120-
wannier90: orm.AbstractCode
121-
122-
123-
class WannierizeBlocksCodes(TypedDict):
124-
"""Codes for the block-by-block wannierization chain (:func:`WannierizeBlocks`)."""
125-
126-
pw: orm.AbstractCode
127-
pw2wannier90: orm.AbstractCode
128-
wannier90: orm.AbstractCode
129-
wannierjl: NotRequired[
130-
Annotated[
131-
orm.AbstractCode,
132-
SocketMeta(help="Needed for block splitting via parallel transport."),
133-
]
134-
]
135-
136-
137110
def validate_projection_type(projection_type: WannierProjectionType) -> None:
138111
"""Reject projection types the block wannierization does not support.
139112

‎src/aiida_koopmans/workgraphs/codes.py‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
"""Per-chain code requirements, declared as TypedDicts.
2+
3+
One TypedDict per chain, holding exactly the codes that chain wires:
4+
required members for what always runs, ``NotRequired`` members for codes a
5+
setting turns on. Each conditional member is annotated with a
6+
``SocketMeta(help="Needed for ...")`` phrased in input-file vocabulary —
7+
the keyword a user writes is the anchor.
8+
9+
These are the graph-input types of the chain entry points, and the
10+
dispatcher's source of truth for which codes to load and pass: required
11+
keys come from ``__required_keys__``, the purpose strings from
12+
``typing.get_type_hints(..., include_extras=True)``. The module imports no
13+
chain code, so that introspection stays cheap.
14+
"""
15+
16+
# No ``from __future__ import annotations`` here: with stringified
17+
# annotations, ``TypedDict.__required_keys__`` cannot see the
18+
# ``NotRequired`` qualifiers and reports every member as required
19+
# (python/cpython#97727, still true on 3.13).
20+
from typing import Annotated, NotRequired, TypedDict
21+
22+
from aiida import orm
23+
from node_graph.socket_meta import SocketMeta
24+
25+
#: Annotation for the DSCF codes only the periodic Wannier route consumes.
26+
WannierRouteCode = Annotated[
27+
orm.AbstractCode,
28+
SocketMeta(help="Needed for init_orbitals: mlwfs or projwfs."),
29+
]
30+
31+
32+
class PwBandsCodes(TypedDict):
33+
"""Codes for the DFT bands chain (``RunPwBands``)."""
34+
35+
pw: orm.AbstractCode
36+
37+
38+
class DielectricCodes(TypedDict):
39+
"""Codes for the dielectric chain (``DielectricTask``)."""
40+
41+
pw: orm.AbstractCode
42+
ph: orm.AbstractCode
43+
44+
45+
class WannierizeCodes(TypedDict):
46+
"""Codes for the whole-manifold wannierization chain.
47+
48+
Shared by ``Wannierize`` and ``OptimizeWannierization``, whose
49+
upstream builders wire the same scf / nscf / pw2wannier90 / wannier90
50+
steps.
51+
"""
52+
53+
pw: orm.AbstractCode
54+
pw2wannier90: orm.AbstractCode
55+
wannier90: orm.AbstractCode
56+
projwfc: NotRequired[
57+
Annotated[orm.AbstractCode, SocketMeta(help="Needed for SCDM projections.")]
58+
]
59+
60+
61+
class WannierizeBlockCodes(TypedDict):
62+
"""Codes for one block's wannierization (``WannierizeBlock``).
63+
64+
``pw`` never runs here (the block skips scf / nscf) but the upstream
65+
``Wannier90WorkChain`` builder still requires it to assemble the pw
66+
namespaces this graph discards.
67+
"""
68+
69+
pw: orm.AbstractCode
70+
pw2wannier90: orm.AbstractCode
71+
wannier90: orm.AbstractCode
72+
73+
74+
class WannierizeBlocksCodes(TypedDict):
75+
"""Codes for the block-by-block wannierization chain (``WannierizeBlocks``)."""
76+
77+
pw: orm.AbstractCode
78+
pw2wannier90: orm.AbstractCode
79+
wannier90: orm.AbstractCode
80+
wannierjl: NotRequired[
81+
Annotated[
82+
orm.AbstractCode,
83+
SocketMeta(help="Needed when block_wannierization_threshold is set."),
84+
]
85+
]
86+
87+
88+
class SplitBlockCodes(TypedDict):
89+
"""Codes for the wannierize-and-split path (``WannierizeAndSplitBlock``)."""
90+
91+
pw: orm.AbstractCode
92+
pw2wannier90: orm.AbstractCode
93+
wannier90: orm.AbstractCode
94+
wannierjl: orm.AbstractCode
95+
96+
97+
class MlwfInitCodes(TypedDict):
98+
"""Codes for the Wannier-seeded kcp.x initialisation (``MlwfInitialization``)."""
99+
100+
pw: orm.AbstractCode
101+
pw2wannier90: orm.AbstractCode
102+
wannier90: orm.AbstractCode
103+
wann2kcp: orm.AbstractCode
104+
merge_evc: orm.AbstractCode
105+
kcp: orm.AbstractCode
106+
107+
108+
class FoldingCodes(TypedDict):
109+
"""Codes for the fold-to-supercell chain (``FoldToSupercell``)."""
110+
111+
wann2kcp: orm.AbstractCode
112+
merge_evc: orm.AbstractCode
113+
114+
115+
class DscfCodes(TypedDict):
116+
"""Codes for the DSCF chain (``KoopmansDSCFWorkflow`` and the trajectory).
117+
118+
``kcp`` runs every DSCF step; the remaining members exist only for the
119+
periodic Wannier-initialised route
120+
(:func:`~aiida_koopmans.workgraphs.mlwf_init.MlwfInitialization`).
121+
"""
122+
123+
kcp: orm.AbstractCode
124+
pw: NotRequired[WannierRouteCode]
125+
pw2wannier90: NotRequired[WannierRouteCode]
126+
wannier90: NotRequired[WannierRouteCode]
127+
wann2kcp: NotRequired[WannierRouteCode]
128+
merge_evc: NotRequired[WannierRouteCode]
129+
130+
131+
class DfptCodes(TypedDict):
132+
"""Codes for the singlepoint DFPT chain (``SinglepointDFPTWorkflow``)."""
133+
134+
pw: orm.AbstractCode
135+
pw2wannier90: orm.AbstractCode
136+
wannier90: orm.AbstractCode
137+
kcw: orm.AbstractCode
138+
ph: NotRequired[Annotated[orm.AbstractCode, SocketMeta(help="Needed for eps_inf: auto.")]]
139+
140+
141+
class PdosCodes(TypedDict):
142+
"""Codes for the projected-DOS chain (``RunPdos``)."""
143+
144+
pw: orm.AbstractCode
145+
dos: orm.AbstractCode
146+
projwfc: orm.AbstractCode

‎src/aiida_koopmans/workgraphs/dfpt.py‎

Lines changed: 1 addition & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,6 @@
6262
from aiida import orm
6363
from aiida_quantumespresso.common.types import SpinType
6464
from aiida_workgraph import dynamic, task
65-
from aiida_workgraph.socket_spec import SocketMeta
6665

6766
from aiida_koopmans.calculations.kcw import (
6867
KcwHamCalculation,
@@ -84,6 +83,7 @@
8483
WannierizeBlocks,
8584
WannierizeOverrides,
8685
)
86+
from aiida_koopmans.workgraphs.codes import DfptCodes
8787
from aiida_koopmans.workgraphs.ph import DielectricTask
8888
from aiida_koopmans.workgraphs.pw import RunScfNscf
8989
from aiida_koopmans.workgraphs.utils.wannier_merge import (
@@ -314,16 +314,6 @@ class ChannelResults(TypedDict, total=False):
314314
wann2kc_remote_folder: orm.RemoteData
315315

316316

317-
class DfptCodes(TypedDict):
318-
"""Codes for the singlepoint DFPT chain (:func:`SinglepointDFPTWorkflow`)."""
319-
320-
pw: orm.AbstractCode
321-
pw2wannier90: orm.AbstractCode
322-
wannier90: orm.AbstractCode
323-
kcw: orm.AbstractCode
324-
ph: NotRequired[Annotated[orm.AbstractCode, SocketMeta(help="Needed for eps_inf='auto'.")]]
325-
326-
327317
class KoopmansDFPTOutputs(TypedDict):
328318
"""Outputs of :func:`SinglepointDFPTWorkflow`.
329319

‎src/aiida_koopmans/workgraphs/folding.py‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
from aiida_koopmans.projections import ProjectionBlock
5151
from aiida_koopmans.spin import SpinChannel
5252
from aiida_koopmans.workgraphs.block_wannierize import WannierizeBlockOutputs
53+
from aiida_koopmans.workgraphs.codes import FoldingCodes
5354
from aiida_koopmans.workgraphs.utils.wannier_merge import MergeGroup, merge_dest_filename
5455

5556
Wann2kcpTask = task(Wann2kcpCalculation)
@@ -85,13 +86,6 @@ def extract_wannier_files(retrieved: orm.FolderData) -> dict:
8586
return outputs
8687

8788

88-
class FoldingCodes(TypedDict):
89-
"""Codes for the fold-to-supercell chain (:func:`FoldToSupercell`)."""
90-
91-
wann2kcp: orm.AbstractCode
92-
merge_evc: orm.AbstractCode
93-
94-
9589
class FoldToSupercellOutputs(TypedDict, total=False):
9690
"""Folded manifold wavefunctions, one ``SinglefileData`` per kcp.x file.
9791

‎src/aiida_koopmans/workgraphs/kcp.py‎

Lines changed: 1 addition & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@
2929
from aiida_pseudo.data.pseudo.upf import UpfData
3030
from aiida_quantumespresso.workflows.protocols.utils import recursive_merge
3131
from aiida_workgraph import dynamic, task
32-
from aiida_workgraph.socket_spec import SocketMeta
3332

3433
from aiida_koopmans.calculations.kcp import KcpCalculation
3534
from aiida_koopmans.calculations.kcp_inputs import build_kcp_inputs
@@ -52,6 +51,7 @@
5251
WannierizeBlockOutputs,
5352
WannierizeOverrides,
5453
)
54+
from aiida_koopmans.workgraphs.codes import DscfCodes
5555
from aiida_koopmans.workgraphs.convert_spin import convert_spin1_to_spin2
5656
from aiida_koopmans.workgraphs.variational_orbitals import (
5757
assign_orbital_groups,
@@ -90,29 +90,6 @@ class KIFinalOutputs(TypedDict):
9090
remote_folder: orm.RemoteData
9191

9292

93-
#: Annotation for the DSCF codes only the periodic Wannier route consumes.
94-
WannierRouteCode = Annotated[
95-
orm.AbstractCode,
96-
SocketMeta(help="Needed for the Wannier-initialised route (init_orbitals='mlwfs'/'projwfs')."),
97-
]
98-
99-
100-
class DscfCodes(TypedDict):
101-
"""Codes for the DSCF chain (:func:`KoopmansDSCFWorkflow` and the trajectory).
102-
103-
``kcp`` runs every DSCF step; the remaining members exist only for the
104-
periodic Wannier-initialised route
105-
(:func:`~aiida_koopmans.workgraphs.mlwf_init.MlwfInitialization`).
106-
"""
107-
108-
kcp: orm.AbstractCode
109-
pw: NotRequired[WannierRouteCode]
110-
pw2wannier90: NotRequired[WannierRouteCode]
111-
wannier90: NotRequired[WannierRouteCode]
112-
wann2kcp: NotRequired[WannierRouteCode]
113-
merge_evc: NotRequired[WannierRouteCode]
114-
115-
11693
class KoopmansDSCFOutputs(TypedDict):
11794
"""A full KI-DSCF workflow's outputs, with the Wannier-route-only extras.
11895

‎src/aiida_koopmans/workgraphs/mlwf_init.py‎

Lines changed: 1 addition & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
WannierizeBlocks,
4646
WannierizeOverrides,
4747
)
48+
from aiida_koopmans.workgraphs.codes import MlwfInitCodes
4849
from aiida_koopmans.workgraphs.folding import FoldToSupercell, enumerate_fold_targets
4950
from aiida_koopmans.workgraphs.kcp import (
5051
KcpStep,
@@ -67,17 +68,6 @@
6768
}
6869

6970

70-
class MlwfInitCodes(TypedDict):
71-
"""Codes for the Wannier-seeded kcp.x initialisation (:func:`MlwfInitialization`)."""
72-
73-
pw: orm.AbstractCode
74-
pw2wannier90: orm.AbstractCode
75-
wannier90: orm.AbstractCode
76-
wann2kcp: orm.AbstractCode
77-
merge_evc: orm.AbstractCode
78-
kcp: orm.AbstractCode
79-
80-
8171
class MlwfInitializationOutputs(TypedDict):
8272
"""Outputs of :func:`MlwfInitialization`.
8373

‎src/aiida_koopmans/workgraphs/pdos.py‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,7 @@
1616
validate_parallelization,
1717
)
1818
from aiida_koopmans.workgraphs import inject_pseudo_family
19-
20-
21-
class PdosCodes(TypedDict):
22-
"""Codes for the projected-DOS chain (:func:`RunPdos`)."""
23-
24-
pw: orm.AbstractCode
25-
dos: orm.AbstractCode
26-
projwfc: orm.AbstractCode
19+
from aiida_koopmans.workgraphs.codes import PdosCodes
2720

2821

2922
class PdosOutputs(TypedDict):

‎src/aiida_koopmans/workgraphs/ph.py‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
validate_parallelization,
2727
)
2828
from aiida_koopmans.workgraphs import pin_kpoints
29+
from aiida_koopmans.workgraphs.codes import DielectricCodes
2930
from aiida_koopmans.workgraphs.pw import PwBaseStep
3031

3132

@@ -42,13 +43,6 @@ class DielectricOutputs(DielectricConstant):
4243
ph_output_parameters: dict
4344

4445

45-
class DielectricCodes(TypedDict):
46-
"""Codes for the dielectric chain (:func:`DielectricTask`)."""
47-
48-
pw: orm.AbstractCode
49-
ph: orm.AbstractCode
50-
51-
5246
PhBaseTask = task(PhBaseWorkChain)
5347

5448

0 commit comments

Comments
 (0)