Add CommonRelaxInputGenerator subclass for DFT based codes

aiida_common_workflows/common/types.py

@@ -2,7 +2,7 @@
 """Module with basic type definitions."""
 from enum import Enum
 
-__all__ = ('ElectronicType', 'SpinType', 'RelaxType')
+__all__ = ('ElectronicType', 'SpinType', 'RelaxType', 'OccupationType', 'XcFunctionalType')
 
 
 class RelaxType(Enum):
@@ -24,12 +24,30 @@ class SpinType(Enum):
     NONE = 'none'
     COLLINEAR = 'collinear'
     NON_COLLINEAR = 'non_collinear'
-    SPIN_ORBIT = 'spin_orbit'
 
 
 class ElectronicType(Enum):
     """Enumeration of known electronic types."""
 
-    AUTOMATIC = 'automatic'
+    UNKNOWN = 'unknown'
     METAL = 'metal'
     INSULATOR = 'insulator'
+
+
+class OccupationType(Enum):
+    """Enumeration of known methods of treating electronic occupations."""
+
+    FIXED = 'fixed'
+    TETRAHEDRON = 'tetrahedron'
+    GAUSSIAN = 'gaussian'
+    FERMI_DIRAC = 'fermi-dirac'
+    METHFESSEL_PAXTON = 'methfessel-paxton'
+    MARZARI_VANDERBILT = 'marzari-vanderbilt'
+
+
+class XcFunctionalType(Enum):
+    """Enumeration of known exchange-correlation functional types."""
+
+    LDA = 'lda'
+    PBE = 'pbe'
+    PBESOL = 'pbesol'

aiida_common_workflows/workflows/relax/generator.py

@@ -5,7 +5,7 @@
 from aiida import orm
 from aiida import plugins
 
-from aiida_common_workflows.common import ElectronicType, RelaxType, SpinType
+from aiida_common_workflows.common import ElectronicType, RelaxType, SpinType, OccupationType, XcFunctionalType
 from aiida_common_workflows.generators import ChoiceType, InputGenerator
 
 __all__ = ('CommonRelaxInputGenerator',)
@@ -37,6 +37,12 @@ def define(cls, spec):
             help='The protocol to use for the automated input generation. This value indicates the level of precision '
             'of the results and computational cost that the input parameters will be selected for.',
         )
+        spec.input(
+            'spin_orbit',
+            valid_type=bool,
+            default=False,
+            help='Whether to apply spin-orbit coupling.',
+        )
         spec.input(
             'spin_type',
             valid_type=SpinType,
@@ -109,3 +115,81 @@ def define(cls, spec):
             required=False,
             help='Options for the geometry optimization calculation jobs.',
         )
+
+
+class CommonDftRelaxInputGenerator(CommonRelaxInputGenerator, metaclass=abc.ABCMeta):
+    """Input generator for the common relax workflow.
+
+    .. note:: This class is a subclass of the ``CommonRelaxInputGenerator`` but defines some additional inputs that are
+        common to a number of implementations.
+
+    This class should be subclassed by implementations for specific quantum engines. After calling the super, they can
+    modify the ports defined here in the base class as well as add additional custom ports.
+    """
+
+    @staticmethod
+    def validate_kpoints_shift(value, _):
+        """Validate the ``kpoints_shift`` input."""
+        if not isinstance(value, list) or len(value) != 3 or any(not isinstance(element, float) for element in value):
+            return f'The `kpoints_shift` argument should be a list of three floats, but got: `{value}`.'
+
+    @staticmethod
+    def validate_inputs(value, _):
+        """Docs."""
+        if value['spin_orbit'] is True and value['spin_type'] == SpinType.NONE:
+            return '`spin_type` cannot be `SpinType.NONE` for `spin_orbit = True`.'
+
+        smearing_broadening = value['smearing_broadening']
+        occupation_type = value['occupation_type']
+
+        if smearing_broadening is not None and occupation_type not in [
+            OccupationType.FIXED, OccupationType.TETRAHEDRON
+        ]:
+            return f'cannot define `smearing_broadening` for `occupation_type = {occupation_type}.'
+
+    @classmethod
+    def define(cls, spec):
+        """Define the specification of the input generator.
+
+        The ports defined on the specification are the inputs that will be accepted by the ``get_builder`` method.
+        """
+        super().define(spec)
+        spec.inputs.validator = cls.validate_inputs
+        spec.input(
+            'occupation_type',
+            valid_type=OccupationType,
+            serializer=OccupationType,
+            default=OccupationType.FIXED,
+            help='The way to treat electronic occupations.',
+        )
+        spec.input(
+            'smearing_broadening',
+            valid_type=float,
+            required=False,
+            help='The broadening of the smearing in eV. Should only be specified if a smearing method is defined for'
+            'the `occupation_type` input.',
+        )
+        spec.input(
+            'xc_functional',
+            valid_type=XcFunctionalType,
+            serializer=XcFunctionalType,
+            default=XcFunctionalType.PBE,
+            help='The functional for the exchange-correlation to be used.',
+        )
+        spec.input(
+            'kpoints_distance',
+            valid_type=float,
+            required=False,
+            help='The desired minimum distance between k-points in reciprocal space in 1/Å. The implementation will'
+            'guarantee that a k-point mesh is generated for which the distances between all adjacent k-points along '
+            'each cell vector are at most this distance. It is therefore possible that the distance is smaller than '
+            'requested along certain directions.',
+        )
+        spec.input(
+            'kpoints_shift',
+            valid_type=list,
+            validator=cls.validate_kpoints_shift,
+            required=False,
+            help='Optional shift to apply to all k-points of the k-point mesh. Should be a list of three floats where '
+            'each float is a number between 0 and 1.',
+        )

aiida_common_workflows/workflows/relax/quantum_espresso/generator.py

@@ -7,7 +7,7 @@
 from aiida_common_workflows.common import ElectronicType, RelaxType, SpinType
 from aiida_common_workflows.generators import ChoiceType, CodeType
 
-from ..generator import CommonRelaxInputGenerator
+from ..generator import CommonDftRelaxInputGenerator
 
 __all__ = ('QuantumEspressoCommonRelaxInputGenerator',)
 
@@ -62,7 +62,7 @@ def create_magnetic_allotrope(structure, magnetization_per_site):
     return (allotrope, allotrope_magnetic_moments)
 
 
-class QuantumEspressoCommonRelaxInputGenerator(CommonRelaxInputGenerator):
+class QuantumEspressoCommonRelaxInputGenerator(CommonDftRelaxInputGenerator):
     """Input generator for the common relax workflow implementation of Quantum ESPRESSO."""
 
     def __init__(self, *args, **kwargs):