✨ EpwPrepWorkChain: add support for SOC

src/aiida_epw/workflows/prep.py

@@ -8,7 +8,7 @@
 from aiida.engine import WorkChain, ToContext, if_
 from aiida_quantumespresso.workflows.ph.base import PhBaseWorkChain
 from aiida_quantumespresso.workflows.protocols.utils import ProtocolMixin
-from aiida_quantumespresso.common.types import ElectronicType
+from aiida_quantumespresso.common.types import ElectronicType, SpinType
 
 from aiida_quantumespresso.calculations.functions.create_kpoints_from_distance import (
     create_kpoints_from_distance,
@@ -485,6 +485,7 @@ def get_builder_from_protocol(
         bands_kpoints=None,
         workflow_type="mob",
         electronic_type=ElectronicType.METAL,
+        spin_type=SpinType.NONE,
         bandplot=False,
         w90_chk_to_ukk_script=None,
         **kwargs,
@@ -501,6 +502,7 @@ def get_builder_from_protocol(
         :param electronic_type: indicate the electronic character of the system through ``ElectronicType`` instance.
             Use ``ElectronicType.INSULATOR`` for valence bands only (e.g., for insulators/semiconductors),
             or ``ElectronicType.METAL`` (default) to include conduction bands.
+        :param spin_type: spin treatment passed to the Wannier90 builder. Defaults to ``SpinType.NONE``.
         :param kwargs: additional keyword arguments that will be passed to the ``get_builder_from_protocol`` of all the
             sub processes that are called by this workchain.
         :return: a process builder instance with all inputs defined ready for launch.
@@ -556,6 +558,7 @@ def get_builder_from_protocol(
                 reference_bands=reference_bands,
                 bands_kpoints=bands_kpoints,
                 electronic_type=electronic_type,
+                spin_type=spin_type,
                 **kwargs,
             )
             w90_bands.separate_plotting = False
@@ -616,6 +619,7 @@ def get_builder_from_protocol(
                 electronic_type=electronic_type,
                 projection_type=wannier_projection_type,
                 bands_kpoints=bands_kpoints,
+                spin_type=spin_type,
                 **kwargs,
             )
         elif wannier_projection_type == WannierProjectionType.ANALYTIC:
@@ -632,6 +636,7 @@ def get_builder_from_protocol(
                 electronic_type=electronic_type,
                 projection_type=wannier_projection_type,
                 bands_kpoints=bands_kpoints,
+                spin_type=spin_type,
                 **kwargs,
             )
         else:

tests/workflows/test_prep.py

@@ -357,6 +357,7 @@ def test_get_builder_from_protocol_w90_script(
         Wannier90BandsWorkChain,
     )
     from aiida_quantumespresso.workflows.ph.base import PhBaseWorkChain
+    from aiida_quantumespresso.common.types import SpinType
     from aiida_epw.workflows.base import EpwBaseWorkChain
 
     from plumpy.ports import Port, PortNamespace
@@ -383,17 +384,23 @@ def pop(self, key, default=None):
             return super().pop(key, default)
 
     # We need to return an instance that can behave like a dict/builder
+    w90_builder_kwargs = []
+
+    def mock_w90_get_builder(*args, **kwargs):
+        w90_builder_kwargs.append(kwargs)
+        return mock_get_builder(*args, **kwargs)
+
     def mock_get_builder(*args, **kwargs):
         builder = MockBuilder()
         if "w90_chk_to_ukk_script" in kwargs:
             builder.w90_chk_to_ukk_script = kwargs["w90_chk_to_ukk_script"]
         return builder
 
     monkeypatch.setattr(
-        Wannier90OptimizeWorkChain, "get_builder_from_protocol", mock_get_builder
+        Wannier90OptimizeWorkChain, "get_builder_from_protocol", mock_w90_get_builder
     )
     monkeypatch.setattr(
-        Wannier90BandsWorkChain, "get_builder_from_protocol", mock_get_builder
+        Wannier90BandsWorkChain, "get_builder_from_protocol", mock_w90_get_builder
     )
     monkeypatch.setattr(PhBaseWorkChain, "get_builder_from_protocol", mock_get_builder)
     monkeypatch.setattr(EpwBaseWorkChain, "get_builder_from_protocol", mock_get_builder)
@@ -436,8 +443,11 @@ def mock_get_builder(*args, **kwargs):
         w90_chk_to_ukk_script=w90_script,
         wannier_projection_type=WannierProjectionType.ATOMIC_PROJECTORS_QE,
         reference_bands=reference_bands,
+        spin_type=SpinType.SPIN_ORBIT,
     )
 
+    assert w90_builder_kwargs[0]["spin_type"] is SpinType.SPIN_ORBIT
+
     # Check that it is also propagated to epw_base builder
     assert builder.epw_base.w90_chk_to_ukk_script == w90_script