(PR #37) Updating the schemas and cycling workchain

The battery schema has been updated to account for most, if not all robot output fields. This was primarily motivated by the need for electrode weights during visualization (capacity normalization), then expanded to account for all fields in case needed in the future (and also, for good measure). The cycling workchain was adjusted accordingly. 

Unrelated to the schemas but also included here, the new group label field was added to the workchain to support auto-generating experiment groups on submission.

aiida_aurora/data/battery.py

@@ -54,6 +54,7 @@ def validate(self, parameters_dict):  # pylint: disable=no-self-use
         # Manual fix to convert date-times to ISO string format
         # TODO integrate this into the data schema
         d["metadata"]["creation_datetime"] = d["metadata"]["creation_datetime"].isoformat()
+        d["metadata"].pop("groups", None)
         return d
 
     def get_json(self):

aiida_aurora/schemas/battery.py

@@ -1,107 +1,101 @@
 from datetime import datetime
 from enum import Flag
-from typing import Literal, Optional
+from typing import Literal, Optional, Set
 
-from pydantic import BaseModel, NonNegativeFloat, PositiveFloat, root_validator
+from pydantic import BaseModel, NonNegativeFloat, PositiveInt
 
 from .utils import extract_schema_types
 
 
-class BatteryComposition(BaseModel):  # TypedDict?
-    description: Optional[str] = None
-    cathode: Optional[str] = None
-    anode: Optional[str] = None
-    electrolyte: Optional[str] = None
-
-    class Config:
-        # exclude fields from export, to avoid validation errors when reloaded
-        fields = {
-            'cathode': {
-                'exclude': True
-            },
-            'anode': {
-                'exclude': True
-            },
-            'electrolyte': {
-                'exclude': True
-            },
-        }
-
-    @root_validator
-    def validate_composition(cls, values):
-        """
-        Check that components are not specified if 'description' is specified
-        then build components from a 'description' string, or vice versa.
-        # TODO: what to do if 'description' is not in the C|E|A format?
-        """
-        if values['description']:
-            if any(values[key] for key in ('cathode', 'anode', 'electrolyte')):
-                raise ValueError("You cannot specify a 'description' and any component at the same time.")
-            values['description'] = values['description'].strip()
-            components = list(map(str.strip, values['description'].split('|')))
-            if len(components) == 3:
-                values['cathode'], values['electrolyte'], values['anode'] = components
-            else:
-                values['cathode'], values['electrolyte'], values['anode'] = (None, None, None)
-                # raise ValueError(
-                # "Composition 'description' does not have 3 components (i.e. {cathode}|{electrolyte}|{anode}).")
-        elif any(values[key] for key in ('cathode', 'anode', 'electrolyte')):
-            for key in ('cathode', 'anode', 'electrolyte'):
-                values[key] = values[key].strip()
-            values['description'] = f"{values['cathode']}|{values['electrolyte']}|{values['anode']}"
-        else:
-            raise ValueError("You must specify either a string 'description' or the components.")
-        return values
-
-
-class BatteryCapacity(BaseModel):  # TypedDict?
-    nominal: PositiveFloat
+class Component(BaseModel):
+    description: Optional[str]
+
+
+class Diameter(BaseModel):
+    nominal: NonNegativeFloat
     actual: Optional[NonNegativeFloat]
-    units: Literal["mAh", "Ah"]
+    units: Literal["mm"] = "mm"
 
 
-class BatteryMetadata(BaseModel):
-    name: str
-    creation_datetime: datetime
-    creation_process: str
+class Capacity(BaseModel):
+    nominal: NonNegativeFloat
+    actual: Optional[NonNegativeFloat]
+    units: Literal["mAh", "Ah"] = "mAh"
 
 
-class BatterySpecs(BaseModel):
-    """
-    Battery specification schema.
-    """
-    manufacturer: str
-    composition: BatteryComposition
-    form_factor: str
-    capacity: BatteryCapacity
+class Electrolyte(Component):
+    formula: str
+    position: PositiveInt
+    amount: NonNegativeFloat
 
-    # manufacturer, form_factor:
-    # should we use a Literal or a validator to check that they are one of the available ones?
 
-    # add pre-validator to specify capacity as str (e.g. '4.8 mAh')?
+class ElectrodeWeight(BaseModel):
+    total: NonNegativeFloat
+    collector: NonNegativeFloat
+    net: Optional[NonNegativeFloat]
+    units: Literal["mg", "g"] = "mg"
 
 
-class BatterySample(BatterySpecs):
-    """
-    Battery sample schema.
-    """
-    battery_id: int
-    metadata: BatteryMetadata
+class Electrode(Component):
+    formula: str
+    position: PositiveInt
+    diameter: Diameter
+    weight: ElectrodeWeight
+    capacity: Capacity
+
+
+class Separator(Component):
+    name: str  # ? use `Literal` of available?
+    diameter: Diameter
+
+
+class Spacer(Component):
+    value: NonNegativeFloat
+    units: Literal["mm"] = "mm"
+
+
+class Composition(BaseModel):
+    description: Optional[str]
+    anode: Electrode
+    cathode: Electrode
+    electrolyte: Electrolyte
+    separator: Separator
+    spacer: Spacer
+
+
+class BatterySpecs(BaseModel):
+    case: str  # ? use `Literal` of available?
+    manufacturer: str  # ? use `Literal` of available?
+    composition: Composition
+    capacity: Capacity
+    np_ratio: Optional[str]
+
+
+class BatteryMetadata(BaseModel):
+    name: str
+    groups: Set[str] = {"all-samples"}
+    batch: str = ""
+    subbatch: str = "0"
+    creation_datetime: datetime
+    creation_process: str
 
 
 class ChargeState(Flag):
-    """Defines the charge state of a battery."""
     CHARGED = True
     DISCHARGED = False
 
 
-class BatteryState(BatterySample):
-    """
-    Battery state schema.
-    """
-    state_id: int
+class BatteryState(BaseModel):
+    used = False
     charged: ChargeState = ChargeState.CHARGED
 
 
+class BatterySample(BaseModel):
+    id: int
+    # state: BatteryState  # TODO move to metadata?
+    specs: BatterySpecs
+    metadata: BatteryMetadata
+
+
 BatterySpecsJsonTypes = extract_schema_types(BatterySpecs)
 BatterySampleJsonTypes = extract_schema_types(BatterySample)

aiida_aurora/schemas/cycling.py

@@ -584,8 +584,3 @@ def move_step_forward(self, i):
         if i < self.n_steps - 1:
             j = i + 1
             self.method[i], self.method[j] = self.method[j], self.method[i]
-
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, ElectroChemSequence):
-            return NotImplemented
-        return self.name == other.name

