testing

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/dgbowl/converters/sample.py

@@ -3,23 +3,23 @@
 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
+    return SampleSchema(
+        name=sample.metadata.name,
+        capacity=sample.specs.capacity.nominal,
+    )

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

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(