(PR #50) Adding support for new cycling workflow

This PR adds frontend support for the new cycling workflow. A new protocol (multi-)selector has been added, with the previous custom protocol widget moved to a "create" tab allowing users to create/save new protocols to be selected from the "select" tab. The tomato settings widget was modified to support multiple protocols, each with its own settings/monitors. The experiment builder was modified w.r.t. these changes.

Finally, the submission process was modified to support the above features, including a newly improved input preview section providing a quick overview of selected protocols/settings/monitoring.

.gitignore

@@ -130,3 +130,4 @@ dmypy.json
 
 # Project-specific
 .aurora-*.ipynb
+*.json

INSTALLATION_GUIDE.md

@@ -234,8 +234,7 @@
   ```
   verdi group create BatterySamples
   verdi group create CyclingSpecs
-  verdi group create CalcJobs
-  verdi group create MonitorJobs
+  verdi group create WorkChains
   ```
 
 ## 4. Aurora app installation

aurora/common/models/battery_experiment.py

@@ -1,6 +1,6 @@
 import json
 import logging
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import pandas as pd
 from aiida_aurora.schemas.battery import (BatterySampleJsonTypes,
@@ -41,12 +41,14 @@ def __init__(self):
 
         self.selected_samples = pd.DataFrame()
         self.selected_protocol = ElectroChemSequence(method=[])
+        self.selected_protocols: Dict[int, ElectroChemSequence] = {}
+
         self.add_protocol_step()  # initialize sequence with first default step
 
     def reset_inputs(self):
         """Resets all inputs."""
-        # Not implemented yet...
-        return None
+        self.reset_selected_samples()
+        self.reset_selected_protocols()
 
     # ----------------------------------------------------------------------#
     # METHODS RELATED TO OBSERVABLES
@@ -368,6 +370,10 @@ def display_query_results(self, query: dict) -> None:
     # ----------------------------------------------------------------------#
     # METHODS RELATED TO PROTOCOLS
     # ----------------------------------------------------------------------#
+    def reset_selected_protocols(self):
+        """Resets all inputs."""
+        self.selected_protocols = {}
+
     def _count_technique_occurencies(self, technique):
         return [type(step)
                 for step in self.selected_protocol.method].count(technique)
@@ -428,6 +434,50 @@ def save_protocol(self):
         except OSError as err:
             print(f"Failed to save protocols => '{str(err)}'")
 
+    def add_selected_protocols(self, indices: List[int]) -> None:
+        """Add selected protocols to local cache.
+
+        Parameters
+        ----------
+        `indices` : `List[int]`
+            The `SelectMultiple` indices of the selected protocols.
+        """
+        for index in indices:
+            self.add_selected_protocol(index)
+
+    def add_selected_protocol(self, index: int) -> None:
+        """Add selected protocol to local cache.
+
+        Parameters
+        ----------
+        `index` : `int`
+             The `SelectMultiple` index of the selected protocol.
+        """
+        if index not in self.selected_protocols:
+            self.selected_protocols[index] = self.available_protocols[index]
+
+    def remove_selected_protocols(self, indices: List[int]) -> None:
+        """Remove selected protocol from local cache.
+
+        Parameters
+        ----------
+        `index` : `int`
+             The `SelectMultiple` index of the selected protocol.
+        """
+        for index in indices:
+            self.remove_selected_protocol(index)
+
+    def remove_selected_protocol(self, index: int) -> None:
+        """Remove selected protocol from local cache.
+
+        Parameters
+        ----------
+        `index` : `int`
+             The `SelectMultiple` index of the selected protocol.
+        """
+        if index in self.selected_protocols:
+            del self.selected_protocols[index]
+
     def update_available_protocols(
         self,
         source_file=AVAILABLE_PROTOCOLS_FILE,

aurora/engine/submit.py

@@ -1,82 +1,167 @@
-# from time import sleep
+from typing import Dict, List
 
-import aiida
 from aiida.engine import submit
 from aiida.manage.configuration import load_profile
-from aiida.orm import Dict, load_code, load_group
+from aiida.orm import Dict as AiiDADict
+from aiida.orm import List as AiiDAList
+from aiida.orm import load_code, load_group
+from aiida_aurora.data import (BatterySampleData, CyclingSpecsData,
+                               TomatoSettingsData)
+from aiida_aurora.schemas.battery import BatterySample
+from aiida_aurora.schemas.cycling import ElectroChemSequence
+from aiida_aurora.schemas.dgbowl import Tomato_0p2
+from aiida_aurora.workflows import CyclingSequenceWorkChain
 
 load_profile()
 
-BatterySampleData = aiida.plugins.DataFactory('aurora.batterysample')
-CyclingSpecsData = aiida.plugins.DataFactory('aurora.cyclingspecs')
-TomatoSettingsData = aiida.plugins.DataFactory('aurora.tomatosettings')
-BatteryCyclerExperiment = aiida.plugins.CalculationFactory('aurora.cycler')
+SAMPLES_GROUP = load_group("BatterySamples")
+PROTOCOLS_GROUP = load_group("CyclingSpecs")
+WORKFLOWS_GROUP = load_group("WorkChains")
 
-GROUP_SAMPLES = load_group("BatterySamples")
-GROUP_METHODS = load_group("CyclingSpecs")
-GROUP_CALCJOBS = load_group("CalcJobs")
 
-
-def submit_experiment(sample,
-                      method,
-                      tomato_settings,
-                      monitor_settings,
-                      code_name,
-                      sample_node_label="",
-                      method_node_label="",
-                      calcjob_node_label=""):
+def submit_experiment(
+    sample: BatterySample,
+    protocols: List[ElectroChemSequence],
+    settings: List[Tomato_0p2],
+    monitors: List[dict],
+    code_name: str,
+    sample_node_label="",
+    protocol_node_label="",
+    workchain_node_label="",
+) -> CyclingSequenceWorkChain:
     """
     sample : `aiida_aurora.schemas.battery.BatterySample`
     method : `aiida_aurora.schemas.cycling.ElectroChemSequence`
     tomato_settings : `aiida_aurora.schemas.dgbowl_schemas.Tomato_0p2`
     """
 
+    inputs = get_inputs(
+        sample,
+        sample_node_label,
+        code_name,
+        protocols,
+        settings,
+        monitors,
+        protocol_node_label,
+    )
+
+    return submit_job(inputs, workchain_node_label)
+
+
+def get_inputs(
+    sample: BatterySample,
+    sample_node_label: str,
+    code_name: str,
+    protocols: List[ElectroChemSequence],
+    settings: List[Tomato_0p2],
+    monitors: List[dict],
+    protocol_node_label: str,
+) -> dict:
+    """Prepare input dictionaries for workflow."""
+
+    inputs = {
+        "battery_sample": build_sample_node(sample, sample_node_label),
+        "tomato_code": load_code(code_name),
+        "protocol_order": AiiDAList(),
+        "protocols": {},
+        "control_settings": {},
+        "monitor_settings": {},
+    }
+
+    # push cycler locking (if requested) to final workflow step
+    if any(not ts.unlock_when_done for ts in settings):
+        for ts in settings:
+            ts.unlock_when_done = True
+        settings[-1].unlock_when_done = False
+
+    for protocol, settings, protocol_monitors in zip(
+            protocols,
+            settings,
+            monitors,
+    ):
+
+        step = protocol.name
+
+        inputs["protocol_order"].append(step)
+
+        inputs["protocols"][step] = build_protocol_node(
+            protocol,
+            protocol_node_label,
+        )
+
+        inputs["control_settings"][step] = build_settings_node(settings)
+
+        inputs["monitor_settings"][step] = build_monitors_input(
+            protocol,
+            protocol_monitors,
+        )
+
+    return inputs
+
+
+def build_sample_node(
+    sample: BatterySample,
+    sample_node_label: str,
+) -> BatterySampleData:
+    """Construct an AiiDA data node from battery sample data."""
     sample_node = BatterySampleData(sample.dict())
-    if sample_node_label:
-        sample_node.label = sample_node_label
+    sample_node.label = sample_node_label
     sample_node.store()
-    GROUP_SAMPLES.add_nodes(sample_node)
+    SAMPLES_GROUP.add_nodes(sample_node)
+    return sample_node
+
+
+def build_protocol_node(
+    protocol: ElectroChemSequence,
+    protocol_node_label: str,
+) -> CyclingSpecsData:
+    """Construct an AiiDA data node from cycling protocol data."""
+    protocol_node = CyclingSpecsData(protocol.dict())
+    protocol_node.label = protocol_node_label
+    protocol_node.store()
+    PROTOCOLS_GROUP.add_nodes(protocol_node)
+    return protocol_node
+
 
-    method_node = CyclingSpecsData(method.dict())
-    if method_node_label:
-        method_node.label = method_node_label
-    method_node.store()
-    GROUP_METHODS.add_nodes(method_node)
+def build_settings_node(settings: Tomato_0p2) -> TomatoSettingsData:
+    """Construct an AiiDA data node from tomato settings data."""
+    settings_node = TomatoSettingsData(settings.dict())
+    settings_node.label = ""
+    settings_node.store()
+    return settings_node
 
-    tomato_settings_node = TomatoSettingsData(tomato_settings.dict())
-    tomato_settings_node.label = ""  # TODO? write default name generator - e.g. "tomato-True-monitor_600"
-    tomato_settings_node.store()
 
-    code = load_code(code_name)
+def build_monitors_input(
+    protocol: ElectroChemSequence,
+    protocol_monitors: Dict[str, dict],
+) -> dict:
+    """Construct a dictionary of `Dict` monitors for the protocol."""
 
-    builder = BatteryCyclerExperiment.get_builder()
-    builder.battery_sample = sample_node
-    builder.code = code
-    builder.protocol = method_node
-    builder.control_settings = tomato_settings_node
+    monitors: Dict[str, dict] = {}
 
-    if monitor_settings:
+    for label, monitor_settings in protocol_monitors.items():
         refresh_rate = monitor_settings.pop("refresh_rate", 600)
-        builder.monitors = {
-            "capacity":
-            Dict({
+        monitor_name = f"{protocol.name}_{label}"
+        monitors[monitor_name] = AiiDADict(
+            dict={
                 "entry_point": "aurora.monitors.capacity_threshold",
                 "minimum_poll_interval": refresh_rate,
                 "kwargs": {
                     "settings": monitor_settings,
                     "filename": "snapshot.json",
                 },
-            }),
-        }
+            })
 
-    job = submit(builder)
-    job.label = calcjob_node_label
-    print(f"Job <{job.pk}> submitted to AiiDA...")
-    GROUP_CALCJOBS.add_nodes(job)
+    return monitors
 
-    if monitor_settings:
-        job.set_extra('monitored', True)
-    else:
-        job.set_extra('monitored', False)
 
-    return job
+def submit_job(
+    inputs: dict,
+    workchain_node_label: str,
+) -> CyclingSequenceWorkChain:
+    """docstring"""
+    workchain = submit(CyclingSequenceWorkChain, **inputs)
+    workchain.label = workchain_node_label
+    print(f"Workflow <{workchain.pk}> submitted to AiiDA...")
+    WORKFLOWS_GROUP.add_nodes(workchain)
+    return workchain

aurora/experiment/protocols/__init__.py

@@ -1,7 +1,9 @@
 from .custom import CyclingCustom
+from .selector import ProtocolSelector
 from .standard import CyclingStandard
 
 __all__ = [
     "CyclingCustom",
+    "ProtocolSelector",
     "CyclingStandard",
 ]