aiida_aurora/schemas/dgbowl/converters/sample.py

@@ -3,23 +3,24 @@
 from aiida_aurora.schemas.battery import BatterySample
 
 
-def batterysample_to_sample_0(batsample: BatterySample, SampleSchema: BaseModel):
+def batterysample_to_sample_0(
+    sample: BatterySample,
+    SampleSchema: BaseModel,
+) -> BaseModel:
     """
     Convert a BatterySample into a Sample.
 
     Compatible with the following dgbowl-schemas payload versions:
         [0.1, 0.2]
     """
 
-    if not isinstance(batsample, BatterySample):
-        if isinstance(batsample, dict):
-            batsample = BatterySample(**batsample)
+    if not isinstance(sample, BatterySample):
+        if isinstance(sample, dict):
+            sample = BatterySample(**sample)
         else:
             raise TypeError()
-    # if batsample.capacity.units == "mAh":
-    # capacity = float(batsample.capacity.nominal) * 0.001
-    # elif batsample.capacity.units == "Ah":
-    # capacity = float(batsample.capacity.nominal)
-
-    sample = SampleSchema(name=batsample.metadata.name, capacity=batsample.capacity.nominal)
-    return sample
+    C = sample.specs.capacity.nominal
+    return SampleSchema(
+        name=sample.metadata.name,
+        capacity=C if sample.specs.capacity.units == "Ah" else C * 0.001,
+    )

aiida_aurora/schemas/utils/__init__.py

@@ -40,7 +40,7 @@ def _make_formatted_dict(my_dict, key_arr, val):
         key = key_arr[i]
         if key not in current:
             if i == len(key_arr) - 1:
-                current[key] = val
+                current[key] = None if val == 'None' or isna(val) else val
             else:
                 current[key] = {}
         else:
@@ -59,7 +59,7 @@ def pd_dataframe_to_formatted_json(df, sep="."):
     result = []
     for _, row in df.iterrows():
         parsed_row = {}
-        for idx, val in row.iteritems():
+        for idx, val in row.items():
             keys = idx.split(sep)
             parsed_row = _make_formatted_dict(parsed_row, keys, val)
         result.append(parsed_row)
@@ -79,17 +79,19 @@ def dict_to_formatted_json(series, sep="."):
 
 def extract_schema_types(model, sep="."):
     """Convert a pydantic schema into a nested dictionary containing types."""
-    SCHEMA_TYPES = {'string': str, 'integer': int, 'number': float}
+    SCHEMA_TYPES = {
+        'string': str,
+        'integer': int,
+        'number': float,
+        'array': object,
+        'boolean': bool,
+    }
     schema_dic = {}
     for name, sdic in model.schema()['properties'].items():
         if '$ref' in sdic:
