👌 Protocols: support pseudo families without recommended cutoffs

src/aiida_quantumespresso/workflows/pw/base.py

@@ -13,6 +13,7 @@
     while_,
 )
 from aiida.plugins import CalculationFactory, GroupFactory
+from aiida_pseudo.groups.mixins import RecommendedCutoffMixin
 
 from aiida_quantumespresso.calculations.functions.create_kpoints_from_distance import (
     create_kpoints_from_distance,
@@ -25,9 +26,7 @@
 from ..protocols.utils import ProtocolMixin
 
 PwCalculation = CalculationFactory('quantumespresso.pw')
-SsspFamily = GroupFactory('pseudo.family.sssp')
-PseudoDojoFamily = GroupFactory('pseudo.family.pseudo_dojo')
-CutoffsPseudoPotentialFamily = GroupFactory('pseudo.family.cutoffs')
+PseudoPotentialFamily = GroupFactory('pseudo.family')
 
 
 class PwBaseWorkChain(ProtocolMixin, BaseRestartWorkChain):
@@ -216,44 +215,59 @@ def get_builder_from_protocol(
         # Update the parameters based on the protocol inputs
         parameters = inputs['pw']['parameters']
 
+        system_overrides = (overrides or {}).get('pw', {}).get('parameters', {}).get('SYSTEM', {})
+        cutoffs_in_overrides = all(key in system_overrides for key in ('ecutwfc', 'ecutrho'))
+
         if overrides and 'pseudos' in overrides.get('pw', {}):
             pseudos = overrides['pw']['pseudos']
 
             if sorted(pseudos.keys()) != sorted(structure.get_kind_names()):
                 raise ValueError(f'`pseudos` override needs one value for each of the {len(structure.kinds)} kinds.')
 
-            system_overrides = overrides['pw'].get('parameters', {}).get('SYSTEM', {})
-
-            if not all(key in system_overrides for key in ('ecutwfc', 'ecutrho')):
+            if not cutoffs_in_overrides:
                 raise ValueError(
                     'When overriding the pseudo potentials, both `ecutwfc` and `ecutrho` cutoffs should be '
                     f'provided in the `overrides`: {overrides}'
                 )
 
         else:
+            query = orm.QueryBuilder().append(PseudoPotentialFamily, filters={'label': pseudo_family})
+
             try:
-                pseudo_set = (
-                    PseudoDojoFamily,
-                    SsspFamily,
-                    CutoffsPseudoPotentialFamily,
-                )
-                pseudo_family = orm.QueryBuilder().append(pseudo_set, filters={'label': pseudo_family}).one()[0]
+                pseudo_family = query.one()[0]
             except exceptions.NotExistent as exception:
                 raise ValueError(
                     f'required pseudo family `{pseudo_family}` is not installed. Please use `aiida-pseudo install` to'
                     'install it.'
                 ) from exception
-
-            try:
-                parameters['SYSTEM']['ecutwfc'], parameters['SYSTEM']['ecutrho'] = (
-                    pseudo_family.get_recommended_cutoffs(structure=structure, unit='Ry')
-                )
-                pseudos = pseudo_family.get_pseudos(structure=structure)
-            except ValueError as exception:
+            except exceptions.MultipleObjectsError as exception:
+                matches = ', '.join(f'`{family}`' for family in query.all(flat=True))
                 raise ValueError(
-                    f'failed to obtain recommended cutoffs for pseudo family `{pseudo_family}`: {exception}'
+                    f'the label `{pseudo_family}` matches more than one installed pseudo family: {matches}. Please '
+                    'delete or relabel all but one, or pass the `pseudos` in the `overrides` instead.'
                 ) from exception
 
+            # Families that do not define recommended cutoffs can still be used, as long as the `overrides` provide
+            # both cutoffs themselves, since these are applied further down and take precedence anyway.
+            if not cutoffs_in_overrides:
+                if not isinstance(pseudo_family, RecommendedCutoffMixin):
+                    raise ValueError(
+                        f'pseudo family `{pseudo_family}` cannot recommend cutoffs. Provide both `ecutwfc` and '
+                        '`ecutrho` in the `overrides`, or use a family that recommends them.'
+                    )
+
+                try:
+                    parameters['SYSTEM']['ecutwfc'], parameters['SYSTEM']['ecutrho'] = (
+                        pseudo_family.get_recommended_cutoffs(structure=structure, unit='Ry')
+                    )
+                except ValueError as exception:
+                    raise ValueError(
+                        f'failed to obtain recommended cutoffs for pseudo family `{pseudo_family}`: {exception} '
+                        'If the family does not define any, specify both `ecutwfc` and `ecutrho` in the `overrides`.'
+                    ) from exception
+
+            pseudos = pseudo_family.get_pseudos(structure=structure)
+
         parameters['CONTROL']['etot_conv_thr'] = natoms * meta_parameters['etot_conv_thr_per_atom']
         parameters['ELECTRONS']['conv_thr'] = natoms * meta_parameters['conv_thr_per_atom']
 

tests/conftest.py

@@ -172,6 +172,57 @@ def pseudo_family(generate_upf_data):
     return family
 
 
+@pytest.fixture(scope='session')
+def generate_silicon_pseudo_family(generate_upf_data):
+    """Return a factory that creates a silicon-only pseudo potential family of a given type and label."""
+    from aiida_pseudo.data.pseudo.upf import UpfData
+
+    def _generate_silicon_pseudo_family(family_type, label):
+        upf = generate_upf_data('Si')
+
+        with tempfile.TemporaryDirectory() as directory:
+            dirpath = pathlib.Path(directory)
+
+            with open(dirpath / 'Si.upf', 'w+b') as handle, upf.open(mode='rb') as source:
+                handle.write(source.read())
+                handle.flush()
+
+            return family_type.create_from_folder(dirpath, label, pseudo_type=UpfData)
+
+    return _generate_silicon_pseudo_family
+
+
+@pytest.fixture(scope='session')
+def pseudo_families_without_cutoffs(generate_silicon_pseudo_family):
+    """Create the silicon pseudo potential families that recommend no cutoffs, keyed on their class name.
+
+    A ``CutoffsPseudoPotentialFamily`` can recommend cutoffs but has none set here. A ``PseudoPotentialFamily``, the
+    class that ``aiida-pseudo install family`` installs by default, cannot recommend them at all.
+    """
+    from aiida_pseudo.groups.family import CutoffsPseudoPotentialFamily, PseudoPotentialFamily
+
+    return {
+        family_type.__name__: generate_silicon_pseudo_family(family_type, f'custom/{family_type.__name__}')
+        for family_type in (PseudoPotentialFamily, CutoffsPseudoPotentialFamily)
+    }
+
+
+@pytest.fixture(scope='session')
+def pseudo_families_shared_label(generate_silicon_pseudo_family):
+    """Create two silicon pseudo potential families of different types under one label, and return that label.
+
+    Group labels are unique per type string, so families of different types can share one.
+    """
+    from aiida_pseudo.groups.family import CutoffsPseudoPotentialFamily, PseudoPotentialFamily
+
+    label = 'custom/shared-label'
+
+    for family_type in (PseudoPotentialFamily, CutoffsPseudoPotentialFamily):
+        generate_silicon_pseudo_family(family_type, label)
+
+    return label
+
+
 @pytest.fixture
 def generate_calc_job():
     """Fixture to construct a new `CalcJob` instance and call `prepare_for_submission` for testing `CalcJob` classes.

tests/workflows/protocols/pw/test_base.py

@@ -182,6 +182,21 @@ def test_overrides_pseudo_family(fixture_code, generate_structure):
     assert parameters['SYSTEM']['ecutrho'] == 400.0
 
 
+def test_overrides_cutoffs(fixture_code, generate_structure):
+    """Test cutoffs in the ``overrides`` take precedence over those recommended by the ``pseudo_family``."""
+    code = fixture_code('quantumespresso.pw')
+    structure = generate_structure('silicon')
+
+    # The default family recommends 30 Ry and 240 Ry, see ``test_default``
+    overrides = {'pw': {'parameters': {'SYSTEM': {'ecutwfc': 45.0, 'ecutrho': 180.0}}}}
+
+    builder = PwBaseWorkChain.get_builder_from_protocol(code, structure, overrides=overrides)
+    parameters = builder.pw.parameters.get_dict()
+    assert parameters['SYSTEM']['ecutwfc'] == 45.0
+    assert parameters['SYSTEM']['ecutrho'] == 180.0
+    assert 'Si' in builder.pw.pseudos
+
+
 def test_magnetization_overrides(fixture_code, generate_structure):
     """Test magnetization ``overrides`` for the ``PwBaseWorkChain.get_builder_from_protocol`` method."""
     code = fixture_code('quantumespresso.pw')
@@ -418,6 +433,92 @@ def test_pseudos_family_structure_fail(fixture_code, generate_structure):
             structure,
         )
 
+    # If the cutoffs are specified in the ``overrides``, the family is only asked for the pseudo potentials
+    with pytest.raises(ValueError, match=r'does not contain pseudo for element `U`') as exception:
+        PwBaseWorkChain.get_builder_from_protocol(
+            code,
+            structure,
+            overrides={'pw': {'parameters': {'SYSTEM': {'ecutwfc': 30.0, 'ecutrho': 240.0}}}},
+        )
+
+    assert 'recommended cutoffs' not in str(exception.value)
+
+
+@pytest.mark.parametrize('family_type', ('PseudoPotentialFamily', 'CutoffsPseudoPotentialFamily'))
+def test_pseudo_family_without_cutoffs(fixture_code, generate_structure, pseudo_families_without_cutoffs, family_type):
+    """Test a ``pseudo_family`` without recommended cutoffs, where the cutoffs are specified in the ``overrides``."""
+    code = fixture_code('quantumespresso.pw')
+    structure = generate_structure('silicon')
+    family = pseudo_families_without_cutoffs[family_type]
+
+    builder = PwBaseWorkChain.get_builder_from_protocol(
+        code,
+        structure,
+        overrides={
+            'pseudo_family': family.label,
+            'pw': {'parameters': {'SYSTEM': {'ecutwfc': 30.0, 'ecutrho': 240.0}}},
+        },
+    )
+    parameters = builder.pw.parameters.get_dict()
+
+    assert builder.pw.pseudos['Si'].uuid == family.get_pseudo(element='Si').uuid
+    assert parameters['SYSTEM']['ecutwfc'] == 30.0
+    assert parameters['SYSTEM']['ecutrho'] == 240.0
+
+
+@pytest.mark.parametrize('family_type', ('PseudoPotentialFamily', 'CutoffsPseudoPotentialFamily'))
+@pytest.mark.parametrize('system_overrides', ({}, {'ecutwfc': 30.0}, {'ecutrho': 240.0}))
+def test_pseudo_family_without_cutoffs_fail(
+    fixture_code, generate_structure, pseudo_families_without_cutoffs, system_overrides, family_type
+):
+    """Test a ``pseudo_family`` without recommended cutoffs fails if the ``overrides`` do not specify both cutoffs."""
+    code = fixture_code('quantumespresso.pw')
+    structure = generate_structure('silicon')
+
+    with pytest.raises(ValueError, match=r'both `ecutwfc` and `ecutrho` in the `overrides`'):
+        PwBaseWorkChain.get_builder_from_protocol(
+            code,
+            structure,
+            overrides={
+                'pseudo_family': pseudo_families_without_cutoffs[family_type].label,
+                'pw': {'parameters': {'SYSTEM': system_overrides}},
+            },
+        )
+
+
+def test_pseudo_family_shared_label_fail(fixture_code, generate_structure, pseudo_families_shared_label):
+    """Test a ``pseudo_family`` label that is shared by families of different types is refused."""
+    code = fixture_code('quantumespresso.pw')
+    structure = generate_structure('silicon')
+
+    with pytest.raises(ValueError, match=r'matches more than one installed pseudo family') as exception:
+        PwBaseWorkChain.get_builder_from_protocol(
+            code,
+            structure,
+            overrides={
+                'pseudo_family': pseudo_families_shared_label,
+                'pw': {'parameters': {'SYSTEM': {'ecutwfc': 30.0, 'ecutrho': 240.0}}},
+            },
+        )
+
+    for family_type in ('PseudoPotentialFamily', 'CutoffsPseudoPotentialFamily'):
+        assert f'`{family_type}<{pseudo_families_shared_label}>`' in str(exception.value)
+
+
+@pytest.mark.usefixtures('pseudo_families_shared_label')
+def test_pseudo_family_unique_label(fixture_code, generate_structure):
+    """Test a uniquely labelled ``pseudo_family`` still resolves while a shared label is installed."""
+    code = fixture_code('quantumespresso.pw')
+    structure = generate_structure('silicon')
+
+    builder = PwBaseWorkChain.get_builder_from_protocol(
+        code, structure, overrides={'pseudo_family': 'SSSP/1.3/PBEsol/efficiency'}
+    )
+    parameters = builder.pw.parameters.get_dict()
+
+    assert parameters['SYSTEM']['ecutwfc'] == 30.0
+    assert parameters['SYSTEM']['ecutrho'] == 240.0
+
 
 def test_options(fixture_code, generate_structure):
     """Test specifying ``options`` for the ``get_builder_from_protocol()`` method."""