Adding support for common protocols in VASP

Author: giovannipizzi

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