feat(batch): typed LoadPolicy/CellSpec + resolve_specs (#699)

cellpy/batch/__init__.py

@@ -23,6 +23,13 @@
     write_journal,
 )
 from cellpy.batch.layout import BatchPaths, ensure_dirs
+from cellpy.batch.policy import (
+    CellSpec,
+    LoadPolicy,
+    SourcePreference,
+    parse_argument,
+    resolve_specs,
+)
 
 __all__ = [
     "Journal",
@@ -33,4 +40,9 @@
     "journal_from_custom_json",
     "BatchPaths",
     "ensure_dirs",
+    "SourcePreference",
+    "LoadPolicy",
+    "CellSpec",
+    "resolve_specs",
+    "parse_argument",
 ]

cellpy/batch/policy.py

@@ -0,0 +1,184 @@
+"""Typed batch loading options + spec resolution (batch v3, #699).
+
+Replaces the kwargs tunnels of the legacy ``CyclingExperiment.update`` (79
+``kwargs.pop/get`` calls, precedence documented in a single docstring) with two
+dataclasses and one pure function:
+
+- :class:`LoadPolicy` -- batch-wide loading knobs (``force_cellpy``/``force_raw``/
+  ``force_recalc``/... collapse into typed fields).
+- :class:`CellSpec` -- fully resolved per-cell loading instructions.
+- :func:`resolve_specs` -- the *single* place journal rows, policy-level
+  overrides and per-cell overrides merge, with the precedence the legacy code
+  smeared across ~200 lines of ``update()``:
+
+      journal row  <  journal ``argument``  <  policy overrides  <  per-cell
+
+The legacy merge was ``{**cell_spec_page, **kwargs, **cell_spec}``
+(batch_experiments.py:361); this reproduces it as a tested, pure function.
+"""
+
+from __future__ import annotations
+
+import ast
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Mapping
+
+from cellpy.batch.journal import FILENAME, Journal
+
+
+class SourcePreference(str, Enum):
+    """Which source a cell is loaded from."""
+
+    AUTO = "auto"  # cellpy file if fresh, else raw (today's default)
+    CELLPY_ONLY = "cellpy_only"  # replaces force_cellpy=True
+    RAW_ONLY = "raw_only"  # replaces force_raw_file=True
+
+
+@dataclass
+class LoadPolicy:
+    """Batch-wide loading options (one object instead of a kwargs tunnel)."""
+
+    source: SourcePreference = SourcePreference.AUTO
+    recalc: bool = False  # replaces force_recalc
+    max_cycle: int | None = None
+    accept_errors: bool = True  # errors collected, not raised
+    all_in_memory: bool = False
+    skip_bad_cells: bool = False
+    selector: dict | None = None  # forwarded to the cellpy-file loader
+    loader_kwargs: dict = field(default_factory=dict)  # the one escape hatch
+    #: batch-level per-field overrides applied to every cell (e.g. {"mass": 1.0}).
+    overrides: dict = field(default_factory=dict)
+
+
+@dataclass
+class CellSpec:
+    """Fully resolved per-cell loading instructions."""
+
+    label: str
+    raw_files: list = field(default_factory=list)
+    cellpy_file: Any | None = None
+    instrument: str | None = None
+    model: str | None = None
+    mass: float | None = None
+    nom_cap: float | None = None
+    area: float | None = None
+    cycle_mode: str | None = None
+    #: leftover per-cell knobs (recalc, data_points, ...) not mapped to a field.
+    overrides: dict = field(default_factory=dict)
+
+
+#: Journal columns that map directly to a typed :class:`CellSpec` field.
+_SPEC_FIELDS = ("instrument", "model", "mass", "nom_cap", "area", "cycle_mode")
+
+
+def _is_nan(value: Any) -> bool:
+    return isinstance(value, float) and value != value
+
+
+def _clean(value: Any) -> Any:
+    """Normalise pandas/polars null-ish values to ``None``."""
+    if value is None or _is_nan(value):
+        return None
+    return value
+
+
+def _coerce_scalar(value: Any) -> Any:
+    """Coerce a string spec value the way the legacy update() did."""
+    if not isinstance(value, str):
+        return value
+    low = value.strip().lower()
+    if low == "true":
+        return True
+    if low == "false":
+        return False
+    if low in ("none", ""):
+        return None
+    try:
+        return ast.literal_eval(value)
+    except (ValueError, SyntaxError):
+        return value
+
+
+def parse_argument(argument: Any) -> dict:
+    """Parse a journal ``argument`` cell into a dict of coerced values.
+
+    Accepts a dict (``{"recalc": "False"}``), the compact string form
+    (``"recalc=False;data_points=(1, 10000)"``), or null-ish -> ``{}``.
+    """
+    if argument is None or _is_nan(argument):
+        return {}
+    if isinstance(argument, dict):
+        return {key: _coerce_scalar(val) for key, val in argument.items()}
+    if isinstance(argument, str):
+        text = argument.strip()
+        if not text:
+            return {}
+        parsed: dict = {}
+        for part in text.split(";"):
+            if "=" not in part:
+                continue
+            key, val = part.split("=", 1)
+            parsed[key.strip()] = _coerce_scalar(val.strip())
+        return parsed
+    return {}
+
+
+def _as_list(value: Any) -> list:
+    if value is None:
+        return []
+    if isinstance(value, list):
+        return list(value)
+    return [value]
+
+
+def resolve_specs(
+    journal: Journal,
+    policy: LoadPolicy | None = None,
+    per_cell: Mapping[str, Mapping[str, Any]] | None = None,
+) -> list[CellSpec]:
+    """Resolve one :class:`CellSpec` per cell in ``journal``.
+
+    Precedence (later wins): journal columns < journal ``argument`` <
+    ``policy.overrides`` < ``per_cell[label]``.
+    """
+    policy = policy or LoadPolicy()
+    per_cell = per_cell or {}
+
+    specs: list[CellSpec] = []
+    for row in journal.pages.iter_rows(named=True):
+        label = row[FILENAME]
+
+        journal_fields = {
+            field_name: _clean(row.get(field_name))
+            for field_name in _SPEC_FIELDS
+            if field_name in row
+        }
+        # cell_type stands in for cycle_mode when the latter is absent
+        if not journal_fields.get("cycle_mode") and _clean(row.get("cell_type")):
+            journal_fields["cycle_mode"] = _clean(row.get("cell_type"))
+
+        argument = parse_argument(row.get("argument"))
+        overrides = dict(per_cell.get(label, {}))
+
+        merged = {**journal_fields, **argument, **policy.overrides, **overrides}
+
+        specs.append(
+            CellSpec(
+                label=label,
+                raw_files=_as_list(_clean(row.get("raw_file_names"))),
+                cellpy_file=_clean(row.get("cellpy_file_name")),
+                instrument=merged.get("instrument"),
+                model=merged.get("model"),
+                mass=merged.get("mass"),
+                nom_cap=merged.get("nom_cap"),
+                area=merged.get("area"),
+                cycle_mode=merged.get("cycle_mode"),
+                overrides={
+                    key: val
+                    for key, val in merged.items()
+                    if key not in _SPEC_FIELDS
+                },
+            )
+        )
+    return specs

