diff --git a/pyproject.toml b/pyproject.toml index dcc6d9b6..95b46ead 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,8 @@ dependencies = [ "aiida-quantumespresso>=4.4,<5", "aiida-wannier90 @ git+https://github.com/aiidateam/aiida-wannier90.git@v2.2.0", "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 f1b8da53..7852b212 100644 --- a/src/aiida_wannier90_workflows/utils/pseudo/__init__.py +++ b/src/aiida_wannier90_workflows/utils/pseudo/__init__.py @@ -1,6 +1,10 @@ """Utility functions for pseudo potential family.""" +import re import typing as ty +import warnings + +from upf_tools import UPFDict from aiida import orm from aiida.common import exceptions @@ -12,6 +16,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 +74,42 @@ 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. 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``. + + :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 +148,171 @@ 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}" + 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 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 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] = derived return pseudo_orbitals +# 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 _chi_label(entry: ty.Mapping) -> ty.Optional[str]: + """Build the ``n`` + spectroscopic-letter label of one ``PP_CHI`` entry. + + ``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 neither ``l`` nor ``n`` can be determined, so the + caller can fail closed. + """ + try: + l = int(entry.get("l")) + except (TypeError, ValueError): + return None + letter = _L_TO_LETTER.get(l) + if letter is None: + return None + 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) 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()) + chi = upf["pswfc"]["chi"] + 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 + 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)) + + if _occupations_disagree(pseudo, chi): + return None + + return PseudoOrbitals( + filename=getattr(pseudo, "filename", f"{pseudo.element}.upf"), + md5=pseudo.md5, + pswfcs=pswfcs, + semicores=[], + ) + + +def _occupations_disagree( + pseudo: PseudoPotentialData, chi: ty.Sequence[ty.Mapping] +) -> 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 False + 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 False + if not occupations: + return False + occ_sum = sum(occupations) + 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}, 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( 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..177244c1 100644 --- a/src/aiida_wannier90_workflows/workflows/base/wannier90.py +++ b/src/aiida_wannier90_workflows/workflows/base/wannier90.py @@ -242,6 +242,9 @@ 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 +272,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 +380,9 @@ 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 +423,9 @@ 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..9b564fde --- /dev/null +++ b/tests/utils/test_pseudo.py @@ -0,0 +1,227 @@ +"""Unit tests for the :py:mod:`~aiida_wannier90_workflows.utils.pseudo` module.""" + +import pytest + + +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 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 (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. + """ + 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): + n, l = entry[0], entry[1] + label = f"{n}{letters[l]}" if labels is None else labels[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 " + ) + if has_so: + j = entry[2] + 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, 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 the (possibly invalid) UPF content read by the derivation.""" + return self._content + + +@pytest.mark.parametrize( + ("element", "z_valence", "chi", "expected"), + ( + ("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_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 = _derive_pseudo_orbitals_from_upf(pseudo) + assert entry is not None + assert entry["pswfcs"] == expected + assert entry["semicores"] == [] + + +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 + + # 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_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.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 = ( + '\n' + '\n' + ' 0 0.1 0.2 \n' + "\n" + ) + assert _derive_pseudo_orbitals_from_upf(FakePseudo("Si", 4, content)) is None + + +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 ( + _derive_pseudo_orbitals_from_upf(FakePseudo("Au", 19, "not a upf file")) is None + ) + + +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 10.0 but z_valence claims 12 -- e.g. a semicore 3S + # missing from PP_PSWFC. + content = _upf_content( + "Ti", 12, [(3, 1), (3, 2), (4, 0)], occupations=[6.0, 2.0, 2.0] + ) + with pytest.warns(UserWarning, match="disagrees with"): + 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, + ) + + 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_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, [(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 derivation 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", 4, "not a upf file")})