Adding support for custom protocols

docs/source/workflows/base/relax/index.rst

@@ -154,7 +154,7 @@ Only ``structure`` and ``engines`` can be specified as a positional argument, al
 
 * ``spin_type``. (Type: a python string).
   An optional string to specify the spin degree of freedom for the calculation.
-  It accepts the values ‘none’ or ‘collinear’. These will be extended in the future to include, for instance, non-collinear magnetism and spin-orbit coupling.
+  It accepts the values ‘none’, ‘collinear’, ‘non_collinear’ or ‘spin_orbit’.
   The default is to run the calculation without spin polarization.
   To explore the supported spin types for each implementation an inspection method is available:
 
@@ -163,11 +163,15 @@ Only ``structure`` and ``engines`` can be specified as a positional argument, al
     input_generator.get_spin_types()
 
 
-* ``magnetization_per_site``. (Type: Python None or a Python list of floats).
+* ``magnetization_per_site``. (Type: Python None or a Python list of floats/list/tuple/sequence).
   An input devoted to the initial magnetization specifications.
   It accepts a list where each entry refers to an atomic site in the structure.
-  The quantity is passed as the spin polarization in units of electrons, meaning the difference between spin up and spin down electrons for the site.
+  For collinear calculations each entry can be a single float.
+  The quantity is then passed as the spin polarization in units of electrons, meaning the difference between spin up and spin down electrons for the site.
   This also corresponds to the magnetization of the site in Bohr magnetons (μB).
+  For non-collinear spin calculations each entry can either be a single float or a list/tuple/sequence of 3 floats.
+  Passing a single float ``value`` is equivalent to passing a sequence ``(0., 0., value)``.
+  The quantity is then passed as a Cartesian magnetization-vector (x, y, z) in units of Bohr magnetons (μB).
   The default for this input is the Python value None and, in case of calculations with spin, the None value signals that the implementation should automatically decide an appropriate default initial magnetization.
   The implementation of such choice is code-dependent and described in the supplementary material of the `S. P. Huber et al., npj Comput. Mater. 7, 136 (2021)`_.
 

src/aiida_common_workflows/utils/__init__.py

@@ -0,0 +1,18 @@
+""" Some common util functions to be used in other modules."""
+
+
+def to_spherical(v):
+    """Convert a vector from Cartesian (x, y, z) to
+    spherical (r, theta, phi) coordinates.
+
+    Note that the angles are returned in degrees.
+    """
+    import numpy as np
+
+    x, y, z = v
+    r = np.linalg.norm(v)
+    rxy = np.hypot(x, y)
+    theta = np.degrees(np.arctan2(rxy, z))
+    phi = np.degrees(np.arctan2(y, x))
+
+    return r, theta, phi

src/aiida_common_workflows/workflows/relax/abinit/generator.py

@@ -51,7 +51,9 @@ def define(cls, spec):
             (ElectronicType.METAL, ElectronicType.INSULATOR, ElectronicType.UNKNOWN)
         )
         spec.inputs['engines']['relax']['code'].valid_type = CodeType('abinit')