-            # print(f"  {name} is a class {sdic['$ref']}")
-            sub_schema = extract_schema_types(
-                getattr(battery_schemas, sdic['$ref'].split('/')[-1])
-            )  # call extract_schema for the $ref class
+            sub_schema = extract_schema_types(getattr(battery_schemas, sdic['$ref'].split('/')[-1]))
             for key, value in sub_schema.items():
                 schema_dic[f'{name}.{key}'] = value
         elif 'type' in sdic:
-            # print(f"{name} is of type {sdic['type']}")
             schema_dic[name] = SCHEMA_TYPES[sdic['type']]
     return schema_dic

aiida_aurora/workflows/cycling_sequence.py

@@ -1,8 +1,9 @@
 from aiida import orm
 from aiida.engine import ToContext, WorkChain, append_, while_
-from aiida.plugins import CalculationFactory, DataFactory
+from aiida.plugins import DataFactory
+
+from aiida_aurora.calculations import BatteryCyclerExperiment
 
-CyclerCalcjob = CalculationFactory('aurora.cycler')
 CyclingSpecsData = DataFactory('aurora.cyclingspecs')
 BatterySampleData = DataFactory('aurora.batterysample')
 TomatoSettingsData = DataFactory('aurora.tomatosettings')
@@ -55,6 +56,12 @@ def define(cls, spec):
             help="List of experiment protocol names in order of execution.",
         )
 
+        spec.input(
+            "group_label",
+            valid_type=orm.Str,
+            help="A label prefix for grouping experiments.",
+        )
+
         spec.input_namespace(
             "protocols",
             dynamic=True,
@@ -120,34 +127,47 @@ def inspect_cycling_step(self):
 
     def run_cycling_step(self):
         """Run the next cycling step."""
-        current_keyname = self.worksteps_keynames.pop(0)
+        protocol_name = self.worksteps_keynames.pop(0)
 
         inputs = {
             'code': self.inputs.tomato_code,
             'battery_sample': self.inputs.battery_sample,
-            'protocol': self.inputs.protocols[current_keyname],
-            'control_settings': self.inputs.control_settings[current_keyname],
+            'protocol': self.inputs.protocols[protocol_name],
+            'control_settings': self.inputs.control_settings[protocol_name],
         }
 
-        has_monitors = current_keyname in self.inputs.monitor_settings
+        has_monitors = protocol_name in self.inputs.monitor_settings
 
         if has_monitors:
-            inputs['monitors'] = self.inputs.monitor_settings[current_keyname]
+            inputs['monitors'] = self.inputs.monitor_settings[protocol_name]
 
-        running = self.submit(CyclerCalcjob, **inputs)
+        running = self.submit(BatteryCyclerExperiment, **inputs)
         sample_name = self.inputs.battery_sample.attributes["metadata"]["name"]
-        running.label = f"{current_keyname} | {sample_name}"
+        running.label = f"{protocol_name} | {sample_name}"
 
         if has_monitors:
             running.set_extra('monitored', True)
         else:
             running.set_extra('monitored', False)
 
-        orm.load_group("CalcJobs").add_nodes(running)
-
-        self.report(f'launching CyclerCalcjob<{running.pk}>')
+        workflow_group = self.inputs.group_label.value
+        experiment_group = f"{workflow_group}/{protocol_name}"
+        for group in (
+            "all-experiments",
+            protocol_name,
+            workflow_group,
+            experiment_group,
+        ):
+            self.add_to_group(running, group)
+
+        self.report(f'launching BatteryCyclerExperiment<{running.pk}>')
         return ToContext(subprocesses=append_(running))
 
+    def add_to_group(self, node: BatteryCyclerExperiment, label: str) -> None:
+        """docstring"""
+        group_label = f"aurora/experiments/{label}"
+        orm.Group.collection.get_or_create(group_label)[0].add_nodes(node)
+
     def gather_results(self):
         """Gather the results from all cycling steps."""
         keynames = list(self.inputs['protocols'].keys())

docs/source/nitpick-exceptions

@@ -20,6 +20,7 @@ py:class dgbowl_schemas.tomato.payload_0_2.tomato.Tomato
 
 py:class pydantic.main.BaseModel
 py:class pydantic.generics.GenericModel
+py:class pydantic.types.PositiveInt
 py:class pydantic.types.PositiveFloat
 py:class pydantic.types.NonNegativeFloat
 

examples/cycling_sequence/test_cycling_sequence.py

@@ -18,15 +18,15 @@ def generate_test_inputs():
     nodes_dict = {}
 
     nodes_dict['sample'] = BatterySampleData({
-        'manufacturer':
+        'specs.manufacturer':
         'fake_maufacturer',
-        'composition':
+        'specs.composition':
         dict(description='C|E|A'),
-        'form_factor':
+        'specs.case':
         'fake_form',
-        'capacity':
+        'specs.capacity':
         dict(nominal=1.0, units='Ah'),
-        'battery_id':
+        'id':
         666,
         'metadata':
         dict(