🐛 Protocols: Fix behaviour for pseudos overrides

There are currently two bugs in the `PwBaseWorkChain.get_builder_from_protocol()`
method when the users overrides the `pw.pseudos` input:

1. For magnetic calculations, the `starting_magnetization` is still constructed based on
   the number of valence electrons (`z_valence`) from the `pseudo_family`.
2. If the user overrides the `pseudos` input but doesn't override the cutoffs
   (`SYSTEM.ecutrho`, `SYSTEM.ecutwfc`), these are still (silently) taken from the
   `pseudo_family`.

Here we fix these two issues by

1. Passing the `z_valences` of the provides `pseudos` override to the
   `get_magnetization` function.
2. Strictly checking the overrides: they should both (i) provide pseudo
   potentials for _all_ the kinds in the input `structure` and (ii) provide both
   energy cutoffs (`SYSTEM.ecutrho`, `SYSTEM.ecutwfc`) in case the `pseudos` are
   overridden.

Moreover, if the user completely specifies the pseudos and cutoffs through the
overrides, it should be fine in case `pseudo_family` is None or False. This is important
in the context of adding custom pseudo potential functionality in the QEapp, see:

aiidalab/aiidalab-qe#1348

Co-authored-by: Edan Bainglass <edan.bainglass@gmail.com>

Author: 2 people

src/aiida_quantumespresso/workflows/pw/base.py

@@ -156,31 +156,47 @@ def get_builder_from_protocol(
 
         natoms = len(structure.sites)
 
-        try:
-            pseudo_set = (PseudoDojoFamily, SsspFamily, CutoffsPseudoPotentialFamily)
-            pseudo_family = orm.QueryBuilder().append(pseudo_set, filters={'label': pseudo_family}).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:
-            cutoff_wfc, cutoff_rho = pseudo_family.get_recommended_cutoffs(structure=structure, unit='Ry')
-            pseudos = pseudo_family.get_pseudos(structure=structure)
-        except ValueError as exception:
-            raise ValueError(
-                f'failed to obtain recommended cutoffs for pseudo family `{pseudo_family}`: {exception}'
-            ) from exception
-
         # Update the parameters based on the protocol inputs
         parameters = inputs['pw']['parameters']
+
+        if overrides and 'pseudos' in overrides['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')):
+                raise ValueError(
+                    'When overriding the pseudo potentials, both `ecutwfc` and `ecutrho` cutoffs should be '
+                    f'provided in the `overrides`: {overrides}'
+                )
+
+        else:
+            try:
+                pseudo_set = (PseudoDojoFamily, SsspFamily, CutoffsPseudoPotentialFamily)
+                pseudo_family = orm.QueryBuilder().append(pseudo_set, filters={'label': pseudo_family}).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:
+                raise ValueError(
+                    f'failed to obtain recommended cutoffs for pseudo family `{pseudo_family}`: {exception}'
+                ) from exception
+
         parameters['CONTROL']['etot_conv_thr'] = natoms * meta_parameters['etot_conv_thr_per_atom']
         parameters['ELECTRONS']['conv_thr'] = natoms * meta_parameters['conv_thr_per_atom']
-        parameters['SYSTEM']['ecutwfc'] = cutoff_wfc
-        parameters['SYSTEM']['ecutrho'] = cutoff_rho
 
-        #If the structure is 2D periodic in the x-y plane, we set assume_isolate to `2D`
+        # If the structure is 2D periodic in the x-y plane, we set assume_isolate to `2D`
         if structure.pbc == (True, True, False):
             parameters['SYSTEM']['assume_isolated'] = '2D'
 