tests/test_batch_v3_policy.py

@@ -0,0 +1,91 @@
+"""Tests for batch v3 policy: LoadPolicy/CellSpec + resolve_specs (#699)."""
+
+import polars as pl
+
+from cellpy.batch import (
+    CellSpec,
+    LoadPolicy,
+    SourcePreference,
+    parse_argument,
+    read_journal,
+    resolve_specs,
+)
+from cellpy.batch.journal import FILENAME, Journal
+
+
+# ---- argument parsing (matches legacy coercion) -------------------------
+
+
+def test_parse_argument_string_form():
+    got = parse_argument("recalc=False;data_points=(1, 10000)")
+    assert got == {"recalc": False, "data_points": (1, 10000)}
+
+
+def test_parse_argument_dict_form_coerces():
+    assert parse_argument({"recalc": "TRUE", "keep": "none"}) == {
+        "recalc": True,
+        "keep": None,
+    }
+
+
+def test_parse_argument_empty_and_null():
+    assert parse_argument(None) == {}
+    assert parse_argument("") == {}
+    assert parse_argument(float("nan")) == {}
+
+
+# ---- resolve_specs precedence: journal < policy < per-cell --------------
+
+
+def _journal(**columns):
+    columns.setdefault(FILENAME, ["a"])
+    return Journal(name="t", project="p", pages=pl.DataFrame(columns))
+
+
+def test_resolve_specs_journal_only():
+    j = _journal(mass=[1.0], instrument=["arbin_res"], argument=["recalc=False"])
+    (spec,) = resolve_specs(j)
+    assert isinstance(spec, CellSpec)
+    assert spec.label == "a"
+    assert spec.mass == 1.0
+    assert spec.instrument == "arbin_res"
+    assert spec.overrides["recalc"] is False
+
+
+def test_resolve_specs_policy_overrides_journal():
+    j = _journal(mass=[1.0], argument=["recalc=False"])
+    policy = LoadPolicy(overrides={"mass": 2.0, "recalc": True})
+    (spec,) = resolve_specs(j, policy=policy)
+    assert spec.mass == 2.0  # policy beats journal
+    assert spec.overrides["recalc"] is True  # policy beats journal argument
+
+
+def test_resolve_specs_per_cell_wins():
+    j = _journal(mass=[1.0], argument=["recalc=False"])
+    policy = LoadPolicy(overrides={"mass": 2.0})
+    (spec,) = resolve_specs(j, policy=policy, per_cell={"a": {"mass": 3.0}})
+    assert spec.mass == 3.0  # per-cell beats policy beats journal
+
+
+def test_resolve_specs_cell_type_is_cycle_mode_fallback():
+    j = _journal(cell_type=["anode"])
+    (spec,) = resolve_specs(j)
+    assert spec.cycle_mode == "anode"
+
+
+def test_resolve_specs_real_journal(parameters):
+    j = read_journal(parameters.journal_file_json_path)
+    specs = resolve_specs(j)
+    assert len(specs) == len(j)
+    assert all(isinstance(s, CellSpec) for s in specs)
+    assert {s.label for s in specs} == set(j.cell_names)
+    # the argument column ("recalc=...") lands in overrides, not a spec field
+    assert any("recalc" in s.overrides for s in specs)
+
+
+def test_loadpolicy_defaults():
+    p = LoadPolicy()
+    assert p.source is SourcePreference.AUTO
+    assert p.recalc is False
+    assert p.accept_errors is True
+    assert p.loader_kwargs == {} and p.overrides == {}