From 445ff0d4b68fae04edbf547cda7413492b9d4453 Mon Sep 17 00:00:00 2001
From: Giovanni Pizzi <gio.piz@gmail.com>
Date: Fri, 28 Nov 2025 16:05:15 +0100
Subject: [PATCH 01/16] Adding support for custom protocols

For now, only implemented for QE and abinit, needs to be implemented for all

Also, fixing a couple of bugs in the ACWF relax for ABINIT
---
 .../workflows/relax/abinit/generator.py       | 38 ++++++-------------
 .../workflows/relax/abinit/workchain.py       |  2 +-
 .../workflows/relax/generator.py              | 18 ++++++++-
 .../relax/quantum_espresso/generator.py       | 12 +++++-
 4 files changed, 39 insertions(+), 31 deletions(-)

diff --git a/src/aiida_common_workflows/workflows/relax/abinit/generator.py b/src/aiida_common_workflows/workflows/relax/abinit/generator.py
index 7778fc45..ab1714f9 100644
--- a/src/aiida_common_workflows/workflows/relax/abinit/generator.py
+++ b/src/aiida_common_workflows/workflows/relax/abinit/generator.py
@@ -51,7 +51,7 @@ 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 +62,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 +71,12 @@ 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')
@@ -187,31 +193,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)
diff --git a/src/aiida_common_workflows/workflows/relax/abinit/workchain.py b/src/aiida_common_workflows/workflows/relax/abinit/workchain.py
index 64b0a79b..65f8a590 100644
--- a/src/aiida_common_workflows/workflows/relax/abinit/workchain.py
+++ b/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
 
 
diff --git a/src/aiida_common_workflows/workflows/relax/generator.py b/src/aiida_common_workflows/workflows/relax/generator.py
index 79456dab..08907872 100644
--- a/src/aiida_common_workflows/workflows/relax/generator.py
+++ b/src/aiida_common_workflows/workflows/relax/generator.py
@@ -17,6 +17,13 @@ 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)
 
 class OptionalRelaxFeatures(OptionalFeature):
     FIXED_MAGNETIZATION = 'fixed_total_cell_magnetization'
@@ -45,7 +52,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 '
@@ -82,6 +89,15 @@ def define(cls, spec):
             'electrons, for the site. This also corresponds to the magnetization of the site in Bohr magnetons '
             '(μB).',
         )
