👌 get_magnetization: Replace pseudo_family input by z_valences

The current `get_magnetization` utility function has an input `pseudo_family` that
expects an `aiida-pseudo` `PseudoPotentialFamily` to extract the number of valence
electrons in the pseudo potential. However, when a user wants to use custom pseudo
potentials e.g. through the `overrides` of the `get_builder_from_protocol()` method,
it would still be beneficial to use the `get_magnetization` function to determine the
correct Quantum ESPRESSO inputs.

Here we make the function more generic by replacing the `pseudo_family` input by
`z_valences`, a simple dictionary that maps each kind name to the corresponding number
of valence electrons in the pseudo potential.

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

Author: 2 people

src/aiida_quantumespresso/workflows/protocols/utils.py

@@ -119,7 +119,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:
@@ -134,7 +134,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.
@@ -144,6 +145,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:
 
@@ -157,8 +160,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
@@ -172,8 +175,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:
@@ -197,8 +200,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]
@@ -216,7 +219,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

@@ -191,7 +191,7 @@ def get_builder_from_protocol(
 
         magnetization = get_magnetization(
             structure=structure,
-            pseudo_family=pseudo_family,
+            z_valences={kind.name: pseudo_family.get_pseudo(element=kind.symbol).z_valence for kind in structure.kinds},
             initial_magnetic_moments=initial_magnetic_moments,
             spin_type=spin_type
         )

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)