From 20c7d371ea7411b01a31c2d018ec4ff7ca9c631c Mon Sep 17 00:00:00 2001 From: Edward Linscott Date: Wed, 8 Jul 2026 13:19:37 +0200 Subject: [PATCH 1/5] Infer valence orbitals for pseudopotentials missing from the semicore tables get_pseudo_orbitals resolved pseudos exclusively by md5 against the bundled semicore tables, so any family without a table (SG15, custom pseudos, ...) failed outright with "Cannot find pseudopotential with md5 ...". Instead of accreting per-family JSON files, resolve in three tiers: an explicit per-kind `overrides` mapping of `PseudoOrbitals` TypedDicts (threaded through Wannier90BaseWorkChain.get_builder_from_protocol as `pseudo_orbitals_overrides`), then the bundled tables, then an aufbau inference from the pseudo's z_valence (cross-checked against the UPF's atomic-wave-function angular momenta when readable, warning emitted, semicores left empty so nothing is silently excluded). The hard error remains only when inference is impossible. The inference and tier precedence are unit-tested in tests/utils/test_pseudo.py with duck-typed pseudo stubs. Co-Authored-By: Claude Fable 5 --- .../utils/pseudo/__init__.py | 171 ++++++++++++++++-- .../workflows/base/wannier90.py | 6 +- tests/utils/test_pseudo.py | 72 ++++++++ 3 files changed, 227 insertions(+), 22 deletions(-) create mode 100644 tests/utils/test_pseudo.py diff --git a/src/aiida_wannier90_workflows/utils/pseudo/__init__.py b/src/aiida_wannier90_workflows/utils/pseudo/__init__.py index f1b8da53..4b8cff43 100644 --- a/src/aiida_wannier90_workflows/utils/pseudo/__init__.py +++ b/src/aiida_wannier90_workflows/utils/pseudo/__init__.py @@ -1,6 +1,7 @@ """Utility functions for pseudo potential family.""" import typing as ty +import warnings from aiida import orm from aiida.common import exceptions @@ -12,6 +13,27 @@ CutoffsPseudoPotentialFamily = GroupFactory("pseudo.family.cutoffs") +# Two-class split only because requires-python is 3.9: merge into a single +# TypedDict with NotRequired[...] fields once python >= 3.11. +class _PseudoOrbitalsTableMeta(ty.TypedDict, total=False): + """Optional provenance metadata carried by the bundled semicore tables.""" + + filename: str + md5: str + + +class PseudoOrbitals(_PseudoOrbitalsTableMeta): + """Valence-orbital description of one pseudopotential. + + * ``pswfcs`` — labels of the pseudo wave functions, e.g. ``["3S", "3P"]``. + * ``semicores`` — the subset of ``pswfcs`` to treat as semicore states + (excluded from the Wannier manifold when ``exclude_semicore`` is on). + """ + + pswfcs: ty.List[str] + semicores: ty.List[str] + + def get_pseudo_and_cutoff( pseudo_family: str, structure: orm.StructureData ) -> ty.Tuple[ty.Mapping[str, PseudoPotentialData], float, float]: @@ -49,24 +71,41 @@ def get_pseudo_and_cutoff( return pseudos, cutoff_wfc, cutoff_rho -def get_pseudo_orbitals(pseudos: ty.Mapping[str, PseudoPotentialData]) -> dict: - """Get the pseudo wave functions contained in the pseudo potential. - - Currently only support the following pseudopotentials installed by `aiida-pseudo`: - * SSSP/1.3/PBE/efficiency - * SSSP/1.3/PBEsol/efficiency - * SSSP/1.1/PBE/efficiency - * SSSP/1.1/PBEsol/efficiency - * PseudoDojo/0.4/LDA/SR/standard/upf - * PseudoDojo/0.4/LDA/SR/stringent/upf - * PseudoDojo/0.4/PBE/SR/standard/upf - * PseudoDojo/0.4/PBE/SR/stringent/upf - * PseudoDojo/0.5/PBE/SR/standard/upf - * PseudoDojo/0.5/PBE/SR/stringent/upf - * PseudoDojo/0.4/PBE/FR/standard/upf - * Pslibrary/1.0.0/relPBE/PAW - ** Pslibrary should be installed manually. - ** Please follow `src/aiida_wannier90_workflows/utils/pseudo/data/__init__.py` +def get_pseudo_orbitals( + pseudos: ty.Mapping[str, PseudoPotentialData], + overrides: ty.Optional[ty.Mapping[str, PseudoOrbitals]] = None, +) -> ty.Dict[str, PseudoOrbitals]: + """Get the valence orbitals (pseudo wave functions) of each pseudopotential. + + Resolution proceeds in three tiers, per kind: + + 1. an entry in ``overrides``, when given; + 2. the bundled semicore tables (matched by md5), which carry curated + ``semicores`` lists for the following families installed by + ``aiida-pseudo``: + + * SSSP/1.3/PBE/efficiency + * SSSP/1.3/PBEsol/efficiency + * SSSP/1.1/PBE/efficiency + * SSSP/1.1/PBEsol/efficiency + * PseudoDojo/0.4/LDA/SR/standard/upf + * PseudoDojo/0.4/LDA/SR/stringent/upf + * PseudoDojo/0.4/PBE/SR/standard/upf + * PseudoDojo/0.4/PBE/SR/stringent/upf + * PseudoDojo/0.5/PBE/SR/standard/upf + * PseudoDojo/0.5/PBE/SR/stringent/upf + * PseudoDojo/0.4/PBE/FR/standard/upf + * Pslibrary/1.0.0/relPBE/PAW + + 3. inference from the pseudo's ``z_valence`` via aufbau filling, with an + empty ``semicores`` list (nothing is auto-excluded) and a warning. + + A ``ValueError`` is raised only when none of the tiers can resolve a + pseudopotential; its message explains what to pass via ``overrides``. + + :param overrides: optional per-kind :class:`PseudoOrbitals` entries that + take precedence over the bundled tables, e.g. ``{"Ti": {"pswfcs": + ["3S", "3P", "4S", "3D"], "semicores": ["3S", "3P"]}}``. """ from .data import load_pseudo_metadata @@ -105,18 +144,110 @@ def get_pseudo_orbitals(pseudos: ty.Mapping[str, PseudoPotentialData]) -> dict: # pseudos dictionary will contain kinds as keys, which may change # e.g. when including Hubbard corrections 'Mn'->'Mn3d' for kind in pseudos: + if overrides is not None and kind in overrides: + pseudo_orbitals[kind] = overrides[kind] + continue for data in pseudo_data: if data.get(pseudos[kind].element, {}).get("md5", "") == pseudos[kind].md5: pseudo_orbitals[kind] = data[pseudos[kind].element] break else: - raise ValueError( - f"Cannot find pseudopotential {kind} with md5 {pseudos[kind].md5}" + inferred = _infer_pseudo_orbitals(pseudos[kind]) + if inferred is None: + raise ValueError( + f"Cannot find pseudopotential {kind} with md5 {pseudos[kind].md5}, " + "and its valence orbitals could not be inferred. Provide the entry " + "explicitly via the `overrides` argument, e.g. " + '{"' + str(kind) + '": {"pswfcs": ["3S", "3P"], "semicores": []}}.' + ) + warnings.warn( + f"Pseudopotential {kind} (md5 {pseudos[kind].md5}) is not in the bundled " + "semicore tables; its valence orbitals were inferred from z_valence " + f"({inferred['pswfcs']}) and no semicore states will be excluded. Pass " + "`overrides` to specify them explicitly." ) + pseudo_orbitals[kind] = inferred return pseudo_orbitals +# Aufbau filling order and per-l occupancies, used to infer the valence +# orbitals of pseudopotentials that are not covered by the bundled tables. +_AUFBAU_ORDER = [ + (1, "S"), (2, "S"), (2, "P"), (3, "S"), (3, "P"), (4, "S"), (3, "D"), + (4, "P"), (5, "S"), (4, "D"), (5, "P"), (6, "S"), (4, "F"), (5, "D"), + (6, "P"), (7, "S"), (5, "F"), (6, "D"), (7, "P"), +] +_L_OCCUPANCY = {"S": 2, "P": 6, "D": 10, "F": 14} +_L_INDEX = {"S": 0, "P": 1, "D": 2, "F": 3} + + +def _infer_pseudo_orbitals(pseudo: PseudoPotentialData) -> ty.Optional[PseudoOrbitals]: + """Infer a ``get_pseudo_orbitals`` entry from the pseudo's ``z_valence``. + + Walks the aufbau filling of the neutral atom from the outermost shell + inward until the pseudo's valence electrons are accounted for; the shells + collected on the way are the valence orbitals (semicore-in-valence + included, e.g. Ti with z_valence 12 yields 3S 3P 4S 3D). ``semicores`` + is left empty: which shells to auto-exclude is a curated judgement, so + nothing is excluded for inferred entries. When the UPF file exposes its + atomic wave functions, their angular-momentum counts are cross-checked + against the inference and a mismatch downgrades the result to ``None``. + + Returns ``None`` when the inference is not possible (no ``z_valence``, + unknown element, or a failed cross-check). + """ + from aiida.common.constants import elements as _aiida_elements + + z_valence = getattr(pseudo, "z_valence", None) + if z_valence is None: + return None + symbol_to_z = {data["symbol"]: z for z, data in _aiida_elements.items()} + z_atom = symbol_to_z.get(pseudo.element) + if z_atom is None: + return None + + filled = [] + remaining = z_atom + for shell_n, shell_l in _AUFBAU_ORDER: + if remaining <= 0: + break + occupancy = min(_L_OCCUPANCY[shell_l], remaining) + filled.append((shell_n, shell_l, occupancy)) + remaining -= occupancy + pswfcs = [] + accounted = 0 + for shell_n, shell_l, occupancy in reversed(filled): + pswfcs.append(f"{shell_n}{shell_l}") + accounted += occupancy + if accounted >= round(z_valence): + break + pswfcs.reverse() + + # Cross-check against the UPF's own atomic wave functions when readable + # (their labels are not exposed, but the angular momenta are). + try: + from upf_to_json import upf_to_json + + upf = upf_to_json(pseudo.get_content(), pseudo.filename)["pseudo_potential"] + upf_l_values = sorted( + wfc["angular_momentum"] for wfc in upf["atomic_wave_functions"] + ) + except Exception: # pylint: disable=broad-except + upf_l_values = None + if upf_l_values: + inferred_l_values = sorted(_L_INDEX[label[-1]] for label in pswfcs) + if inferred_l_values != upf_l_values: + return None + + return PseudoOrbitals( + filename=f"{pseudo.element}.upf", + md5=pseudo.md5, + pswfcs=pswfcs, + semicores=[], + ) + + def get_semicore_list( structure: orm.StructureData, pseudo_orbitals: dict, spin_non_collinear: bool ) -> list: diff --git a/src/aiida_wannier90_workflows/workflows/base/wannier90.py b/src/aiida_wannier90_workflows/workflows/base/wannier90.py index 7b3db456..cfb42bea 100644 --- a/src/aiida_wannier90_workflows/workflows/base/wannier90.py +++ b/src/aiida_wannier90_workflows/workflows/base/wannier90.py @@ -242,6 +242,7 @@ def get_builder_from_protocol( overrides: dict = None, pseudo_family: str = None, external_projectors: dict = None, + pseudo_orbitals_overrides: ty.Optional[ty.Mapping[str, "PseudoOrbitals"]] = None, electronic_type: ElectronicType = ElectronicType.METAL, spin_type: SpinType = SpinType.NONE, projection_type: WannierProjectionType = WannierProjectionType.ATOMIC_PROJECTORS_QE, @@ -269,6 +270,7 @@ def get_builder_from_protocol( get_explicit_kpoints, ) from aiida_wannier90_workflows.utils.pseudo import ( + PseudoOrbitals, get_number_of_projections, get_number_of_projections_ext, get_pseudo_and_cutoff, @@ -376,7 +378,7 @@ def get_builder_from_protocol( num_wann = num_projs if meta_parameters["exclude_semicore"]: - pseudo_orbitals = get_pseudo_orbitals(pseudos) + pseudo_orbitals = get_pseudo_orbitals(pseudos, overrides=pseudo_orbitals_overrides) if projection_type == WannierProjectionType.ATOMIC_PROJECTORS_EXTERNAL: semicore_list = get_semicore_list_ext( structure, external_projectors, pseudo_orbitals, spin_non_collinear @@ -417,7 +419,7 @@ def get_builder_from_protocol( ]: parameters["auto_projections"] = True elif projection_type == WannierProjectionType.ANALYTIC: - pseudo_orbitals = get_pseudo_orbitals(pseudos) + pseudo_orbitals = get_pseudo_orbitals(pseudos, overrides=pseudo_orbitals_overrides) projections = [] if external_projectors is None: for kind in structure.kinds: diff --git a/tests/utils/test_pseudo.py b/tests/utils/test_pseudo.py new file mode 100644 index 00000000..ca501212 --- /dev/null +++ b/tests/utils/test_pseudo.py @@ -0,0 +1,72 @@ +"""Unit tests for the :py:mod:`~aiida_wannier90_workflows.utils.pseudo` module.""" + +import pytest + + +class FakePseudo: + """Duck-typed stand-in for a ``PseudoPotentialData``.""" + + def __init__(self, element, z_valence, md5="0" * 32): + self.element = element + self.z_valence = z_valence + self.md5 = md5 + self.filename = f"{element}.upf" + + def get_content(self): + """Return invalid UPF content, so the wave-function cross-check is skipped.""" + return "not a upf file" + + +@pytest.mark.parametrize( + ("element", "z_valence", "expected"), + ( + ("Si", 4, ["3S", "3P"]), + ("O", 6, ["2S", "2P"]), + # Semicore-in-valence pseudisation: shells collected outermost-inward + # until the valence electrons are accounted for. + ("Ti", 12, ["3S", "3P", "4S", "3D"]), + ("Cu", 19, ["3S", "3P", "4S", "3D"]), + ), +) +def test_infer_pseudo_orbitals(element, z_valence, expected): + """Aufbau inference reproduces the curated-table labels.""" + from aiida_wannier90_workflows.utils.pseudo import _infer_pseudo_orbitals + + entry = _infer_pseudo_orbitals(FakePseudo(element, z_valence)) + assert entry is not None + assert entry["pswfcs"] == expected + assert entry["semicores"] == [] + + +def test_infer_pseudo_orbitals_without_z_valence(): + """No ``z_valence`` means no inference.""" + from aiida_wannier90_workflows.utils.pseudo import _infer_pseudo_orbitals + + assert _infer_pseudo_orbitals(FakePseudo("Si", None)) is None + + +def test_get_pseudo_orbitals_overrides_win(): + """An explicit override bypasses tables and inference.""" + from aiida_wannier90_workflows.utils.pseudo import PseudoOrbitals, get_pseudo_orbitals + + override = PseudoOrbitals(pswfcs=["3S", "3P", "4S", "3D"], semicores=["3S", "3P"]) + result = get_pseudo_orbitals({"Ti": FakePseudo("Ti", 12)}, overrides={"Ti": override}) + assert result["Ti"]["semicores"] == ["3S", "3P"] + + +def test_get_pseudo_orbitals_inference_warns(): + """A pseudo missing from the tables resolves by inference, with a warning.""" + from aiida_wannier90_workflows.utils.pseudo import get_pseudo_orbitals + + with pytest.warns(UserWarning, match="inferred from z_valence"): + result = get_pseudo_orbitals({"Si": FakePseudo("Si", 4)}) + assert result["Si"]["pswfcs"] == ["3S", "3P"] + assert result["Si"]["semicores"] == [] + + +def test_get_pseudo_orbitals_unresolvable_raises(): + """When inference is impossible the error explains the override escape hatch.""" + from aiida_wannier90_workflows.utils.pseudo import get_pseudo_orbitals + + with pytest.raises(ValueError, match="overrides"): + get_pseudo_orbitals({"Xx": FakePseudo("Xx", None)}) From aa5ded49a3404a2050c62bb8bca28a3b9985a0f3 Mon Sep 17 00:00:00 2001 From: Edward Linscott Date: Tue, 14 Jul 2026 15:03:11 +0200 Subject: [PATCH 2/5] Collapse j-split wavefunctions when cross-checking inferred orbitals Fully-relativistic pseudopotentials list one atomic wavefunction per j channel (l > 0 shells appear twice, j = l +/- 1/2), so the raw angular-momentum multiset never matches the per-shell aufbau inference and every FR pseudo outside the bundled tables was rejected. Count each l > 0 shell once per j pair before comparing. Co-Authored-By: Claude Fable 5 --- .../utils/pseudo/__init__.py | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/src/aiida_wannier90_workflows/utils/pseudo/__init__.py b/src/aiida_wannier90_workflows/utils/pseudo/__init__.py index 4b8cff43..b9fe215d 100644 --- a/src/aiida_wannier90_workflows/utils/pseudo/__init__.py +++ b/src/aiida_wannier90_workflows/utils/pseudo/__init__.py @@ -230,9 +230,24 @@ def _infer_pseudo_orbitals(pseudo: PseudoPotentialData) -> ty.Optional[PseudoOrb from upf_to_json import upf_to_json upf = upf_to_json(pseudo.get_content(), pseudo.filename)["pseudo_potential"] - upf_l_values = sorted( - wfc["angular_momentum"] for wfc in upf["atomic_wave_functions"] - ) + wfcs = upf["atomic_wave_functions"] + if wfcs and all( + wfc.get("total_angular_momentum") is not None for wfc in wfcs + ): + # Fully-relativistic pseudos list one wavefunction per j channel + # (l > 0 shells split into j = l +/- 1/2): collapse the split so + # the counts compare against the per-shell inference. + import math + from collections import Counter + + counts = Counter(wfc["angular_momentum"] for wfc in wfcs) + upf_l_values = sorted( + l + for l, count in counts.items() + for _ in range(count if l == 0 else math.ceil(count / 2)) + ) + else: + upf_l_values = sorted(wfc["angular_momentum"] for wfc in wfcs) except Exception: # pylint: disable=broad-except upf_l_values = None if upf_l_values: From 6e50a7ea9cb74cd1cbd8bb650f6e0af1a9233a6d Mon Sep 17 00:00:00 2001 From: Edward Linscott Date: Fri, 17 Jul 2026 10:43:54 +0200 Subject: [PATCH 3/5] Fail closed when inferred valence orbitals cannot be validated against the UPF _infer_pseudo_orbitals returned an unvalidated aufbau guess whenever the UPF cross-check could not run (unparseable content, empty or absent atomic wave functions), so out-of-table heavy and f-block pseudos silently produced wrong orbital sets (e.g. Au z_valence=11 -> ['4F','5D'], Pb z_valence=4 -> ['5D','6P']) instead of raising. The cross-check now reads the UPF angular momenta via upf-tools (added to dependencies), and any failure to read them returns None so get_pseudo_orbitals raises the pre-existing ValueError. Warning text now states the inference was validated against the UPF. Tests drive the cross-check with parseable synthetic UPF content: scalar- and fully-relativistic pass cases, a heavy-element rejection, and a fail-closed unparseable case. Co-Authored-By: Claude Fable 5 --- pyproject.toml | 3 +- .../utils/pseudo/__init__.py | 54 ++++---- tests/utils/test_pseudo.py | 124 ++++++++++++++++-- 3 files changed, 144 insertions(+), 37 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a72995c2..e055c18a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,8 @@ dependencies = [ "aiida-quantumespresso>=4.4", "aiida-wannier90>=2.2", "click>=8.0", - "colorama" + "colorama", + "upf-tools>=0.1.9" ] [project.urls] diff --git a/src/aiida_wannier90_workflows/utils/pseudo/__init__.py b/src/aiida_wannier90_workflows/utils/pseudo/__init__.py index b9fe215d..5d275634 100644 --- a/src/aiida_wannier90_workflows/utils/pseudo/__init__.py +++ b/src/aiida_wannier90_workflows/utils/pseudo/__init__.py @@ -1,11 +1,14 @@ """Utility functions for pseudo potential family.""" +import math import typing as ty import warnings +from collections import Counter from aiida import orm from aiida.common import exceptions from aiida.plugins import DataFactory, GroupFactory +from upf_tools import UPFDict PseudoPotentialData = DataFactory("pseudo") SsspFamily = GroupFactory("pseudo.family.sssp") @@ -162,8 +165,9 @@ def get_pseudo_orbitals( ) warnings.warn( f"Pseudopotential {kind} (md5 {pseudos[kind].md5}) is not in the bundled " - "semicore tables; its valence orbitals were inferred from z_valence " - f"({inferred['pswfcs']}) and no semicore states will be excluded. Pass " + "semicore tables; its valence orbitals were inferred from z_valence and " + "validated against the UPF's atomic wave functions " + f"({inferred['pswfcs']}), and no semicore states will be excluded. Pass " "`overrides` to specify them explicitly." ) pseudo_orbitals[kind] = inferred @@ -224,36 +228,42 @@ def _infer_pseudo_orbitals(pseudo: PseudoPotentialData) -> ty.Optional[PseudoOrb break pswfcs.reverse() - # Cross-check against the UPF's own atomic wave functions when readable - # (their labels are not exposed, but the angular momenta are). + # Cross-check against the UPF's own atomic wave functions (their labels + # are not exposed, but the angular momenta are). The aufbau inference + # above is only trustworthy where it can be validated -- in particular it + # is wrong for f-block and heavy elements, where the reverse-Madelung walk + # mistakes deep-core f shells for valence -- so a UPF whose angular momenta + # cannot be read back (unparseable content, no atomic wave functions, + # unexpected structure) is treated as unresolvable: return None so + # get_pseudo_orbitals raises and asks for an explicit `overrides` entry, + # rather than returning an unvalidated guess. try: - from upf_to_json import upf_to_json - - upf = upf_to_json(pseudo.get_content(), pseudo.filename)["pseudo_potential"] - wfcs = upf["atomic_wave_functions"] - if wfcs and all( - wfc.get("total_angular_momentum") is not None for wfc in wfcs - ): + upf = UPFDict.from_str(pseudo.get_content()) + chi = upf["pswfc"]["chi"] + if not chi: + return None + chi_l_values = [entry["l"] for entry in chi] + if upf["header"]["has_so"]: # Fully-relativistic pseudos list one wavefunction per j channel # (l > 0 shells split into j = l +/- 1/2): collapse the split so # the counts compare against the per-shell inference. - import math - from collections import Counter - - counts = Counter(wfc["angular_momentum"] for wfc in wfcs) + counts = Counter(chi_l_values) upf_l_values = sorted( l for l, count in counts.items() for _ in range(count if l == 0 else math.ceil(count / 2)) ) else: - upf_l_values = sorted(wfc["angular_momentum"] for wfc in wfcs) - except Exception: # pylint: disable=broad-except - upf_l_values = None - if upf_l_values: - inferred_l_values = sorted(_L_INDEX[label[-1]] for label in pswfcs) - if inferred_l_values != upf_l_values: - return None + upf_l_values = sorted(chi_l_values) + except (KeyError, TypeError, ValueError, IndexError, AttributeError, SyntaxError): + # xml.etree ParseError subclasses SyntaxError; a missing PP_PSWFC or + # PP_HEADER section surfaces as KeyError once parsed; unparseable + # content yields a dict without the expected keys. + return None + + inferred_l_values = sorted(_L_INDEX[label[-1]] for label in pswfcs) + if inferred_l_values != upf_l_values: + return None return PseudoOrbitals( filename=f"{pseudo.element}.upf", diff --git a/tests/utils/test_pseudo.py b/tests/utils/test_pseudo.py index ca501212..d1c69921 100644 --- a/tests/utils/test_pseudo.py +++ b/tests/utils/test_pseudo.py @@ -3,41 +3,125 @@ import pytest +def _upf_content(element, z_valence, chi, has_so=False): + """Build minimal UPF v2 content that ``upf_tools`` can parse. + + Only the pieces the valence-orbital cross-check reads are populated with + meaningful values: the ``PP_HEADER`` ``has_so`` flag and one ``PP_CHI`` + per atomic wave function carrying its angular momentum ``l``. + + :param chi: for a scalar-relativistic pseudo, a list of ``l`` values; for a + fully-relativistic pseudo (``has_so=True``), a list of ``(l, j)`` pairs + with one entry per j channel. + """ + has_so_str = "true" if has_so else "false" + chi_blocks, relwfc_blocks = [], [] + for index, entry in enumerate(chi, start=1): + l = entry[0] if has_so else entry + chi_blocks.append( + f' 0.0 0.0 0.0 ' + ) + if has_so: + j = entry[1] + relwfc_blocks.append( + f'' + ) + spin_orb = ( + f'\n{chr(10).join(relwfc_blocks)}\n\n' + if has_so + else "" + ) + return ( + '\n' + f'\n' + ' 0 0.1 0.2 \n' + ' 0 0 0 \n' + ' 0 \n' + f'\n{chr(10).join(chi_blocks)}\n\n' + f'{spin_orb}' + ' 0 0 0 \n' + '\n' + ) + + class FakePseudo: """Duck-typed stand-in for a ``PseudoPotentialData``.""" - def __init__(self, element, z_valence, md5="0" * 32): + def __init__(self, element, z_valence, content="", md5="0" * 32): self.element = element self.z_valence = z_valence self.md5 = md5 self.filename = f"{element}.upf" + self._content = content def get_content(self): - """Return invalid UPF content, so the wave-function cross-check is skipped.""" - return "not a upf file" + """Return the (possibly invalid) UPF content used by the cross-check.""" + return self._content @pytest.mark.parametrize( - ("element", "z_valence", "expected"), + ("element", "z_valence", "chi", "expected"), ( - ("Si", 4, ["3S", "3P"]), - ("O", 6, ["2S", "2P"]), + ("Si", 4, [0, 1], ["3S", "3P"]), + ("O", 6, [0, 1], ["2S", "2P"]), # Semicore-in-valence pseudisation: shells collected outermost-inward # until the valence electrons are accounted for. - ("Ti", 12, ["3S", "3P", "4S", "3D"]), - ("Cu", 19, ["3S", "3P", "4S", "3D"]), + ("Ti", 12, [0, 0, 1, 2], ["3S", "3P", "4S", "3D"]), + ("Cu", 19, [0, 0, 1, 2], ["3S", "3P", "4S", "3D"]), ), ) -def test_infer_pseudo_orbitals(element, z_valence, expected): - """Aufbau inference reproduces the curated-table labels.""" +def test_infer_pseudo_orbitals(element, z_valence, chi, expected): + """Aufbau inference reproduces the curated-table labels and passes the + cross-check against a UPF whose angular momenta agree.""" from aiida_wannier90_workflows.utils.pseudo import _infer_pseudo_orbitals - entry = _infer_pseudo_orbitals(FakePseudo(element, z_valence)) + pseudo = FakePseudo(element, z_valence, _upf_content(element, z_valence, chi)) + entry = _infer_pseudo_orbitals(pseudo) assert entry is not None assert entry["pswfcs"] == expected assert entry["semicores"] == [] +def test_infer_pseudo_orbitals_fully_relativistic(): + """A fully-relativistic pseudo lists one wavefunction per j channel; the + j-split is collapsed before comparing against the per-shell inference.""" + from aiida_wannier90_workflows.utils.pseudo import _infer_pseudo_orbitals + + # 3S (l=0, j=1/2), 3P (l=1, j=1/2), 3P (l=1, j=3/2): raw l-multiset is + # {0, 1, 1}; only after collapsing the p j-channels does it match the + # inferred {0, 1}, so a passing result proves the collapse ran. + content = _upf_content("Si", 4, [(0, 0.5), (1, 0.5), (1, 1.5)], has_so=True) + entry = _infer_pseudo_orbitals(FakePseudo("Si", 4, content)) + assert entry is not None + assert entry["pswfcs"] == ["3S", "3P"] + + +def test_infer_pseudo_orbitals_rejects_contradicting_upf(): + """A heavy element whose aufbau inference disagrees with the UPF's own + wave functions is downgraded to ``None`` rather than trusted. + + Au (z_valence 11) infers ['4F', '5D'] (l-multiset {2, 3}) from the + reverse-Madelung walk, but the real pseudo carries 5d/6s wave functions + (l-multiset {0, 2}); the mismatch must reject the inference. + """ + from aiida_wannier90_workflows.utils.pseudo import _infer_pseudo_orbitals + + content = _upf_content("Au", 11, [0, 2]) + assert _infer_pseudo_orbitals(FakePseudo("Au", 11, content)) is None + + +def test_infer_pseudo_orbitals_fails_closed_when_unparseable(): + """When the UPF cannot be parsed for its angular momenta the inference is + unvalidated and must not be returned (fail closed, not fail open).""" + from aiida_wannier90_workflows.utils.pseudo import _infer_pseudo_orbitals + + assert _infer_pseudo_orbitals(FakePseudo("Au", 11, "not a upf file")) is None + + def test_infer_pseudo_orbitals_without_z_valence(): """No ``z_valence`` means no inference.""" from aiida_wannier90_workflows.utils.pseudo import _infer_pseudo_orbitals @@ -55,11 +139,13 @@ def test_get_pseudo_orbitals_overrides_win(): def test_get_pseudo_orbitals_inference_warns(): - """A pseudo missing from the tables resolves by inference, with a warning.""" + """A pseudo missing from the tables resolves by inference, with a warning + that states the orbitals were validated against the UPF.""" from aiida_wannier90_workflows.utils.pseudo import get_pseudo_orbitals - with pytest.warns(UserWarning, match="inferred from z_valence"): - result = get_pseudo_orbitals({"Si": FakePseudo("Si", 4)}) + pseudo = FakePseudo("Si", 4, _upf_content("Si", 4, [0, 1])) + with pytest.warns(UserWarning, match="validated against the UPF"): + result = get_pseudo_orbitals({"Si": pseudo}) assert result["Si"]["pswfcs"] == ["3S", "3P"] assert result["Si"]["semicores"] == [] @@ -70,3 +156,13 @@ def test_get_pseudo_orbitals_unresolvable_raises(): with pytest.raises(ValueError, match="overrides"): get_pseudo_orbitals({"Xx": FakePseudo("Xx", None)}) + + +def test_get_pseudo_orbitals_heavy_element_raises(): + """A heavy element whose inference contradicts its UPF is unresolvable and + raises rather than silently yielding a wrong orbital set.""" + from aiida_wannier90_workflows.utils.pseudo import get_pseudo_orbitals + + pseudo = FakePseudo("Au", 11, _upf_content("Au", 11, [0, 2])) + with pytest.raises(ValueError, match="overrides"): + get_pseudo_orbitals({"Au": pseudo}) From 136fa58454e3ad4282009c10ce50e643e47e98f5 Mon Sep 17 00:00:00 2001 From: Edward Linscott Date: Fri, 17 Jul 2026 12:02:59 +0200 Subject: [PATCH 4/5] Read out-of-table valence orbitals from the UPF's PP_PSWFC block _infer_pseudo_orbitals guessed the valence orbitals from z_valence via an aufbau (reverse-Madelung) walk and then validated the guess against the UPF's angular momenta. The walk is wrong for f-block and heavy elements, where it mistakes deep-core f shells for valence (Au z_valence=19 -> ['4F','5D'] instead of 5S 5P 5D 6S; Pb -> ['5D','6P'] instead of 5D 6S 6P), so those pseudos failed validation and raised even though the UPF states its own valence. Replace the guess-then-validate path with _derive_pseudo_orbitals_from_upf, which reads the labels directly from each PP_CHI: the spectroscopic letter from the angular momentum l, and the principal quantum number from the label's leading digits. The label's n is preferred over the PP_CHI n attribute on purpose: some ultrasoft/PAW generators write the pseudo principal number there (radial nodes + l + 1), so a node-free 6s/6p/5d projector carries n=1/2/3 while the label keeps the true 6S/6P/5D. Reading the attribute produced nonsense such as Mg -> ['1S','2P'] and Au.pz-rrkjus -> ['2P','3D','1S']; using the label gives ['3S','3P'] and ['6P','5D','6S']. The attribute is used only when the label has no leading digit (empty or non-standard label). Fully-relativistic j-splits collapse to one label per unique (n, l) in file order. The aufbau walk is deleted; z_valence now only backs a soft occupation-sum sanity check that warns rather than rejects. Fail-closed behaviour is preserved: an unreadable or absent PP_PSWFC returns None so get_pseudo_orbitals raises the pre-existing overrides ValueError. Warning text now says the orbitals were read from the pseudopotential file. Verified by deriving orbitals for all 144 PseudoDojo nc-sr-05 SR UPFs (pswfcs match the curated md5 table and occupations sum to z_valence for every file), plus out-of-table ONCV (sr and FR), ultrasoft/PAW and old v1-format pseudos: heavy elements now come out physically correct, the FR j-split collapses, and a real ONCV pseudo with no PP_PSWFC (He) fails closed. Co-Authored-By: Claude Fable 5 --- .../utils/pseudo/__init__.py | 208 ++++++++++-------- tests/utils/test_pseudo.py | 187 ++++++++++------ 2 files changed, 238 insertions(+), 157 deletions(-) diff --git a/src/aiida_wannier90_workflows/utils/pseudo/__init__.py b/src/aiida_wannier90_workflows/utils/pseudo/__init__.py index 5d275634..833cc857 100644 --- a/src/aiida_wannier90_workflows/utils/pseudo/__init__.py +++ b/src/aiida_wannier90_workflows/utils/pseudo/__init__.py @@ -1,9 +1,8 @@ """Utility functions for pseudo potential family.""" -import math +import re import typing as ty import warnings -from collections import Counter from aiida import orm from aiida.common import exceptions @@ -100,8 +99,9 @@ def get_pseudo_orbitals( * PseudoDojo/0.4/PBE/FR/standard/upf * Pslibrary/1.0.0/relPBE/PAW - 3. inference from the pseudo's ``z_valence`` via aufbau filling, with an - empty ``semicores`` list (nothing is auto-excluded) and a warning. + 3. the pseudo's own ``PP_PSWFC`` block, read via ``upf-tools``: the labels + of its atomic wave functions are the valence orbitals, with an empty + ``semicores`` list (nothing is auto-excluded) and a warning. A ``ValueError`` is raised only when none of the tiers can resolve a pseudopotential; its message explains what to pass via ``overrides``. @@ -155,124 +155,160 @@ def get_pseudo_orbitals( pseudo_orbitals[kind] = data[pseudos[kind].element] break else: - inferred = _infer_pseudo_orbitals(pseudos[kind]) - if inferred is None: + derived = _derive_pseudo_orbitals_from_upf(pseudos[kind]) + if derived is None: raise ValueError( f"Cannot find pseudopotential {kind} with md5 {pseudos[kind].md5}, " - "and its valence orbitals could not be inferred. Provide the entry " - "explicitly via the `overrides` argument, e.g. " + "and its valence orbitals could not be read from the pseudopotential " + "file. Provide the entry explicitly via the `overrides` argument, e.g. " '{"' + str(kind) + '": {"pswfcs": ["3S", "3P"], "semicores": []}}.' ) warnings.warn( f"Pseudopotential {kind} (md5 {pseudos[kind].md5}) is not in the bundled " - "semicore tables; its valence orbitals were inferred from z_valence and " - "validated against the UPF's atomic wave functions " - f"({inferred['pswfcs']}), and no semicore states will be excluded. Pass " + "semicore tables; its valence orbitals were read directly from the " + "pseudopotential file's atomic wave functions " + f"({derived['pswfcs']}), and no semicore states will be excluded. Pass " "`overrides` to specify them explicitly." ) - pseudo_orbitals[kind] = inferred + pseudo_orbitals[kind] = derived return pseudo_orbitals -# Aufbau filling order and per-l occupancies, used to infer the valence -# orbitals of pseudopotentials that are not covered by the bundled tables. -_AUFBAU_ORDER = [ - (1, "S"), (2, "S"), (2, "P"), (3, "S"), (3, "P"), (4, "S"), (3, "D"), - (4, "P"), (5, "S"), (4, "D"), (5, "P"), (6, "S"), (4, "F"), (5, "D"), - (6, "P"), (7, "S"), (5, "F"), (6, "D"), (7, "P"), -] -_L_OCCUPANCY = {"S": 2, "P": 6, "D": 10, "F": 14} -_L_INDEX = {"S": 0, "P": 1, "D": 2, "F": 3} +# Angular-momentum quantum number to spectroscopic letter, used to build a +# pseudo wave function's label from its (n, l) when the UPF's own label field +# is missing or unusable. +_L_TO_LETTER = {0: "S", 1: "P", 2: "D", 3: "F", 4: "G", 5: "H"} -def _infer_pseudo_orbitals(pseudo: PseudoPotentialData) -> ty.Optional[PseudoOrbitals]: - """Infer a ``get_pseudo_orbitals`` entry from the pseudo's ``z_valence``. +def _chi_label(entry: ty.Mapping) -> ty.Optional[str]: + """Build the ``n`` + spectroscopic-letter label of one ``PP_CHI`` entry. - Walks the aufbau filling of the neutral atom from the outermost shell - inward until the pseudo's valence electrons are accounted for; the shells - collected on the way are the valence orbitals (semicore-in-valence - included, e.g. Ti with z_valence 12 yields 3S 3P 4S 3D). ``semicores`` - is left empty: which shells to auto-exclude is a curated judgement, so - nothing is excluded for inferred entries. When the UPF file exposes its - atomic wave functions, their angular-momentum counts are cross-checked - against the inference and a mismatch downgrades the result to ``None``. + ``l`` (the angular momentum) is the physical ground truth and is always + carried by a ``PP_CHI`` entry, so the spectroscopic letter is taken from + it. The principal quantum number ``n`` is taken from the leading digits of + the UPF's own ``label`` field, which carries the true atomic ``n`` (``3S``, + ``6P``, ...) across generators. The ``n`` *attribute* is not reliable for + this: some ultrasoft/PAW generators write the *pseudo* principal quantum + number there -- radial nodes + l + 1 -- so a node-free 6s/6p/5d projector + is labelled n=1/2/3 rather than 6/6/5. The attribute is used only as a + fallback when the label has no leading digit (empty or non-standard label + but intact ``n``/``l``). - Returns ``None`` when the inference is not possible (no ``z_valence``, - unknown element, or a failed cross-check). + Returns ``None`` when neither ``l`` nor ``n`` can be determined, so the + caller can fail closed. """ - from aiida.common.constants import elements as _aiida_elements - - z_valence = getattr(pseudo, "z_valence", None) - if z_valence is None: + try: + l = int(entry.get("l")) + except (TypeError, ValueError): return None - symbol_to_z = {data["symbol"]: z for z, data in _aiida_elements.items()} - z_atom = symbol_to_z.get(pseudo.element) - if z_atom is None: + letter = _L_TO_LETTER.get(l) + if letter is None: return None - - filled = [] - remaining = z_atom - for shell_n, shell_l in _AUFBAU_ORDER: - if remaining <= 0: - break - occupancy = min(_L_OCCUPANCY[shell_l], remaining) - filled.append((shell_n, shell_l, occupancy)) - remaining -= occupancy - pswfcs = [] - accounted = 0 - for shell_n, shell_l, occupancy in reversed(filled): - pswfcs.append(f"{shell_n}{shell_l}") - accounted += occupancy - if accounted >= round(z_valence): - break - pswfcs.reverse() - - # Cross-check against the UPF's own atomic wave functions (their labels - # are not exposed, but the angular momenta are). The aufbau inference - # above is only trustworthy where it can be validated -- in particular it - # is wrong for f-block and heavy elements, where the reverse-Madelung walk - # mistakes deep-core f shells for valence -- so a UPF whose angular momenta - # cannot be read back (unparseable content, no atomic wave functions, - # unexpected structure) is treated as unresolvable: return None so - # get_pseudo_orbitals raises and asks for an explicit `overrides` entry, - # rather than returning an unvalidated guess. + label = entry.get("label") + match = re.match(r"\s*(\d+)", label) if isinstance(label, str) else None + n = match.group(1) if match is not None else entry.get("n") + try: + n = int(n) + except (TypeError, ValueError): + return None + return f"{n}{letter}" + + +def _derive_pseudo_orbitals_from_upf( + pseudo: PseudoPotentialData, +) -> ty.Optional[PseudoOrbitals]: + """Read the valence orbitals of a pseudo directly from its ``PP_PSWFC`` block. + + The UPF states its own valence: each ``PP_CHI`` atomic wave function carries + the angular momentum ``l`` and (via its label or ``n`` attribute) the + principal quantum number, so the pseudo wave function labels are read off + rather than guessed. ``semicores`` is left empty -- which shells to + auto-exclude is a curated judgement, so nothing is excluded for + UPF-derived entries. + + Fully-relativistic pseudos list one ``PP_CHI`` per j channel (each l > 0 + shell splits into j = l +/- 1/2); the split is collapsed by keeping each + (n, l) label once, in first-appearance order -- the order in which + ``pw2wannier90.x``/``projwfc.x`` emit the projections. + + Returns ``None`` -- so ``get_pseudo_orbitals`` fails closed and asks for an + explicit ``overrides`` entry -- when the ``PP_PSWFC`` block cannot be read: + unparseable content, no atomic wave functions, or an entry whose (n, l) + cannot be recovered. + """ try: upf = UPFDict.from_str(pseudo.get_content()) chi = upf["pswfc"]["chi"] - if not chi: - return None - chi_l_values = [entry["l"] for entry in chi] - if upf["header"]["has_so"]: - # Fully-relativistic pseudos list one wavefunction per j channel - # (l > 0 shells split into j = l +/- 1/2): collapse the split so - # the counts compare against the per-shell inference. - counts = Counter(chi_l_values) - upf_l_values = sorted( - l - for l, count in counts.items() - for _ in range(count if l == 0 else math.ceil(count / 2)) - ) - else: - upf_l_values = sorted(chi_l_values) except (KeyError, TypeError, ValueError, IndexError, AttributeError, SyntaxError): # xml.etree ParseError subclasses SyntaxError; a missing PP_PSWFC or # PP_HEADER section surfaces as KeyError once parsed; unparseable # content yields a dict without the expected keys. return None - - inferred_l_values = sorted(_L_INDEX[label[-1]] for label in pswfcs) - if inferred_l_values != upf_l_values: + if not chi: return None + labels = [] + for entry in chi: + label = _chi_label(entry) + if label is None: + # A wave function whose (n, l) cannot be recovered makes the whole + # derivation untrustworthy; fail closed rather than emit a partial + # or wrongly-ordered orbital set. + return None + labels.append(label) + # Collapse fully-relativistic j-splits: keep one label per unique (n, l), + # preserving first-appearance order. + pswfcs = list(dict.fromkeys(labels)) + + _warn_if_occupations_disagree(pseudo, chi) + return PseudoOrbitals( - filename=f"{pseudo.element}.upf", + filename=getattr(pseudo, "filename", f"{pseudo.element}.upf"), md5=pseudo.md5, pswfcs=pswfcs, semicores=[], ) +def _warn_if_occupations_disagree( + pseudo: PseudoPotentialData, chi: ty.Sequence[ty.Mapping] +) -> None: + """Warn when the ``PP_CHI`` occupations do not sum to roughly ``z_valence``. + + The atomic wave function occupations should account for the pseudo's + valence electrons; a wild disagreement means the wave functions were + misread or the pseudo is unusual. This is a soft sanity net -- the labels + remain the ground truth, so a mismatch warns rather than rejecting the + derived orbitals. The check is skipped when either quantity is unavailable. + """ + z_valence = getattr(pseudo, "z_valence", None) + if z_valence is None: + return + try: + occupations = [ + float(entry["occupation"]) + for entry in chi + if entry.get("occupation") is not None + ] + z_valence = float(z_valence) + except (KeyError, TypeError, ValueError): + return + if not occupations: + return + occ_sum = sum(occupations) + # Tolerance grows with z_valence so an ionised reference configuration (a + # few electrons short) does not trip the warning, while a gross misread + # (e.g. half the shells dropped) still does. + if abs(occ_sum - z_valence) > max(1.0, 0.2 * z_valence): + warnings.warn( + f"Pseudopotential {pseudo.element} (md5 {pseudo.md5}): its atomic " + f"wave function occupations sum to {occ_sum:g}, which disagrees with " + f"its z_valence {z_valence:g}. The valence orbitals were still read " + "from the pseudopotential file; double-check them." + ) + + def get_semicore_list( structure: orm.StructureData, pseudo_orbitals: dict, spin_non_collinear: bool ) -> list: diff --git a/tests/utils/test_pseudo.py b/tests/utils/test_pseudo.py index d1c69921..3297c340 100644 --- a/tests/utils/test_pseudo.py +++ b/tests/utils/test_pseudo.py @@ -3,29 +3,43 @@ import pytest -def _upf_content(element, z_valence, chi, has_so=False): +def _upf_content( + element, z_valence, chi, has_so=False, labels=None, occupations=None, n_attrs=None +): """Build minimal UPF v2 content that ``upf_tools`` can parse. - Only the pieces the valence-orbital cross-check reads are populated with - meaningful values: the ``PP_HEADER`` ``has_so`` flag and one ``PP_CHI`` - per atomic wave function carrying its angular momentum ``l``. - - :param chi: for a scalar-relativistic pseudo, a list of ``l`` values; for a - fully-relativistic pseudo (``has_so=True``), a list of ``(l, j)`` pairs - with one entry per j channel. + Only the pieces the valence-orbital derivation reads are populated with + meaningful values: the ``PP_HEADER`` ``has_so``/``z_valence`` fields and one + ``PP_CHI`` per atomic wave function carrying its ``n``, ``l`` and label. + + :param chi: for a scalar-relativistic pseudo, a list of ``(n, l)`` pairs; + for a fully-relativistic pseudo (``has_so=True``), a list of + ``(n, l, j)`` triples with one entry per j channel. + :param labels: optional per-entry override of the ``label`` attribute, to + exercise reconstruction from ``n``/``l`` when the label is unusable. + :param occupations: optional per-entry occupation values (default 1.0 + each), to exercise the occupation/z_valence sanity check. + :param n_attrs: optional per-entry override of the emitted ``n`` attribute + (defaults to the label's ``n``), to exercise UPFs whose ``n`` attribute + is the pseudo (node-based) principal number rather than the atomic one. """ has_so_str = "true" if has_so else "false" + letters = {0: "S", 1: "P", 2: "D", 3: "F"} chi_blocks, relwfc_blocks = [], [] for index, entry in enumerate(chi, start=1): - l = entry[0] if has_so else entry + n, l = entry[0], entry[1] + label = f"{n}{letters[l]}" if labels is None else labels[index - 1] + occupation = 1.0 if occupations is None else occupations[index - 1] + n_attr = n if n_attrs is None else n_attrs[index - 1] + label_attr = "" if label is None else f'label="{label}" ' chi_blocks.append( - f' 0.0 0.0 0.0 ' + f' 0.0 0.0 0.0 ' ) if has_so: - j = entry[1] + j = entry[2] relwfc_blocks.append( - f'' + f'' ) spin_orb = ( f'\n{chr(10).join(relwfc_blocks)}\n\n' @@ -59,78 +73,119 @@ def __init__(self, element, z_valence, content="", md5="0" * 32): self._content = content def get_content(self): - """Return the (possibly invalid) UPF content used by the cross-check.""" + """Return the (possibly invalid) UPF content read by the derivation.""" return self._content @pytest.mark.parametrize( ("element", "z_valence", "chi", "expected"), ( - ("Si", 4, [0, 1], ["3S", "3P"]), - ("O", 6, [0, 1], ["2S", "2P"]), - # Semicore-in-valence pseudisation: shells collected outermost-inward - # until the valence electrons are accounted for. - ("Ti", 12, [0, 0, 1, 2], ["3S", "3P", "4S", "3D"]), - ("Cu", 19, [0, 0, 1, 2], ["3S", "3P", "4S", "3D"]), + ("Si", 4, [(3, 0), (3, 1)], ["3S", "3P"]), + ("O", 6, [(2, 0), (2, 1)], ["2S", "2P"]), + ("H", 1, [(1, 0)], ["1S"]), + # Semicore-in-valence pseudisation: the PP_CHI labels are read off in + # the order they appear in the file (which is the order the curated + # tables and QE use), not re-sorted into an aufbau order. + ("Ti", 12, [(3, 0), (3, 1), (3, 2), (4, 0)], ["3S", "3P", "3D", "4S"]), + ("Cu", 19, [(3, 0), (3, 1), (3, 2), (4, 0)], ["3S", "3P", "3D", "4S"]), + # Heavy element the old aufbau walk got wrong (it inferred ['4F','5D']): + # reading PP_CHI gives the physical 5s/5p/5d/6s valence. + ("Au", 19, [(5, 0), (5, 1), (5, 2), (6, 0)], ["5S", "5P", "5D", "6S"]), ), ) -def test_infer_pseudo_orbitals(element, z_valence, chi, expected): - """Aufbau inference reproduces the curated-table labels and passes the - cross-check against a UPF whose angular momenta agree.""" - from aiida_wannier90_workflows.utils.pseudo import _infer_pseudo_orbitals +def test_derive_pseudo_orbitals(element, z_valence, chi, expected): + """The valence orbitals are read directly from the UPF's PP_PSWFC labels.""" + from aiida_wannier90_workflows.utils.pseudo import _derive_pseudo_orbitals_from_upf pseudo = FakePseudo(element, z_valence, _upf_content(element, z_valence, chi)) - entry = _infer_pseudo_orbitals(pseudo) + entry = _derive_pseudo_orbitals_from_upf(pseudo) assert entry is not None assert entry["pswfcs"] == expected assert entry["semicores"] == [] -def test_infer_pseudo_orbitals_fully_relativistic(): - """A fully-relativistic pseudo lists one wavefunction per j channel; the - j-split is collapsed before comparing against the per-shell inference.""" - from aiida_wannier90_workflows.utils.pseudo import _infer_pseudo_orbitals +def test_derive_pseudo_orbitals_fully_relativistic(): + """A fully-relativistic pseudo lists one PP_CHI per j channel; the j-split + is collapsed to one label per (n, l), preserving order.""" + from aiida_wannier90_workflows.utils.pseudo import _derive_pseudo_orbitals_from_upf + + # Pb-like: 5d (j=3/2, 5/2), 6s (j=1/2), 6p (j=1/2, 3/2). The raw label list + # has 5D and 6P twice each; collapsing the j-channels yields the unique set. + chi = [(5, 2, 1.5), (5, 2, 2.5), (6, 0, 0.5), (6, 1, 0.5), (6, 1, 1.5)] + content = _upf_content("Pb", 14, chi, has_so=True) + entry = _derive_pseudo_orbitals_from_upf(FakePseudo("Pb", 14, content)) + assert entry is not None + assert entry["pswfcs"] == ["5D", "6S", "6P"] + + +def test_derive_pseudo_orbitals_reconstructs_unusable_label(): + """When a PP_CHI's label field is empty or non-standard, the label is + reconstructed from its ``n`` and ``l`` rather than failing.""" + from aiida_wannier90_workflows.utils.pseudo import _derive_pseudo_orbitals_from_upf - # 3S (l=0, j=1/2), 3P (l=1, j=1/2), 3P (l=1, j=3/2): raw l-multiset is - # {0, 1, 1}; only after collapsing the p j-channels does it match the - # inferred {0, 1}, so a passing result proves the collapse ran. - content = _upf_content("Si", 4, [(0, 0.5), (1, 0.5), (1, 1.5)], has_so=True) - entry = _infer_pseudo_orbitals(FakePseudo("Si", 4, content)) + # First label empty, second label non-standard: both must be rebuilt from + # the intact n/l attributes. + content = _upf_content("Si", 4, [(3, 0), (3, 1)], labels=[None, "3p_orbital"]) + entry = _derive_pseudo_orbitals_from_upf(FakePseudo("Si", 4, content)) assert entry is not None assert entry["pswfcs"] == ["3S", "3P"] -def test_infer_pseudo_orbitals_rejects_contradicting_upf(): - """A heavy element whose aufbau inference disagrees with the UPF's own - wave functions is downgraded to ``None`` rather than trusted. +def test_derive_pseudo_orbitals_label_wins_over_node_based_n(): + """Some ultrasoft/PAW UPFs write the pseudo (node-based) principal number in + the ``n`` attribute -- a node-free 6s/6p/5d projector gets n=1/2/3 -- while + the label keeps the true atomic n. The label must win.""" + from aiida_wannier90_workflows.utils.pseudo import _derive_pseudo_orbitals_from_upf - Au (z_valence 11) infers ['4F', '5D'] (l-multiset {2, 3}) from the - reverse-Madelung walk, but the real pseudo carries 5d/6s wave functions - (l-multiset {0, 2}); the mismatch must reject the inference. - """ - from aiida_wannier90_workflows.utils.pseudo import _infer_pseudo_orbitals + # Au.pz-rrkjus_aewfc-like: labels 6P/5D/6S with node-based n attributes + # 2/3/1. Reading the n attribute would give the nonsense ['2P','3D','1S']. + content = _upf_content( + "Au", 11, [(6, 1), (5, 2), (6, 0)], + occupations=[0.0, 10.0, 1.0], n_attrs=[2, 3, 1], + ) + entry = _derive_pseudo_orbitals_from_upf(FakePseudo("Au", 11, content)) + assert entry is not None + assert entry["pswfcs"] == ["6P", "5D", "6S"] + + +def test_derive_pseudo_orbitals_no_pswfc_returns_none(): + """A UPF without a PP_PSWFC block cannot be read and fails closed.""" + from aiida_wannier90_workflows.utils.pseudo import _derive_pseudo_orbitals_from_upf - content = _upf_content("Au", 11, [0, 2]) - assert _infer_pseudo_orbitals(FakePseudo("Au", 11, content)) is None + content = ( + '\n' + '\n' + ' 0 0.1 0.2 \n' + '\n' + ) + assert _derive_pseudo_orbitals_from_upf(FakePseudo("Si", 4, content)) is None -def test_infer_pseudo_orbitals_fails_closed_when_unparseable(): - """When the UPF cannot be parsed for its angular momenta the inference is - unvalidated and must not be returned (fail closed, not fail open).""" - from aiida_wannier90_workflows.utils.pseudo import _infer_pseudo_orbitals +def test_derive_pseudo_orbitals_unparseable_returns_none(): + """Unparseable content cannot be read for its wave functions (fail closed, + not fail open).""" + from aiida_wannier90_workflows.utils.pseudo import _derive_pseudo_orbitals_from_upf - assert _infer_pseudo_orbitals(FakePseudo("Au", 11, "not a upf file")) is None + assert _derive_pseudo_orbitals_from_upf(FakePseudo("Au", 19, "not a upf file")) is None -def test_infer_pseudo_orbitals_without_z_valence(): - """No ``z_valence`` means no inference.""" - from aiida_wannier90_workflows.utils.pseudo import _infer_pseudo_orbitals +def test_derive_pseudo_orbitals_warns_on_occupation_mismatch(): + """A gross disagreement between the summed PP_CHI occupations and z_valence + warns, but the labels are still returned (labels are the ground truth).""" + from aiida_wannier90_workflows.utils.pseudo import _derive_pseudo_orbitals_from_upf - assert _infer_pseudo_orbitals(FakePseudo("Si", None)) is None + # Occupations sum to 2.0 but z_valence claims 12: well past the tolerance. + content = _upf_content( + "Ti", 12, [(3, 0), (3, 1), (3, 2), (4, 0)], occupations=[0.5, 0.5, 0.5, 0.5] + ) + with pytest.warns(UserWarning, match="disagrees with"): + entry = _derive_pseudo_orbitals_from_upf(FakePseudo("Ti", 12, content)) + assert entry["pswfcs"] == ["3S", "3P", "3D", "4S"] def test_get_pseudo_orbitals_overrides_win(): - """An explicit override bypasses tables and inference.""" + """An explicit override bypasses tables and UPF derivation.""" from aiida_wannier90_workflows.utils.pseudo import PseudoOrbitals, get_pseudo_orbitals override = PseudoOrbitals(pswfcs=["3S", "3P", "4S", "3D"], semicores=["3S", "3P"]) @@ -138,31 +193,21 @@ def test_get_pseudo_orbitals_overrides_win(): assert result["Ti"]["semicores"] == ["3S", "3P"] -def test_get_pseudo_orbitals_inference_warns(): - """A pseudo missing from the tables resolves by inference, with a warning - that states the orbitals were validated against the UPF.""" +def test_get_pseudo_orbitals_derivation_warns(): + """A pseudo missing from the tables resolves by UPF derivation, with a + warning that states the orbitals were read from the pseudopotential file.""" from aiida_wannier90_workflows.utils.pseudo import get_pseudo_orbitals - pseudo = FakePseudo("Si", 4, _upf_content("Si", 4, [0, 1])) - with pytest.warns(UserWarning, match="validated against the UPF"): + pseudo = FakePseudo("Si", 4, _upf_content("Si", 4, [(3, 0), (3, 1)])) + with pytest.warns(UserWarning, match="read directly from the pseudopotential file"): result = get_pseudo_orbitals({"Si": pseudo}) assert result["Si"]["pswfcs"] == ["3S", "3P"] assert result["Si"]["semicores"] == [] def test_get_pseudo_orbitals_unresolvable_raises(): - """When inference is impossible the error explains the override escape hatch.""" - from aiida_wannier90_workflows.utils.pseudo import get_pseudo_orbitals - - with pytest.raises(ValueError, match="overrides"): - get_pseudo_orbitals({"Xx": FakePseudo("Xx", None)}) - - -def test_get_pseudo_orbitals_heavy_element_raises(): - """A heavy element whose inference contradicts its UPF is unresolvable and - raises rather than silently yielding a wrong orbital set.""" + """When derivation is impossible the error explains the override escape hatch.""" from aiida_wannier90_workflows.utils.pseudo import get_pseudo_orbitals - pseudo = FakePseudo("Au", 11, _upf_content("Au", 11, [0, 2])) with pytest.raises(ValueError, match="overrides"): - get_pseudo_orbitals({"Au": pseudo}) + get_pseudo_orbitals({"Xx": FakePseudo("Xx", 4, "not a upf file")}) From a574b99c458fc6d8a88a57922e207aebcd205c7d Mon Sep 17 00:00:00 2001 From: Edward Linscott Date: Fri, 17 Jul 2026 13:29:02 +0200 Subject: [PATCH 5/5] Fail closed when PP_CHI occupations disagree with z_valence A disagreement can indicate a valence orbital missing from PP_PSWFC, which previously produced only a warning while returning a possibly incomplete orbital set. Zero of 1202 real UPFs tested trip the check. Co-Authored-By: Claude Fable 5 --- .../utils/pseudo/__init__.py | 50 ++++++++------- tests/utils/test_pseudo.py | 64 +++++++++++-------- 2 files changed, 66 insertions(+), 48 deletions(-) diff --git a/src/aiida_wannier90_workflows/utils/pseudo/__init__.py b/src/aiida_wannier90_workflows/utils/pseudo/__init__.py index 833cc857..7852b212 100644 --- a/src/aiida_wannier90_workflows/utils/pseudo/__init__.py +++ b/src/aiida_wannier90_workflows/utils/pseudo/__init__.py @@ -4,10 +4,11 @@ import typing as ty import warnings +from upf_tools import UPFDict + from aiida import orm from aiida.common import exceptions from aiida.plugins import DataFactory, GroupFactory -from upf_tools import UPFDict PseudoPotentialData = DataFactory("pseudo") SsspFamily = GroupFactory("pseudo.family.sssp") @@ -233,9 +234,10 @@ def _derive_pseudo_orbitals_from_upf( ``pw2wannier90.x``/``projwfc.x`` emit the projections. Returns ``None`` -- so ``get_pseudo_orbitals`` fails closed and asks for an - explicit ``overrides`` entry -- when the ``PP_PSWFC`` block cannot be read: - unparseable content, no atomic wave functions, or an entry whose (n, l) - cannot be recovered. + explicit ``overrides`` entry -- when the ``PP_PSWFC`` block cannot be read + (unparseable content, no atomic wave functions, or an entry whose (n, l) + cannot be recovered) or when the wave function occupations do not sum to + ``z_valence``, which can indicate a missing valence orbital. """ try: upf = UPFDict.from_str(pseudo.get_content()) @@ -261,7 +263,8 @@ def _derive_pseudo_orbitals_from_upf( # preserving first-appearance order. pswfcs = list(dict.fromkeys(labels)) - _warn_if_occupations_disagree(pseudo, chi) + if _occupations_disagree(pseudo, chi): + return None return PseudoOrbitals( filename=getattr(pseudo, "filename", f"{pseudo.element}.upf"), @@ -271,20 +274,21 @@ def _derive_pseudo_orbitals_from_upf( ) -def _warn_if_occupations_disagree( +def _occupations_disagree( pseudo: PseudoPotentialData, chi: ty.Sequence[ty.Mapping] -) -> None: - """Warn when the ``PP_CHI`` occupations do not sum to roughly ``z_valence``. - - The atomic wave function occupations should account for the pseudo's - valence electrons; a wild disagreement means the wave functions were - misread or the pseudo is unusual. This is a soft sanity net -- the labels - remain the ground truth, so a mismatch warns rather than rejecting the - derived orbitals. The check is skipped when either quantity is unavailable. +) -> bool: + """Return True when the ``PP_CHI`` occupations do not sum to ``z_valence``. + + The atomic wave function occupations must account for the pseudo's valence + electrons; a disagreement can mean a valence orbital is missing from + ``PP_PSWFC``, in which case the derived labels would be silently + incomplete. The derivation therefore fails closed on any disagreement + beyond numerical noise. The check is skipped when either quantity is + unavailable. """ z_valence = getattr(pseudo, "z_valence", None) if z_valence is None: - return + return False try: occupations = [ float(entry["occupation"]) @@ -293,20 +297,20 @@ def _warn_if_occupations_disagree( ] z_valence = float(z_valence) except (KeyError, TypeError, ValueError): - return + return False if not occupations: - return + return False occ_sum = sum(occupations) - # Tolerance grows with z_valence so an ionised reference configuration (a - # few electrons short) does not trip the warning, while a gross misread - # (e.g. half the shells dropped) still does. - if abs(occ_sum - z_valence) > max(1.0, 0.2 * z_valence): + if abs(occ_sum - z_valence) > 1.0e-3: warnings.warn( f"Pseudopotential {pseudo.element} (md5 {pseudo.md5}): its atomic " f"wave function occupations sum to {occ_sum:g}, which disagrees with " - f"its z_valence {z_valence:g}. The valence orbitals were still read " - "from the pseudopotential file; double-check them." + f"its z_valence {z_valence:g}, so a valence orbital may be missing " + "from PP_PSWFC. Refusing the derived orbitals; provide an explicit " + "`overrides` entry instead." ) + return True + return False def get_semicore_list( diff --git a/tests/utils/test_pseudo.py b/tests/utils/test_pseudo.py index 3297c340..9b564fde 100644 --- a/tests/utils/test_pseudo.py +++ b/tests/utils/test_pseudo.py @@ -17,8 +17,9 @@ def _upf_content( ``(n, l, j)`` triples with one entry per j channel. :param labels: optional per-entry override of the ``label`` attribute, to exercise reconstruction from ``n``/``l`` when the label is unusable. - :param occupations: optional per-entry occupation values (default 1.0 - each), to exercise the occupation/z_valence sanity check. + :param occupations: optional per-entry occupation values (by default the + ``z_valence`` electrons are spread evenly so the sum agrees), to + exercise the occupation/z_valence sanity check. :param n_attrs: optional per-entry override of the emitted ``n`` attribute (defaults to the label's ``n``), to exercise UPFs whose ``n`` attribute is the pseudo (node-based) principal number rather than the atomic one. @@ -29,20 +30,22 @@ def _upf_content( for index, entry in enumerate(chi, start=1): n, l = entry[0], entry[1] label = f"{n}{letters[l]}" if labels is None else labels[index - 1] - occupation = 1.0 if occupations is None else occupations[index - 1] + occupation = ( + float(z_valence) / len(chi) + if occupations is None + else occupations[index - 1] + ) n_attr = n if n_attrs is None else n_attrs[index - 1] label_attr = "" if label is None else f'label="{label}" ' chi_blocks.append( f' 0.0 0.0 0.0 ' + f"{label_attr}> 0.0 0.0 0.0 " ) if has_so: j = entry[2] - relwfc_blocks.append( - f'' - ) + relwfc_blocks.append(f'') spin_orb = ( - f'\n{chr(10).join(relwfc_blocks)}\n\n' + f"\n{chr(10).join(relwfc_blocks)}\n\n" if has_so else "" ) @@ -54,11 +57,11 @@ def _upf_content( f'is_paw="false" has_so="{has_so_str}" l_max="3" l_max_rho="0"/>\n' ' 0 0.1 0.2 \n' ' 0 0 0 \n' - ' 0 \n' - f'\n{chr(10).join(chi_blocks)}\n\n' - f'{spin_orb}' + " 0 \n" + f"\n{chr(10).join(chi_blocks)}\n\n" + f"{spin_orb}" ' 0 0 0 \n' - '\n' + "\n" ) @@ -140,8 +143,11 @@ def test_derive_pseudo_orbitals_label_wins_over_node_based_n(): # Au.pz-rrkjus_aewfc-like: labels 6P/5D/6S with node-based n attributes # 2/3/1. Reading the n attribute would give the nonsense ['2P','3D','1S']. content = _upf_content( - "Au", 11, [(6, 1), (5, 2), (6, 0)], - occupations=[0.0, 10.0, 1.0], n_attrs=[2, 3, 1], + "Au", + 11, + [(6, 1), (5, 2), (6, 0)], + occupations=[0.0, 10.0, 1.0], + n_attrs=[2, 3, 1], ) entry = _derive_pseudo_orbitals_from_upf(FakePseudo("Au", 11, content)) assert entry is not None @@ -157,7 +163,7 @@ def test_derive_pseudo_orbitals_no_pswfc_returns_none(): '\n' ' 0 0.1 0.2 \n' - '\n' + "\n" ) assert _derive_pseudo_orbitals_from_upf(FakePseudo("Si", 4, content)) is None @@ -167,29 +173,37 @@ def test_derive_pseudo_orbitals_unparseable_returns_none(): not fail open).""" from aiida_wannier90_workflows.utils.pseudo import _derive_pseudo_orbitals_from_upf - assert _derive_pseudo_orbitals_from_upf(FakePseudo("Au", 19, "not a upf file")) is None + assert ( + _derive_pseudo_orbitals_from_upf(FakePseudo("Au", 19, "not a upf file")) is None + ) -def test_derive_pseudo_orbitals_warns_on_occupation_mismatch(): - """A gross disagreement between the summed PP_CHI occupations and z_valence - warns, but the labels are still returned (labels are the ground truth).""" +def test_derive_pseudo_orbitals_fails_closed_on_occupation_mismatch(): + """A disagreement between the summed PP_CHI occupations and z_valence can + mean a valence orbital is missing from PP_PSWFC, so the derivation warns + and fails closed rather than returning a possibly-incomplete set.""" from aiida_wannier90_workflows.utils.pseudo import _derive_pseudo_orbitals_from_upf - # Occupations sum to 2.0 but z_valence claims 12: well past the tolerance. + # Occupations sum to 10.0 but z_valence claims 12 -- e.g. a semicore 3S + # missing from PP_PSWFC. content = _upf_content( - "Ti", 12, [(3, 0), (3, 1), (3, 2), (4, 0)], occupations=[0.5, 0.5, 0.5, 0.5] + "Ti", 12, [(3, 1), (3, 2), (4, 0)], occupations=[6.0, 2.0, 2.0] ) with pytest.warns(UserWarning, match="disagrees with"): - entry = _derive_pseudo_orbitals_from_upf(FakePseudo("Ti", 12, content)) - assert entry["pswfcs"] == ["3S", "3P", "3D", "4S"] + assert _derive_pseudo_orbitals_from_upf(FakePseudo("Ti", 12, content)) is None def test_get_pseudo_orbitals_overrides_win(): """An explicit override bypasses tables and UPF derivation.""" - from aiida_wannier90_workflows.utils.pseudo import PseudoOrbitals, get_pseudo_orbitals + from aiida_wannier90_workflows.utils.pseudo import ( + PseudoOrbitals, + get_pseudo_orbitals, + ) override = PseudoOrbitals(pswfcs=["3S", "3P", "4S", "3D"], semicores=["3S", "3P"]) - result = get_pseudo_orbitals({"Ti": FakePseudo("Ti", 12)}, overrides={"Ti": override}) + result = get_pseudo_orbitals( + {"Ti": FakePseudo("Ti", 12)}, overrides={"Ti": override} + ) assert result["Ti"]["semicores"] == ["3S", "3P"]