diff --git a/src/aiida_wannier90_workflows/utils/pseudo/__init__.py b/src/aiida_wannier90_workflows/utils/pseudo/__init__.py index f1b8da53..7134285d 100644 --- a/src/aiida_wannier90_workflows/utils/pseudo/__init__.py +++ b/src/aiida_wannier90_workflows/utils/pseudo/__init__.py @@ -7,35 +7,68 @@ from aiida.plugins import DataFactory, GroupFactory PseudoPotentialData = DataFactory("pseudo") +PseudoPotentialFamily = GroupFactory("pseudo.family") SsspFamily = GroupFactory("pseudo.family.sssp") PseudoDojoFamily = GroupFactory("pseudo.family.pseudo_dojo") CutoffsPseudoPotentialFamily = GroupFactory("pseudo.family.cutoffs") -def get_pseudo_and_cutoff( - pseudo_family: str, structure: orm.StructureData -) -> ty.Tuple[ty.Mapping[str, PseudoPotentialData], float, float]: - """Get pseudo potential and cutoffs of a given pseudo family and structure. +def _load_pseudo_family( + label: str, pseudo_set: ty.Tuple[type, ...] = (PseudoPotentialFamily,) +): + """Return the stored pseudopotential family with the given label. - :param pseudo_family: [description] - :param structure: [description] - :raises ValueError: [description] - :raises ValueError: [description] - :return: [description] + :param pseudo_set: the family classes to search. The default accepts every + ``aiida-pseudo`` family; pass a narrower set when the caller needs a + capability only some families have. + :raises ValueError: no family in ``pseudo_set`` carries that label. """ try: - pseudo_set = (PseudoDojoFamily, SsspFamily, CutoffsPseudoPotentialFamily) - pseudo_family = ( - orm.QueryBuilder() - .append(pseudo_set, filters={"label": pseudo_family}) - .one()[0] - ) + return orm.QueryBuilder().append(pseudo_set, filters={"label": label}).one()[0] except exceptions.NotExistent as exception: raise ValueError( - f"required pseudo family `{pseudo_family}` is not installed. Please use `aiida-pseudo install` to" + f"required pseudo family `{label}` is not installed. Please use `aiida-pseudo install` to" "install it." ) from exception + +def get_pseudos( + pseudo_family: str, structure: orm.StructureData +) -> ty.Mapping[str, PseudoPotentialData]: + """Get the pseudo potential of each kind of a structure, from a pseudo family. + + Any ``aiida-pseudo`` family serves, including one built by the user with + ``aiida-pseudo install family``: no recommended cutoffs are needed. + + :param pseudo_family: label of the family to take the pseudos from. + :param structure: the structure whose kinds are to be covered. + :raises ValueError: the family is not installed, or does not cover every + kind of the structure. + """ + return _load_pseudo_family(pseudo_family).get_pseudos(structure=structure) + + +def get_pseudo_and_cutoff( + pseudo_family: str, structure: orm.StructureData +) -> ty.Tuple[ty.Mapping[str, PseudoPotentialData], float, float]: + """Get the pseudo potentials and the recommended cutoffs of a pseudo family. + + Only families that recommend cutoffs qualify, i.e. SSSP, PseudoDojo and any + family with a cutoff stringency defined. Use :func:`get_pseudos` when the + cutoffs are not needed. + + :param pseudo_family: label of the family to take the pseudos from. + :param structure: the structure whose kinds are to be covered. + :raises ValueError: the family is not installed among those that recommend + cutoffs. + :raises ValueError: the family recommends no cutoffs for this structure. + :return: the pseudos per kind, the wave-function cutoff and the charge + density cutoff, both in Ry. + """ + pseudo_family = _load_pseudo_family( + pseudo_family, (PseudoDojoFamily, SsspFamily, CutoffsPseudoPotentialFamily) + ) + try: cutoff_wfc, cutoff_rho = pseudo_family.get_recommended_cutoffs( structure=structure, unit="Ry" diff --git a/src/aiida_wannier90_workflows/workflows/base/wannier90.py b/src/aiida_wannier90_workflows/workflows/base/wannier90.py index 7b3db456..26e32ca7 100644 --- a/src/aiida_wannier90_workflows/workflows/base/wannier90.py +++ b/src/aiida_wannier90_workflows/workflows/base/wannier90.py @@ -271,8 +271,8 @@ def get_builder_from_protocol( from aiida_wannier90_workflows.utils.pseudo import ( get_number_of_projections, get_number_of_projections_ext, - get_pseudo_and_cutoff, get_pseudo_orbitals, + get_pseudos, get_semicore_list, get_semicore_list_ext, get_wannier_number_of_bands, @@ -323,7 +323,7 @@ def get_builder_from_protocol( if pseudo_family is None: pseudo_family = meta_parameters["pseudo_family"] - pseudos, _, _ = get_pseudo_and_cutoff(pseudo_family, structure) + pseudos = get_pseudos(pseudo_family, structure) if external_projectors is not None: if projection_type in [ diff --git a/tests/conftest.py b/tests/conftest.py index 16fc499e..f5fec392 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -152,6 +152,41 @@ def pseudos(aiida_profile, generate_upf_data, generate_upf_data_soc): return sssp, dojo +@pytest.fixture(scope="session") +def cutoffs_family_without_stringency( + pseudos, generate_upf_data +): # pylint: disable=unused-argument + """Create a cutoffs family with no stringency, so it recommends no cutoffs. + + The shape a user gets after grouping pseudos of their own into a + ``CutoffsPseudoPotentialFamily`` without calling ``set_cutoffs``. Ordered + after ``pseudos``, which resets the profile. + """ + from aiida.plugins import GroupFactory + + family = GroupFactory("pseudo.family.cutoffs")(label="NoStringency/1.0") + family.store() + family.add_nodes([generate_upf_data("Si")]) + + return family + + +@pytest.fixture(scope="session") +def plain_pseudo_family(pseudos, generate_upf_data): # pylint: disable=unused-argument + """Create a plain family, the shape ``aiida-pseudo install family`` produces. + + Cutoffs are not even representable on this family class. Ordered after + ``pseudos``, which resets the profile. + """ + from aiida.plugins import GroupFactory + + family = GroupFactory("pseudo.family")(label="MyPseudos/local") + family.store() + family.add_nodes([generate_upf_data("Si")]) + + return family + + @pytest.fixture(scope="session") def generate_upf_data(filepath_fixtures): """Return a `UpfData` instance for the given element a file for which should exist in `tests/fixtures/pseudos`.""" diff --git a/tests/utils/test_pseudo.py b/tests/utils/test_pseudo.py new file mode 100644 index 00000000..8d3b6321 --- /dev/null +++ b/tests/utils/test_pseudo.py @@ -0,0 +1,66 @@ +"""Tests for the :py:mod:`~aiida_wannier90_workflows.utils.pseudo` module.""" + +import pytest + +from aiida_wannier90_workflows.utils.pseudo import get_pseudo_and_cutoff, get_pseudos + + +def test_get_pseudos_cutoffs_family_without_stringency( + cutoffs_family_without_stringency, generate_structure +): + """A family that recommends no cutoffs still provides its pseudos.""" + pseudos = get_pseudos( + cutoffs_family_without_stringency.label, generate_structure("Si") + ) + + assert sorted(pseudos) == ["Si"] + + +def test_get_pseudos_plain_family(plain_pseudo_family, generate_structure): + """A family that cannot carry cutoffs at all still provides its pseudos.""" + pseudos = get_pseudos(plain_pseudo_family.label, generate_structure("Si")) + + assert sorted(pseudos) == ["Si"] + + +def test_get_pseudos_family_with_cutoffs(pseudos, generate_structure): + """A family that does recommend cutoffs is served by the same entry point.""" + sssp, _ = pseudos + + assert sorted(get_pseudos(sssp.label, generate_structure("Si"))) == ["Si"] + + +def test_get_pseudos_family_not_installed(generate_structure): + """An unknown label names itself in the error.""" + with pytest.raises(ValueError, match="`NotAFamily/1.0` is not installed"): + get_pseudos("NotAFamily/1.0", generate_structure("Si")) + + +def test_get_pseudo_and_cutoff_returns_cutoffs(pseudos, generate_structure): + """The cutoffs of a family that has them are unchanged.""" + sssp, _ = pseudos + + found, cutoff_wfc, cutoff_rho = get_pseudo_and_cutoff( + sssp.label, generate_structure("Si") + ) + + assert sorted(found) == ["Si"] + assert (cutoff_wfc, cutoff_rho) == (30.0, 240.0) + + +def test_get_pseudo_and_cutoff_requires_cutoffs( + cutoffs_family_without_stringency, generate_structure +): + """Asking for cutoffs a family does not have still raises.""" + with pytest.raises(ValueError, match="failed to obtain recommended cutoffs"): + get_pseudo_and_cutoff( + cutoffs_family_without_stringency.label, generate_structure("Si") + ) + + +def test_get_pseudo_and_cutoff_rejects_plain_family( + plain_pseudo_family, generate_structure +): + """A family that cannot carry cutoffs is not a candidate for this function.""" + with pytest.raises(ValueError, match="`MyPseudos/local` is not installed"): + get_pseudo_and_cutoff(plain_pseudo_family.label, generate_structure("Si")) diff --git a/tests/workflows/protocols/base/test_wannier90.py b/tests/workflows/protocols/base/test_wannier90.py index 5b19bbe8..9c7df7d8 100644 --- a/tests/workflows/protocols/base/test_wannier90.py +++ b/tests/workflows/protocols/base/test_wannier90.py @@ -210,3 +210,33 @@ def test_metadata_overrides( ) data_regression.check(serialize_builder(builder)) + + +@pytest.mark.parametrize( + "family_fixture", ("cutoffs_family_without_stringency", "plain_pseudo_family") +) +def test_pseudo_family_without_cutoffs( + fixture_code, generate_structure, request, family_fixture +): + """A pseudo family that recommends no cutoffs builds the same inputs. + + This builder counts bands and projections from the pseudos and never uses + the cutoffs, so a family that has none must serve it as well as one that + does. + """ + code = fixture_code("wannier90.wannier90") + structure = generate_structure("Si") + family = request.getfixturevalue(family_fixture) + + builder = Wannier90BaseWorkChain.get_builder_from_protocol( + code, structure=structure, pseudo_family=family.label + ) + reference = Wannier90BaseWorkChain.get_builder_from_protocol( + code, structure=structure + ) + + assert isinstance(builder, ProcessBuilder) + assert ( + builder.wannier90.parameters.get_dict() + == reference.wannier90.parameters.get_dict() + )