+        spec.input(
+            'custom_protocol',
+            valid_type=dict,
+            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',
             valid_type=OptionalFeatureType(float),
diff --git a/src/aiida_common_workflows/workflows/relax/quantum_espresso/generator.py b/src/aiida_common_workflows/workflows/relax/quantum_espresso/generator.py
index dfa3b09a..a8e9142e 100644
--- a/src/aiida_common_workflows/workflows/relax/quantum_espresso/generator.py
+++ b/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()

From 82f77d967068a469a81b6965f6f0366fac0a2fe2 Mon Sep 17 00:00:00 2001
From: Giovanni Pizzi <gio.piz@gmail.com>
Date: Wed, 3 Dec 2025 15:29:02 +0100
Subject: [PATCH 02/16] Adding possibility to specify cutoffs in the protocol

This allows to override the protocol ones by the launching user
by using a custom protocol
---
 .../workflows/relax/abinit/generator.py       | 22 ++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/aiida_common_workflows/workflows/relax/abinit/generator.py b/src/aiida_common_workflows/workflows/relax/abinit/generator.py
index ab1714f9..2e6598cc 100644
--- a/src/aiida_common_workflows/workflows/relax/abinit/generator.py
+++ b/src/aiida_common_workflows/workflows/relax/abinit/generator.py
@@ -93,15 +93,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': {

From 0323343f75ea1ff14770174e3e438e92b831f85a Mon Sep 17 00:00:00 2001
From: Giovanni Pizzi <gio.piz@gmail.com>
Date: Wed, 3 Dec 2025 15:46:11 +0100
Subject: [PATCH 03/16] Fixing pre-commit issues

---
 .../workflows/relax/abinit/generator.py                   | 8 ++++++--
 src/aiida_common_workflows/workflows/relax/generator.py   | 5 +++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/aiida_common_workflows/workflows/relax/abinit/generator.py b/src/aiida_common_workflows/workflows/relax/abinit/generator.py
index 2e6598cc..0b9345d8 100644
--- a/src/aiida_common_workflows/workflows/relax/abinit/generator.py
+++ b/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', 'custom'))
+        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.
@@ -73,7 +75,9 @@ def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR091
 
         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`.')
+                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))
diff --git a/src/aiida_common_workflows/workflows/relax/generator.py b/src/aiida_common_workflows/workflows/relax/generator.py
index 08907872..a97035d0 100644
--- a/src/aiida_common_workflows/workflows/relax/generator.py
+++ b/src/aiida_common_workflows/workflows/relax/generator.py
@@ -18,13 +18,14 @@ def validate_inputs(value, _):
         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`.'
+        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`.'
+        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)
 
+
 class OptionalRelaxFeatures(OptionalFeature):
     FIXED_MAGNETIZATION = 'fixed_total_cell_magnetization'
 

From 62957f2fc6b12166ca9d59932e268c0aeff024b9 Mon Sep 17 00:00:00 2001
From: Giovanni Pizzi <giovanni.pizzi@psi.ch>
Date: Thu, 11 Dec 2025 14:10:17 +0100
Subject: [PATCH 04/16] Adding support for common protocols in VASP

---
 .../workflows/relax/vasp/generator.py            | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/src/aiida_common_workflows/workflows/relax/vasp/generator.py b/src/aiida_common_workflows/workflows/relax/vasp/generator.py
index b5f9fbec..0f7a4aa8 100644
--- a/src/aiida_common_workflows/workflows/relax/vasp/generator.py
+++ b/src/aiida_common_workflows/workflows/relax/vasp/generator.py
@@ -1,4 +1,5 @@
 """Implementation of `aiida_common_workflows.common.relax.generator.CommonRelaxInputGenerator` for VASP."""
+import copy
 import pathlib
 import typing as t
 
@@ -54,7 +55,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.
@@ -65,6 +68,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)
@@ -75,7 +79,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()

From 0ecc634cb02197b04c737fde97999465919e237f Mon Sep 17 00:00:00 2001
From: Giovanni Pizzi <giovanni.pizzi@psi.ch>
Date: Thu, 11 Dec 2025 14:13:21 +0100
Subject: [PATCH 05/16] Fixing default protocol

---
 src/aiida_common_workflows/workflows/relax/generator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aiida_common_workflows/workflows/relax/generator.py b/src/aiida_common_workflows/workflows/relax/generator.py
index a97035d0..a22b53f3 100644
--- a/src/aiida_common_workflows/workflows/relax/generator.py
+++ b/src/aiida_common_workflows/workflows/relax/generator.py
@@ -92,7 +92,7 @@ def define(cls, spec):
         )
         spec.input(
             'custom_protocol',
-            valid_type=dict,
+            valid_type=OptionalFeatureType(dict),
             non_db=True,
             required=False,
             default=None,

From 419a11aa8e2d903d6881861c867a45869d59a116 Mon Sep 17 00:00:00 2001
From: Giovanni Pizzi <gio.piz@gmail.com>
Date: Tue, 16 Dec 2025 09:55:54 +0100
Subject: [PATCH 06/16] Fix the type of a port

---
 src/aiida_common_workflows/workflows/relax/generator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aiida_common_workflows/workflows/relax/generator.py b/src/aiida_common_workflows/workflows/relax/generator.py
index a22b53f3..69479553 100644
--- a/src/aiida_common_workflows/workflows/relax/generator.py
+++ b/src/aiida_common_workflows/workflows/relax/generator.py
@@ -92,7 +92,7 @@ def define(cls, spec):
         )
         spec.input(
             'custom_protocol',
-            valid_type=OptionalFeatureType(dict),
+            valid_type=(dict, type(None)),
             non_db=True,
             required=False,
             default=None,

From af692834e8710058460056ae0cf5a4e3f9e997c8 Mon Sep 17 00:00:00 2001
From: Giovanni Pizzi <giovanni.pizzi@psi.ch>
Date: Thu, 18 Dec 2025 15:35:10 +0100
Subject: [PATCH 07/16] Adding support for explicit k-points in VASP generator

---
 .../workflows/relax/vasp/generator.py                    | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/aiida_common_workflows/workflows/relax/vasp/generator.py b/src/aiida_common_workflows/workflows/relax/vasp/generator.py
index 47a123cf..013a861f 100644
--- a/src/aiida_common_workflows/workflows/relax/vasp/generator.py
+++ b/src/aiida_common_workflows/workflows/relax/vasp/generator.py
@@ -183,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

From a739c24e4174cfc9398db65990fb545e31e93cc8 Mon Sep 17 00:00:00 2001
From: "A.H. Kole" <git.public.rqd55@slmail.me>
Date: Wed, 20 May 2026 00:19:20 +0200
Subject: [PATCH 08/16] Draft support for non-collinear spins for the
 CommonRelaxInputGenerator

---
 docs/source/workflows/base/relax/index.rst    |  5 +++--
 .../workflows/relax/generator.py              | 22 +++++++++++++++----
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/docs/source/workflows/base/relax/index.rst b/docs/source/workflows/base/relax/index.rst
index e9597da6..93d31657 100644
--- a/docs/source/workflows/base/relax/index.rst
+++ b/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,12 @@ 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/vectors).
   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.
   This also corresponds to the magnetization of the site in Bohr magnetons (μB).
+  For non-collinear spin calculations, a float-valued magnetization is interpreted as a (0, 0, value) vector for that site.
   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)`_.
 
diff --git a/src/aiida_common_workflows/workflows/relax/generator.py b/src/aiida_common_workflows/workflows/relax/generator.py
index 69479553..532453a0 100644
--- a/src/aiida_common_workflows/workflows/relax/generator.py
+++ b/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
 
@@ -25,6 +26,17 @@ def validate_inputs(value, _):
 
     # 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:
+        for mag in value.get('magnetization_per_site'):
+            if isinstance(mag, Sequence) and value.get('spin_type') not in [
+                SpinType.NON_COLLINEAR,
+                SpinType.SPIN_ORBIT,
+            ]:
+                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'
@@ -85,10 +97,12 @@ 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',

From 42f4c2e1ce34f08cf5fa6b5a784920bf565ef1bb Mon Sep 17 00:00:00 2001
From: "A.H. Kole" <git.public.rqd55@slmail.me>
Date: Wed, 20 May 2026 00:23:53 +0200
Subject: [PATCH 09/16] Add non-collinear spin to siesta relax workchain

---
 .../workflows/relax/siesta/generator.py       | 29 +++++++++++++++++-
 tests/workflows/relax/test_siesta.py          | 30 +++++++++++++++++++
 2 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/src/aiida_common_workflows/workflows/relax/siesta/generator.py b/src/aiida_common_workflows/workflows/relax/siesta/generator.py
index a958f827..bc0e08b0 100644
--- a/src/aiida_common_workflows/workflows/relax/siesta/generator.py
+++ b/src/aiida_common_workflows/workflows/relax/siesta/generator.py
@@ -1,5 +1,6 @@
 """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
@@ -15,6 +16,20 @@
 StructureData = plugins.DataFactory('core.structure')
 
 
+def _to_spherical(v):
+    from math import copysign
+
+    import numpy as np
+
+    x, y, z = v
+    r = np.sqrt(x * x + y * y + z * z)
+    rxy = np.sqrt(x * x + y * y)
+    theta = (180 / np.pi) * np.arccos(z / r) if r > 0 else 0
+    phi = (180 / np.pi) * copysign(np.arccos(x / rxy), y) if rxy > 0 else 0
+
+    return r, theta, phi
+
+
 class SiestaCommonRelaxInputGenerator(CommonRelaxInputGenerator):
     """Generator of inputs for the SiestaCommonRelaxWorkChain"""
 
@@ -63,7 +78,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)
         )
@@ -136,6 +153,16 @@ def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR091
                     in_spin_card += f' {i+1} {magn} \n'
                 in_spin_card += '%endblock dm-init-spin'
                 parameters['%block dm-init-spin'] = in_spin_card
+            if spin_type in [SpinType.NON_COLLINEAR, SpinType.SPIN_ORBIT]:
+                in_spin_card = '\n'
+                for i, magn in enumerate(magnetization_per_site):
+                    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
 
         # Basis
         basis = self._get_basis(protocol, structure)
diff --git a/tests/workflows/relax/test_siesta.py b/tests/workflows/relax/test_siesta.py
index a9501343..140dfa18 100644
--- a/tests/workflows/relax/test_siesta.py
+++ b/tests/workflows/relax/test_siesta.py
@@ -1,6 +1,8 @@
 """Tests for the :mod:`aiida_common_workflows.workflows.relax.siesta` module."""
+import numpy as np
 import pytest
 from aiida import engine, plugins
+from aiida_common_workflows.common import SpinType
 
 
 @pytest.fixture
@@ -71,3 +73,31 @@ def test_supported_spin_types(generator, default_builder_inputs):
         inputs['spin_type'] = spin_type
         builder = generator.get_builder(**inputs)
         assert isinstance(builder, engine.ProcessBuilder)
+
+
+@pytest.mark.usefixtures('psml_family')
+def test_magnetization_per_site(generator, default_builder_inputs):
+    """Test the ``magnetization_per_site`` keyword argument."""
+    inputs = default_builder_inputs
+
+    magnetization_per_site = [1.0]
+    for spin_type in [SpinType.COLLINEAR, SpinType.NON_COLLINEAR]:
+        builder = generator.get_builder(magnetization_per_site=magnetization_per_site, spin_type=spin_type, **inputs)
+        magn = float(builder['parameters']['%block dm-init-spin'].split()[1])
+        assert np.isclose(magn, 1.0)
+
+    magnetization_per_site = [(1.0, 0.0, 0.0)]
+    with pytest.raises(ValueError):
+        builder = generator.get_builder(
+            magnetization_per_site=magnetization_per_site, spin_type=SpinType.COLLINEAR, **inputs
+        )
+
+    magnetization_per_site = [(1.0, 0.0, 0.0)]
+    builder = generator.get_builder(
+        magnetization_per_site=magnetization_per_site, spin_type=SpinType.NON_COLLINEAR, **inputs
+    )
+    dm_init_string = builder['parameters']['%block dm-init-spin'].split()
+    r, theta, phi = float(dm_init_string[1]), float(dm_init_string[2]), float(dm_init_string[3])
+    assert np.isclose(r, 1.0)
+    assert np.isclose(theta, 90.0)
+    assert np.isclose(phi, 0.0)

From 6147cc6c1b0d4af895b0d31ea6725f538d7122a0 Mon Sep 17 00:00:00 2001
From: "A.H. Kole" <git.public.rqd55@slmail.me>
Date: Wed, 20 May 2026 01:00:32 +0200
Subject: [PATCH 10/16] Correctly set spin parameter for siesta relax workchain
 for non-collinear spins

---
 .../workflows/relax/siesta/generator.py                       | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/aiida_common_workflows/workflows/relax/siesta/generator.py b/src/aiida_common_workflows/workflows/relax/siesta/generator.py
index bc0e08b0..6bea64b6 100644
--- a/src/aiida_common_workflows/workflows/relax/siesta/generator.py
+++ b/src/aiida_common_workflows/workflows/relax/siesta/generator.py
@@ -142,6 +142,10 @@ def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR091
         # ... spin options (including initial magentization) ...
         if spin_type == SpinType.COLLINEAR:
             parameters['spin'] = 'polarized'
+        if spin_type == SpinType.NON_COLLINEAR:
+            parameters['spin'] = 'non-colinear'
+        if spin_type == SpinType.SPIN_ORBIT:
+            parameters['spin'] = 'spin-orbit'
         if magnetization_per_site is not None:
             if spin_type == SpinType.NONE:
                 import warnings

From adb903fdc3e567e24ef06c8e75a1ddce2184e492 Mon Sep 17 00:00:00 2001
From: "A.H. Kole" <git.public.rqd55@slmail.me>
Date: Wed, 20 May 2026 14:51:46 +0200
Subject: [PATCH 11/16] Move to_spherical to separate util module + improve
 robustness of spin_type handling in siesta relax workchain

---
 src/aiida_common_workflows/utils/__init__.py  | 18 ++++++++++++
 .../workflows/relax/siesta/generator.py       | 29 ++++++-------------
 2 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/src/aiida_common_workflows/utils/__init__.py b/src/aiida_common_workflows/utils/__init__.py
index e69de29b..41ed7857 100644
--- a/src/aiida_common_workflows/utils/__init__.py
+++ b/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
diff --git a/src/aiida_common_workflows/workflows/relax/siesta/generator.py b/src/aiida_common_workflows/workflows/relax/siesta/generator.py
index 6bea64b6..adb092da 100644
--- a/src/aiida_common_workflows/workflows/relax/siesta/generator.py
+++ b/src/aiida_common_workflows/workflows/relax/siesta/generator.py
@@ -8,6 +8,7 @@
 
 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
 
@@ -16,20 +17,6 @@
 StructureData = plugins.DataFactory('core.structure')
 
 
-def _to_spherical(v):
-    from math import copysign
-
-    import numpy as np
-
-    x, y, z = v
-    r = np.sqrt(x * x + y * y + z * z)
-    rxy = np.sqrt(x * x + y * y)
-    theta = (180 / np.pi) * np.arccos(z / r) if r > 0 else 0
-    phi = (180 / np.pi) * copysign(np.arccos(x / rxy), y) if rxy > 0 else 0
-
-    return r, theta, phi
-
-
 class SiestaCommonRelaxInputGenerator(CommonRelaxInputGenerator):
     """Generator of inputs for the SiestaCommonRelaxWorkChain"""
 
@@ -140,28 +127,30 @@ 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'
-        if spin_type == SpinType.NON_COLLINEAR:
+        elif spin_type == SpinType.NON_COLLINEAR:
             parameters['spin'] = 'non-colinear'
-        if spin_type == SpinType.SPIN_ORBIT:
+        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:
+            elif spin_type == SpinType.COLLINEAR:
                 in_spin_card = '\n'
                 for i, magn in enumerate(magnetization_per_site):
                     in_spin_card += f' {i+1} {magn} \n'
                 in_spin_card += '%endblock dm-init-spin'
                 parameters['%block dm-init-spin'] = in_spin_card
-            if spin_type in [SpinType.NON_COLLINEAR, SpinType.SPIN_ORBIT]:
+            elif spin_type in [SpinType.NON_COLLINEAR, SpinType.SPIN_ORBIT]:
                 in_spin_card = '\n'
                 for i, magn in enumerate(magnetization_per_site):
                     if isinstance(magn, Sequence):
-                        r, theta, phi = _to_spherical(magn)
+                        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'

From 0b66c74088634f2e0e9d5bd56b8c92f0271af7f7 Mon Sep 17 00:00:00 2001
From: "A.H. Kole" <git.public.rqd55@slmail.me>
Date: Wed, 20 May 2026 16:20:10 +0200
Subject: [PATCH 12/16] Clarifying documentation for 'magnetization_per_site'

---
 docs/source/workflows/base/relax/index.rst | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/docs/source/workflows/base/relax/index.rst b/docs/source/workflows/base/relax/index.rst
index 93d31657..6a1cfba3 100644
--- a/docs/source/workflows/base/relax/index.rst
+++ b/docs/source/workflows/base/relax/index.rst
@@ -163,12 +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/vectors).
+* ``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, a float-valued magnetization is interpreted as a (0, 0, value) vector for that site.
+  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)`_.
 

From 1b68aab61214c365cae6ba37665fe444413b6c29 Mon Sep 17 00:00:00 2001
From: "A.H. Kole" <git.public.rqd55@slmail.me>
Date: Fri, 22 May 2026 15:39:46 +0200
Subject: [PATCH 13/16] Refactor validation of magnetization_per_site

---
 .../workflows/relax/generator.py                       | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/aiida_common_workflows/workflows/relax/generator.py b/src/aiida_common_workflows/workflows/relax/generator.py
index 532453a0..189218aa 100644
--- a/src/aiida_common_workflows/workflows/relax/generator.py
+++ b/src/aiida_common_workflows/workflows/relax/generator.py
@@ -27,12 +27,12 @@ def validate_inputs(value, _):
     # 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:
+    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) and value.get('spin_type') not in [
-                SpinType.NON_COLLINEAR,
-                SpinType.SPIN_ORBIT,
-            ]:
+            if isinstance(mag, Sequence):
                 return (
                     'a vector valued magnetization is only allowed if `spin_type` is `NON_COLLINEAR` or `SPIN_ORBIT`.'
                 )

From 72294433397a1717487339c659c5468a970f609c Mon Sep 17 00:00:00 2001
From: "A.H. Kole" <git.public.rqd55@slmail.me>
Date: Fri, 22 May 2026 15:43:12 +0200
Subject: [PATCH 14/16] Refactor setting magnetization in siesta relax input
 generator

---
 .../workflows/relax/siesta/generator.py                   | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/src/aiida_common_workflows/workflows/relax/siesta/generator.py b/src/aiida_common_workflows/workflows/relax/siesta/generator.py
index adb092da..002d1d23 100644
--- a/src/aiida_common_workflows/workflows/relax/siesta/generator.py
+++ b/src/aiida_common_workflows/workflows/relax/siesta/generator.py
@@ -140,13 +140,7 @@ def _construct_builder(self, **kwargs) -> engine.ProcessBuilder:  # noqa: PLR091
                 import warnings
 
                 warnings.warn('`magnetization_per_site` will be ignored as `spin_type` is set to SpinType.NONE')
-            elif spin_type == SpinType.COLLINEAR:
-                in_spin_card = '\n'
-                for i, magn in enumerate(magnetization_per_site):
-                    in_spin_card += f' {i+1} {magn} \n'
-                in_spin_card += '%endblock dm-init-spin'
-                parameters['%block dm-init-spin'] = in_spin_card
-            elif spin_type in [SpinType.NON_COLLINEAR, SpinType.SPIN_ORBIT]:
+            else:
                 in_spin_card = '\n'
                 for i, magn in enumerate(magnetization_per_site):
                     if isinstance(magn, Sequence):

From 3de1c3a38d282efb8c3d103e716fe1917662b2cb Mon Sep 17 00:00:00 2001
From: "A.H. Kole" <git.public.rqd55@slmail.me>
Date: Fri, 22 May 2026 15:44:38 +0200
Subject: [PATCH 15/16] Fix punctuation in docs

---
 docs/source/workflows/base/relax/index.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/source/workflows/base/relax/index.rst b/docs/source/workflows/base/relax/index.rst
index 6a1cfba3..61d1ad93 100644
--- a/docs/source/workflows/base/relax/index.rst
+++ b/docs/source/workflows/base/relax/index.rst
@@ -169,7 +169,7 @@ Only ``structure`` and ``engines`` can be specified as a positional argument, al
   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
+  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.

From baa40ce357507a904df7e32a1aa26d91c946f60d Mon Sep 17 00:00:00 2001
From: "A.H. Kole" <git.public.rqd55@slmail.me>
Date: Fri, 22 May 2026 16:02:36 +0200
Subject: [PATCH 16/16] Add check for magnetization_per_site that line contains
 correct number of entries

---
 tests/workflows/relax/test_siesta.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tests/workflows/relax/test_siesta.py b/tests/workflows/relax/test_siesta.py
index 140dfa18..24d740ec 100644
--- a/tests/workflows/relax/test_siesta.py
+++ b/tests/workflows/relax/test_siesta.py
@@ -83,7 +83,9 @@ def test_magnetization_per_site(generator, default_builder_inputs):
     magnetization_per_site = [1.0]
     for spin_type in [SpinType.COLLINEAR, SpinType.NON_COLLINEAR]:
         builder = generator.get_builder(magnetization_per_site=magnetization_per_site, spin_type=spin_type, **inputs)
-        magn = float(builder['parameters']['%block dm-init-spin'].split()[1])
+        magn_string = builder['parameters']['%block dm-init-spin'].splitlines()[1].split()
+        assert len(magn_string) == 2
+        magn = float(magn_string[1])
         assert np.isclose(magn, 1.0)
 
     magnetization_per_site = [(1.0, 0.0, 0.0)]
@@ -96,8 +98,9 @@ def test_magnetization_per_site(generator, default_builder_inputs):
     builder = generator.get_builder(
         magnetization_per_site=magnetization_per_site, spin_type=SpinType.NON_COLLINEAR, **inputs
     )
-    dm_init_string = builder['parameters']['%block dm-init-spin'].split()
-    r, theta, phi = float(dm_init_string[1]), float(dm_init_string[2]), float(dm_init_string[3])
+    magn_string = builder['parameters']['%block dm-init-spin'].splitlines()[1].split()
+    assert len(magn_string) == 4
+    r, theta, phi = float(magn_string[1]), float(magn_string[2]), float(magn_string[3])
     assert np.isclose(r, 1.0)
     assert np.isclose(theta, 90.0)
     assert np.isclose(phi, 0.0)