diff --git a/src/aiida_epw/workflows/prep.py b/src/aiida_epw/workflows/prep.py index bf54f109..d5fd4171 100644 --- a/src/aiida_epw/workflows/prep.py +++ b/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: diff --git a/tests/workflows/test_prep.py b/tests/workflows/test_prep.py index c717e9cb..65d37bba 100644 --- a/tests/workflows/test_prep.py +++ b/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,6 +384,12 @@ 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: @@ -390,10 +397,10 @@ def mock_get_builder(*args, **kwargs): 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