Get pseudos without requiring recommended cutoffs

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"

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 [

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`."""

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"))

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()
+    )