@@ -191,9 +207,9 @@ def get_builder_from_protocol(
 
         magnetization = get_magnetization(
             structure=structure,
-            z_valences={kind.name: pseudo_family.get_pseudo(element=kind.symbol).z_valence for kind in structure.kinds},
+            z_valences={kind.name: pseudos[kind.name].z_valence for kind in structure.kinds},
             initial_magnetic_moments=initial_magnetic_moments,
-            spin_type=spin_type
+            spin_type=spin_type,
         )
         if spin_type is SpinType.COLLINEAR:
             parameters['SYSTEM']['starting_magnetization'] = magnetization['starting_magnetization']
@@ -217,9 +233,6 @@ def get_builder_from_protocol(
             if parameters.get('SYSTEM', {}).get('tot_magnetization') is not None:
                 parameters.setdefault('SYSTEM', {}).pop('starting_magnetization', None)
 
-            pseudos_overrides = overrides.get('pw', {}).get('pseudos', {})
-            pseudos = recursive_merge(pseudos, pseudos_overrides)
-
         metadata = inputs['pw']['metadata']
 
         if options:

tests/workflows/protocols/pw/test_base.py

@@ -203,19 +203,100 @@ def test_parallelization_overrides(fixture_code, generate_structure):
 
 def test_pseudos_overrides(fixture_code, generate_structure, generate_upf_data):
     """Test specifying ``pw.pseudos`` ``overrides`` for the ``get_builder_from_protocol()`` method."""
-    code = fixture_code('quantumespresso.pw')
-    structure = generate_structure('silicon')
-    silicon_pseudo = generate_upf_data('Si')
-
-    overrides = {'pw': {'pseudos': {'Si': silicon_pseudo}}}
-    builder = PwBaseWorkChain.get_builder_from_protocol(
-        code,
-        structure,
-        overrides=overrides,
-    )
-    pseudos = builder.pw.pseudos  # pylint: disable=no-member
+    structure = generate_structure('cobalt-prim')
+    upf_data = generate_upf_data('Co', z_valence=5.0)
+
+    for overrides in (
+        {
+            'pw': {
+                'pseudos': {
+                    'Co': upf_data,
+                },
+                'parameters': {
+                    'SYSTEM': {
+                        'ecutrho': 123,
+                        'ecutwfc': 456,
+                    }
+                }
+            }
+        },
+        {
+            'pseudo_family': None,
+            'pw': {
+                'pseudos': {
+                    'Co': upf_data,
+                },
+                'parameters': {
+                    'SYSTEM': {
+                        'ecutrho': 123,
+                        'ecutwfc': 456,
+                    }
+                }
+            }
+        },
+    ):
+        builder = PwBaseWorkChain.get_builder_from_protocol(
+            code=fixture_code('quantumespresso.pw'),
+            structure=structure,
+            overrides=overrides,
+            spin_type=SpinType.COLLINEAR
+        )
+        pseudos = builder.pw.pseudos  # pylint: disable=no-member
+        parameters = builder.pw.parameters.get_dict()
+
+        assert pseudos['Co'] == upf_data
+        assert parameters['SYSTEM']['ecutrho'] == 123
+        assert parameters['SYSTEM']['ecutwfc'] == 456
+        assert parameters['SYSTEM']['starting_magnetization']['Co'] == 1.0
+
+
+def test_pseudos_overrides_raises(fixture_code, generate_structure, generate_upf_data):
+    """Test specifying ``pw.pseudos`` ``overrides`` for the ``get_builder_from_protocol()`` method.
+
+    This test checks that the `get_builder_from_protocol` raises when the user overrides the `pseudos` but does not
+    override the cutoffs.
+    """
+    structure = generate_structure('cobalt-prim')
+    upf_data = generate_upf_data('Co')
+
+    for params in (
+        ({
+            'pw': {
+                'pseudos': {
+                    'Co': upf_data
+                }
+            }
+        }, ValueError, 'both `ecutwfc` and `ecutrho` cutoffs should be'),
+        ({
+            'pw': {
+                'pseudos': {
+                    'Si': upf_data
+                }
+            }
+        }, ValueError, '`pseudos` override needs one value'),
+        ({
+            'pseudo_family': None,
+            'pw': {
+                'pseudos': {
+                    'Co': upf_data
+                }
+            }
+        }, ValueError, 'both `ecutwfc` and `ecutrho` cutoffs should be'),
+        ({
+            'pseudo_family': None,
+            'pw': {
+                'pseudos': {
+                    'Si': upf_data
+                }
+            }
+        }, ValueError, '`pseudos` override needs one value'),
+    ):
+        overrides, error_type, error_msg = params
 
-    assert pseudos['Si'] == silicon_pseudo
+        with pytest.raises(error_type, match=error_msg):
+            PwBaseWorkChain.get_builder_from_protocol(
+                code=fixture_code('quantumespresso.pw'), structure=structure, overrides=overrides
+            )
 
 
 def test_pseudos_family_structure_fail(fixture_code, generate_structure):