-        spec.inputs['protocol'].valid_type = ChoiceType(('fast', 'moderate', 'precise', 'verification-PBE-v1'))
+        spec.inputs['protocol'].valid_type = ChoiceType(
+            ('fast', 'moderate', 'precise', 'verification-PBE-v1', 'custom')
+        )
 
     def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR0912,PLR0915
         """Construct a process builder based on the provided keyword arguments.
@@ -62,6 +64,7 @@ def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR091
         structure = kwargs['structure']
         engines = kwargs['engines']
         protocol = kwargs['protocol']
+        custom_protocol = kwargs.get('custom_protocol', None)
         spin_type = kwargs['spin_type']
         relax_type = kwargs['relax_type']
         electronic_type = kwargs['electronic_type']
@@ -70,7 +73,14 @@ def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR091
         threshold_stress = kwargs.get('threshold_stress', None)
         reference_workchain = kwargs.get('reference_workchain', None)
 
-        protocol = copy.deepcopy(self.get_protocol(protocol))
+        if protocol == 'custom':
+            if custom_protocol is None:
+                raise ValueError(
+                    'the `custom_protocol` input must be provided when the `protocol` input is set to `custom`.'
+                )
+            protocol = copy.deepcopy(custom_protocol)
+        else:
+            protocol = copy.deepcopy(self.get_protocol(protocol))
         code = engines['relax']['code']
 
         pseudo_family_label = protocol.pop('pseudo_family')
@@ -87,15 +97,31 @@ def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR091
         recommended_ecut_wfc, recommended_ecut_rho = pseudo_family.get_recommended_cutoffs(
             structure=structure, stringency=cutoff_stringency, unit='Eh'
         )
+
+        # In both cases, if the protocol "hardcodes" the cutoff(s),
+        # I use that instead of the one from the pseudopotential family
+        # since it probably means the user really wanted that cutoff.
+        # I use try/except since I need to go deep into a dictionary and
+        # it is easier than using dict.get() a lot of times.
+        try:
+            protocol_ecut = protocol['base']['abinit']['parameters']['ecut']
+        except KeyError:
+            protocol_ecut = None
+
+        try:
+            protocol_pawecutdg = protocol['base']['abinit']['parameters']['pawecutdg']
+        except KeyError:
+            protocol_pawecutdg = None
+
         if pseudo_type == 'pseudo.jthxml':
             # JTH XML are PAW; we need `pawecutdg`
             cutoff_parameters = {
-                'ecut': np.ceil(recommended_ecut_wfc),
-                'pawecutdg': np.ceil(recommended_ecut_rho),
+                'ecut': protocol_ecut if protocol_ecut is not None else np.ceil(recommended_ecut_wfc),
+                'pawecutdg': protocol_pawecutdg if protocol_pawecutdg is not None else np.ceil(recommended_ecut_rho),
             }
         else:
             # All others are NC; no need for `pawecutdg`
-            cutoff_parameters = {'ecut': recommended_ecut_wfc}
+            cutoff_parameters = {'ecut': protocol_ecut if protocol_ecut is not None else np.ceil(recommended_ecut_wfc)}
 
         override = {
             'abinit': {
@@ -187,31 +213,9 @@ def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR091
                 warnings.warn(f'input magnetization per site was None, setting it to {magnetization_per_site}')
             magnetization_per_site = np.array(magnetization_per_site)
 
-            sum_is_zero = np.isclose(sum(magnetization_per_site), 0.0)
-            all_are_zero = np.all(np.isclose(magnetization_per_site, 0.0))
-            non_zero_mags = magnetization_per_site[~np.isclose(magnetization_per_site, 0.0)]
-            all_non_zero_pos = np.all(non_zero_mags > 0.0)
-            all_non_zero_neg = np.all(non_zero_mags < 0.0)
-
-            if all_are_zero:  # non-magnetic
-                warnings.warn(
-                    'all of the initial magnetizations per site are close to zero; doing a non-spin-polarized '
-                    'calculation'
-                )
-            elif (sum_is_zero and not all_are_zero) or (
-                not all_non_zero_pos and not all_non_zero_neg
-            ):  # antiferromagnetic
-                print('Detected antiferromagnetic!')
-                builder.abinit['parameters']['nsppol'] = 1  # antiferromagnetic system
-                builder.abinit['parameters']['nspden'] = 2  # scalar spin-magnetization in the z-axis
-                builder.abinit['parameters']['spinat'] = [[0.0, 0.0, mag] for mag in magnetization_per_site]
-            elif not all_are_zero and (all_non_zero_pos or all_non_zero_neg):  # ferromagnetic
-                print('Detected ferromagnetic!')
-                builder.abinit['parameters']['nsppol'] = 2  # collinear spin-polarization
-                builder.abinit['parameters']['nspden'] = 2  # scalar spin-magnetization in the z-axis
-                builder.abinit['parameters']['spinat'] = [[0.0, 0.0, mag] for mag in magnetization_per_site]
-            else:
-                raise ValueError(f'Initial magnetization {magnetization_per_site} is ambiguous')
+            builder.abinit['parameters']['nsppol'] = 2  # collinear spin-polarization
+            builder.abinit['parameters']['nspden'] = 2  # scalar spin-magnetization in the z-axis
+            builder.abinit['parameters']['spinat'] = [[0.0, 0.0, mag] for mag in magnetization_per_site]
         elif spin_type == SpinType.NON_COLLINEAR:
             if magnetization_per_site is None:
                 magnetization_per_site = get_initial_magnetization(structure)

src/aiida_common_workflows/workflows/relax/abinit/workchain.py

@@ -25,7 +25,7 @@ def get_stress(parameters):
 def get_forces(parameters):
     """Return the forces array from the given parameters node."""
     forces = orm.ArrayData()
-    forces.set_array(name='forces', array=np.array(parameters.base.attributes.get('forces')))
+    forces.set_array(name='forces', array=np.array(parameters.base.attributes.get('cart_forces')))
     return forces
 
 

src/aiida_common_workflows/workflows/relax/generator.py

@@ -1,5 +1,6 @@
 """Module with base input generator for the common structure relax workchains."""
 import abc
+from collections.abc import Sequence
 
 from aiida import orm, plugins
 
@@ -17,6 +18,25 @@ def validate_inputs(value, _):
     if value.get('magnetization_per_site') is not None and value.get('fixed_total_cell_magnetization') is not None:
         return 'the inputs `magnetization_per_site` and ' '`fixed_total_cell_magnetization` are mutually exclusive.'
 
+    if value.get('protocol') == 'custom' and value.get('custom_protocol') is None:
+        return 'the `custom_protocol` input must be provided when the `protocol` input is set to `custom`.'
+
+    if value.get('protocol') != 'custom' and value.get('custom_protocol') is not None:
+        return 'the `custom_protocol` input can only be provided when the `protocol` input is set to `custom`.'
+
+    # TODO: ensure all plugins actually honor this new custom_protocol input! (only QE implemented for now)
+
+    # Validate non-collinear spin type if magnetization per site is vector-valued
+    if value.get('magnetization_per_site') is not None and value.get('spin_type') not in [
+        SpinType.NON_COLLINEAR,
+        SpinType.SPIN_ORBIT,
+    ]:
+        for mag in value.get('magnetization_per_site'):
+            if isinstance(mag, Sequence):
+                return (
+                    'a vector valued magnetization is only allowed if `spin_type` is `NON_COLLINEAR` or `SPIN_ORBIT`.'
+                )
+
 
 class OptionalRelaxFeatures(OptionalFeature):
     FIXED_MAGNETIZATION = 'fixed_total_cell_magnetization'
@@ -45,7 +65,7 @@ def define(cls, spec):
         )
         spec.input(
             'protocol',
-            valid_type=ChoiceType(('fast', 'moderate', 'precise')),
+            valid_type=ChoiceType(('fast', 'moderate', 'precise', 'custom')),
             default='moderate',
             non_db=True,
             help='The protocol to use for the automated input generation. This value indicates the level of precision '
@@ -77,10 +97,21 @@ def define(cls, spec):
             valid_type=list,
             required=False,
             non_db=True,
-            help='The initial magnetization of the system. Should be a list of floats, where each float represents the '
-            'spin polarization in units of electrons, meaning the difference between spin up and spin down '
-            'electrons, for the site. This also corresponds to the magnetization of the site in Bohr magnetons '
-            '(μB).',
+            help='The initial magnetization of the system. Should be a list of floats/vectors, where each '
+            'float/vector represents the spin polarization in units of electrons, meaning the difference '
+            'between spin up and spin down electrons, for the site. This also corresponds to the '
+            'magnetization of the site in Bohr magnetons (μB). '
+            'If a single float is given for a site with non-collinear spins, '
+            'this is interpreted as a (0, 0, value) vector.',
+        )
+        spec.input(
+            'custom_protocol',
+            valid_type=(dict, type(None)),
+            non_db=True,
+            required=False,
+            default=None,
+            help='A custom protocol dictionary that can be provided when the `protocol` input is set to `custom`. '
+            'In that case, this dictionary will be used to override the default protocol settings.',
         )
         spec.input(
             'fixed_total_cell_magnetization',

src/aiida_common_workflows/workflows/relax/quantum_espresso/generator.py

@@ -97,7 +97,7 @@ def define(cls, spec):
         """
         super().define(spec)
         spec.inputs['protocol'].valid_type = ChoiceType(
-            ('fast', 'balanced', 'stringent', 'moderate', 'precise', 'verification-PBE-v1')
+            ('fast', 'balanced', 'stringent', 'moderate', 'precise', 'verification-PBE-v1', 'custom')
         )
         spec.inputs['spin_type'].valid_type = ChoiceType((SpinType.NONE, SpinType.COLLINEAR))
         spec.inputs['relax_type'].valid_type = ChoiceType(
@@ -155,7 +155,15 @@ def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR091
         # Currently, the `aiida-quantumespresso` workflows will expect one of the basic protocols to be passed to the
         # `get_builder_from_protocol()` method. Here, we switch to using the default protocol for the
         # `aiida-quantumespresso` plugin and pass the local protocols as `overrides`.
-        if (
+        if protocol == 'custom':
+            custom_protocol = kwargs.get('custom_protocol', None)
+            if custom_protocol is None:
+                raise ValueError(
+                    'The `custom_protocol` input must be provided when the `protocol` input is set to `custom`.'
+                )
+            overrides = custom_protocol
+            protocol = self._default_protocol
+        elif (
             protocol not in self.process_class._process_class.get_available_protocols()
             and self.process_class._process_class._check_if_alias(protocol)
             not in self.process_class._process_class.get_available_protocols()

src/aiida_common_workflows/workflows/relax/siesta/generator.py

@@ -1,12 +1,14 @@
 """Implementation of `aiida_common_workflows.common.relax.generator.CommonRelaxInputGenerator` for SIESTA."""
 import os
+from collections.abc import Sequence
 
 import yaml
 from aiida import engine, orm, plugins
 from aiida.common import exceptions
 
 from aiida_common_workflows.common import ElectronicType, RelaxType, SpinType
 from aiida_common_workflows.generators import ChoiceType, CodeType
+from aiida_common_workflows.utils import to_spherical
 
 from ..generator import CommonRelaxInputGenerator
 
@@ -63,7 +65,9 @@ def define(cls, spec):
         """
         super().define(spec)
         spec.inputs['protocol'].valid_type = ChoiceType(('fast', 'moderate', 'precise', 'verification-PBE-v1'))
-        spec.inputs['spin_type'].valid_type = ChoiceType((SpinType.NONE, SpinType.COLLINEAR))
+        spec.inputs['spin_type'].valid_type = ChoiceType(
+            (SpinType.NONE, SpinType.COLLINEAR, SpinType.NON_COLLINEAR, SpinType.SPIN_ORBIT)
+        )
         spec.inputs['relax_type'].valid_type = ChoiceType(
             (RelaxType.NONE, RelaxType.POSITIONS, RelaxType.POSITIONS_CELL, RelaxType.POSITIONS_SHAPE)
         )
@@ -123,17 +127,27 @@ def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR091
         if threshold_stress:
             parameters['md-max-stress-tol'] = str(threshold_stress) + ' eV/Ang**3'
         # ... spin options (including initial magentization) ...
-        if spin_type == SpinType.COLLINEAR:
+        if spin_type == SpinType.NONE:
+            parameters['spin'] = 'non-polarized'
+        elif spin_type == SpinType.COLLINEAR:
             parameters['spin'] = 'polarized'
+        elif spin_type == SpinType.NON_COLLINEAR:
+            parameters['spin'] = 'non-colinear'
+        elif spin_type == SpinType.SPIN_ORBIT:
+            parameters['spin'] = 'spin-orbit'
         if magnetization_per_site is not None:
             if spin_type == SpinType.NONE:
                 import warnings
 
                 warnings.warn('`magnetization_per_site` will be ignored as `spin_type` is set to SpinType.NONE')
-            if spin_type == SpinType.COLLINEAR:
+            else:
                 in_spin_card = '\n'
                 for i, magn in enumerate(magnetization_per_site):
-                    in_spin_card += f' {i+1} {magn} \n'
+                    if isinstance(magn, Sequence):
+                        r, theta, phi = to_spherical(magn)
+                        in_spin_card += f' {i+1} {r} {theta} {phi} \n'
+                    else:
+                        in_spin_card += f' {i+1} {magn} \n'
                 in_spin_card += '%endblock dm-init-spin'
                 parameters['%block dm-init-spin'] = in_spin_card
 

src/aiida_common_workflows/workflows/relax/vasp/generator.py

@@ -1,5 +1,5 @@
 """Implementation of `aiida_common_workflows.common.relax.generator.CommonRelaxInputGenerator` for VASP."""
-
+import copy
 import os
 import pathlib
 import typing as t
@@ -56,7 +56,9 @@ def define(cls, spec):
         spec.inputs['relax_type'].valid_type = ChoiceType(tuple(RelaxType))
         spec.inputs['electronic_type'].valid_type = ChoiceType((ElectronicType.METAL, ElectronicType.INSULATOR))
         spec.inputs['engines']['relax']['code'].valid_type = CodeType('vasp.vasp')
-        spec.inputs['protocol'].valid_type = ChoiceType(('fast', 'moderate', 'precise', 'verification-PBE-v1'))
+        spec.inputs['protocol'].valid_type = ChoiceType(
+            ('fast', 'moderate', 'precise', 'verification-PBE-v1', 'custom')
+        )
 
     def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR0912,PLR0915
         """Construct a process builder based on the provided keyword arguments.
@@ -67,6 +69,7 @@ def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR091
         structure = kwargs['structure']
         engines = kwargs['engines']
         protocol = kwargs['protocol']
+        custom_protocol = kwargs.get('custom_protocol', None)
         spin_type = kwargs['spin_type']
         relax_type = kwargs['relax_type']
         magnetization_per_site = kwargs.get('magnetization_per_site', None)
@@ -77,7 +80,15 @@ def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR091
         # Get the protocol that we want to use
         if protocol is None:
             protocol = self._default_protocol
-        protocol = self.get_protocol(protocol)
+
+        if protocol == 'custom':
+            if custom_protocol is None:
+                raise ValueError(
+                    'the `custom_protocol` input must be provided when the `protocol` input is set to `custom`.'
+                )
+            protocol = copy.deepcopy(custom_protocol)
+        else:
+            protocol = copy.deepcopy(self.get_protocol(protocol))
 
         # Set the builder
         builder = self.process_class.get_builder()
@@ -172,7 +183,14 @@ def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR091
                 previous_kpoints.base.attributes.get('mesh'), previous_kpoints.base.attributes.get('offset')
             )
         else:
-            kpoints.set_kpoints_mesh_from_density(protocol['kpoint_distance'])
+            if 'kpoints' in protocol and 'kpoint_distance' in protocol:
+                raise ValueError('Protocol cannot define both `kpoints` and `kpoint_distance` in protocol.')
+            if 'kpoints' not in protocol and 'kpoint_distance' not in protocol:
+                raise ValueError('Protocol must define either `kpoints` or `kpoint_distance` in protocol.')
+            if 'kpoints' in protocol:
+                kpoints.set_kpoints_mesh(protocol['kpoints'])
+            else:
+                kpoints.set_kpoints_mesh_from_density(protocol['kpoint_distance'])
         builder.vasp.kpoints = kpoints
 
         # Set the relax parameters