in-progress

aiida_aurora/workflows/cycling_sequence.py

@@ -1,4 +1,5 @@
 from aiida import orm
+from aiida.common.exceptions import ValidationError
 from aiida.engine import ToContext, WorkChain, append_, while_
 from aiida.plugins import CalculationFactory, DataFactory
 
@@ -8,22 +9,6 @@
 TomatoSettingsData = DataFactory('aurora.tomatosettings')
 
 
-def validate_inputs(inputs, ctx=None):
-    """Validate the inputs of the entire input namespace."""
-
-    error_message = ''
-    for namekey in inputs['techniques'].keys():
-        if namekey not in inputs['control_settings']:
-            error_message += f'namekey {namekey} missing in control_settings\n'
-
-    for namekey in inputs['control_settings'].keys():
-        if namekey not in inputs['techniques']:
-            error_message += f'namekey {namekey} missing in techniques\n'
-
-    if len(error_message) > 0:
-        return error_message
-
-
 class CyclingSequenceWorkChain(WorkChain):
     """This workflow represents a process containing a variable number of steps."""
 
@@ -44,7 +29,7 @@ def define(cls, spec):
         )
 
         spec.input_namespace(
-            "techniques",
+            "protocols",
             dynamic=True,
             valid_type=CyclingSpecsData,
             help="List of experiment specifications."
@@ -55,6 +40,14 @@ def define(cls, spec):
             valid_type=TomatoSettingsData,
             help="List of experiment control settings."
         )
+
+        spec.input_namespace(
+            "monitors",
+            dynamic=True,
+            valid_type=orm.Dict,
+            help="List of battery experiment monitors."
+        )
+
         spec.output_namespace(
             "results",
             dynamic=True,
@@ -64,6 +57,7 @@ def define(cls, spec):
 
         spec.outline(
             cls.setup_workload,
+            cls.validate_inputs,
             while_(cls.has_steps_remaining)(
                 cls.run_cycling_step,
             ),
@@ -72,28 +66,45 @@ def define(cls, spec):
 
     def setup_workload(self):
         """Takes the inputs and wraps them together."""
-        self.worksteps_keynames = list(self.inputs['techniques'].keys())
+        self.protocol_step = list(self.inputs['protocols'].keys())
+
+    def validate_inputs(self) -> str:
+        """Validate the inputs of the entire input namespace."""
+
+        error_message = ""
+
+        for protocol in self.protocol_step:
+
+            if protocol not in self.inputs['control_settings'].keys():
+                error_message += f"{protocol} key missing in control settings\n"
+
+            if protocol not in self.inputs['monitors']:
+                error_message += f'{protocol} key missing in monitors\n'
+
+        if len(error_message) > 0:
+            raise ValidationError(error_message)
 
     def has_steps_remaining(self):
         """Checks if there is any remaining step"""
-        return len(self.worksteps_keynames) > 0
+        return len(self.protocol_step) > 0
 
     def run_cycling_step(self):
         """Description"""
-        current_keyname = self.worksteps_keynames.pop(0)
+        current_keyname = self.protocol_step.pop(0)
         inputs = {
             'code': self.inputs.tomato_code,
             'battery_sample': self.inputs.battery_sample,
-            'technique': self.inputs.techniques[current_keyname],
+            'technique': self.inputs.protocols[current_keyname],
             'control_settings': self.inputs.control_settings[current_keyname],
+            'monitors': {f"{current_keyname}_monitor": self.inputs.monitors[current_keyname]},
         }
         running = self.submit(CyclerCalcjob, **inputs)
         self.report(f'launching CyclerCalcjob<{running.pk}>')
         return ToContext(workchains=append_(running))
 
     def gather_results(self):
         """Description"""
-        keynames = list(self.inputs['techniques'].keys())
+        keynames = list(self.inputs['protocols'].keys())
         if len(self.ctx.workchains) != len(keynames):
             raise RuntimeError('Problem with workchain!')