🐛 get_builder_from_protocol: Fix bugs when overriding pseudopotentials

src/aiida_quantumespresso/workflows/protocols/utils.py

@@ -134,7 +134,7 @@ def get_magnetization_parameters() -> dict:
 
 def get_magnetization(
     structure: StructureData,
-    pseudo_family: PseudoPotentialFamily,
+    z_valences: dict,
     initial_magnetic_moments: Optional[dict] = None,
     spin_type: SpinType = SpinType.COLLINEAR
 ) -> dict:
@@ -149,7 +149,8 @@ def get_magnetization(
     In case the `spin_type` is set to `SpinType.COLLINEAR`, the values for `angle1` and `angle2` will be set to `None`.
 
     :param structure: the structure.
-    :param pseudo_family: pseudopotential family.
+    :param z_valences: dictionary mapping each kind in the structure to the number of valence electrons in the pseudo
+        potential.
     :param initial_magnetic_moments: dictionary mapping each kind in the structure to its magnetic moment.
     :param spin_type: the `SpinType` of the calculation.
     :returns: dictionary of the magnetization.
@@ -159,6 +160,8 @@ def get_magnetization(
         'angle1': {} if spin_type in [SpinType.NON_COLLINEAR, SpinType.SPIN_ORBIT] else None,
         'angle2': {} if spin_type in [SpinType.NON_COLLINEAR, SpinType.SPIN_ORBIT] else None,
     }
+    if sorted(z_valences.keys()) != sorted(structure.get_kind_names()):
+        raise ValueError(f'`z_valences` needs one value for each of the {len(structure.kinds)} kinds.')
 
     if initial_magnetic_moments is not None:
 
@@ -172,8 +175,8 @@ def get_magnetization(
             magmom = initial_magnetic_moments[kind.name]
 
             if isinstance(magmom, (int, float)):
-                magnetization['starting_magnetization'][
-                    kind.name] = magmom / pseudo_family.get_pseudo(element=kind.symbol).z_valence
+                scaled_magmom = magmom / z_valences[kind.name]
+                magnetization['starting_magnetization'][kind.name] = scaled_magmom
 
                 if spin_type in [SpinType.NON_COLLINEAR, SpinType.SPIN_ORBIT]:
                     magnetization['angle1'][kind.name] = 0.0
@@ -187,8 +190,8 @@ def get_magnetization(
                         f'moment of kind `{kind.name}`.'
                     )
 
-                magnetization['starting_magnetization'][
-                    kind.name] = magmom[0] / pseudo_family.get_pseudo(element=kind.symbol).z_valence
+                scaled_magmom = magmom[0] / z_valences[kind.name]
+                magnetization['starting_magnetization'][kind.name] = scaled_magmom
                 magnetization['angle1'][kind.name] = magmom[1]
                 magnetization['angle2'][kind.name] = magmom[2]
             else:
@@ -212,8 +215,8 @@ def get_magnetization(
 
             magmom = kind.get_magmom_coord()
 
-            magnetization['starting_magnetization'][
-                kind.name] = magmom[0] / pseudo_family.get_pseudo(element=kind.symbol).z_valence
+            scaled_magmom = magmom[0] / z_valences[kind.name]
+            magnetization['starting_magnetization'][kind.name] = scaled_magmom
 
             if spin_type in [SpinType.NON_COLLINEAR, SpinType.SPIN_ORBIT]:
                 magnetization['angle1'][kind.name] = magmom[1]
@@ -231,7 +234,7 @@ def get_magnetization(
 
         magnetization['starting_magnetization'][kind.name] = (
             magnetic_parameters['default_magnetization'] if magnetic_moment == 0 else magnetic_moment /
-            pseudo_family.get_pseudo(element=kind.symbol).z_valence
+            z_valences[kind.name]
         )
         if spin_type in [SpinType.NON_COLLINEAR, SpinType.SPIN_ORBIT]:
             magnetization['angle1'][kind.name] = 0.0

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,
-            pseudo_family=pseudo_family,
+            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/conftest.py

@@ -326,10 +326,11 @@ def _generate_calc_job_node(
 def generate_upf_data():
     """Return a `UpfData` instance for the given element a file for which should exist in `tests/fixtures/pseudos`."""
 
-    def _generate_upf_data(element):
-        """Return `UpfData` node."""
-        from aiida_pseudo.data.pseudo import UpfData
-        content = f'<UPF version="2.0.1"><PP_HEADER\nelement="{element}"\nz_valence="4.0"\n/></UPF>\n'
+    from aiida_pseudo.data.pseudo import UpfData
+
+    def _generate_upf_data(element: str, z_valence: float = 4.0) -> UpfData:
+        """Return a `UpfData` node."""
+        content = f'<UPF version="2.0.1"><PP_HEADER\nelement="{element}"\nz_valence="{z_valence}"\n/></UPF>\n'
         stream = io.BytesIO(content.encode('utf-8'))
         return UpfData(stream, filename=f'{element}.upf')
 

tests/workflows/protocols/pw/test_base.py

@@ -218,19 +218,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):

tests/workflows/protocols/test_utils.py

@@ -116,7 +116,6 @@ def test_recursive_merge():
     ),
 )
 def test_get_magnetization(
-    pseudo_family,
     generate_structure,
     structure_id,
     initial_magnetic_moments,
@@ -125,31 +124,43 @@ def test_get_magnetization(
 ):
     """Test the `get_magnetization` function."""
     from aiida_quantumespresso.workflows.protocols.utils import get_magnetization
+    structure = generate_structure(structure_id)
+    z_valences = {kind: 4.0 for kind in structure.get_kind_names()}
 
-    magnetization = get_magnetization(
-        generate_structure(structure_id), pseudo_family, initial_magnetic_moments, spin_type
-    )
+    magnetization = get_magnetization(structure, z_valences, initial_magnetic_moments, spin_type)
 
     assert magnetization == expected_magnetization
 
 
 @pytest.mark.parametrize(
-    'structure_id,initial_magnetic_moments,spin_type,expected_error,error_message',
+    'structure_id,z_valences,initial_magnetic_moments,spin_type,expected_error,error_message',
     (
-        ('silicon', {}, SpinType.COLLINEAR, ValueError, '`initial_magnetic_moments` needs one value for each of'),
+        ('silicon', {}, {
+            'Si': 1.0
+        }, SpinType.COLLINEAR, ValueError, '`z_valences` needs one value for each of'),
+        (
+            'silicon', {
+                'Si': 4.0
+            }, {}, SpinType.COLLINEAR, ValueError, '`initial_magnetic_moments` needs one value for each of'
+        ),
         ('silicon', {
+            'Si': 4.0
+        }, {
             'Si': (1, 2, 3)
         }, SpinType.COLLINEAR, TypeError, 'Spin type is set to '),
-        ('silicon', {
-            'Si': 'zero'
-        }, SpinType.COLLINEAR, TypeError, 'Unrecognised type for magnetic moment'),
+        (
+            'silicon', {
+                'Si': 4.0
+            }, {
+                'Si': 'zero'
+            }, SpinType.COLLINEAR, TypeError, 'Unrecognised type for magnetic moment'
+        ),
     ),
 )
 def test_get_magnetization_failure(
-    pseudo_family, generate_structure, structure_id, initial_magnetic_moments, spin_type, expected_error, error_message
+    generate_structure, structure_id, z_valences, initial_magnetic_moments, spin_type, expected_error, error_message
 ):
     """Test the `get_magnetization` function."""
     from aiida_quantumespresso.workflows.protocols.utils import get_magnetization
-
     with pytest.raises(expected_error, match=error_message):
-        get_magnetization(generate_structure(structure_id), pseudo_family, initial_magnetic_moments, spin_type)
+        get_magnetization(generate_structure(structure_id), z_valences, initial_magnetic_moments, spin_type)