Add CommonRelaxInputGenerator subclass for DFT based codes

Author: sphuber

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', 'SmearingMethodType', 'XcFunctionalType')
 
 
 class RelaxType(Enum):
@@ -33,3 +33,18 @@ class ElectronicType(Enum):
     AUTOMATIC = 'automatic'
     METAL = 'metal'
     INSULATOR = 'insulator'
+
+
+class SmearingMethodType(Enum):
+    """Enumeration of known smearing method types."""
+
+    GAUSSIAN = 'gaussian'
+    FERMI_DIRAC = 'fermi-dirac'
+
+
+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, SmearingMethodType, XcFunctionalType
 from aiida_common_workflows.generators import ChoiceType, InputGenerator
 
 __all__ = ('CommonRelaxInputGenerator',)
@@ -109,3 +109,65 @@ 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}`.'
+
+    @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.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.',
+        )
+        spec.input(
+            'smearing_method',
+            valid_type=SmearingMethodType,
+            serializer=SmearingMethodType,
+            default=SmearingMethodType.GAUSSIAN,
+            help='The method to use for the smearing of the electronic occupations.',
+        )
+        spec.input(
+            'smearing_broadening',
+            valid_type=float,
+            required=False,
+            help='The broadening of the smearing in eV.',
+        )
+        spec.input(
+            'functional',
+            valid_type=XcFunctionalType,
+            serializer=XcFunctionalType,
+            default=XcFunctionalType.PBE,
+            help='The functional for the exchange-correlation to be used.',
+        )

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):