"""Base classes for ionworks.pipeline."""
from typing import Any, NoReturn
import warnings
import pybamm
from pydantic import BaseModel, ConfigDict, Field, field_validator
from .serialize import serialize_parameter_value, serialize_parameters
def _get_element_type(comp: Any) -> str | None:
"""Determine the ``element_type`` discriminator for a pipeline component.
Parameters
----------
comp : Any
A pipeline element instance (``DirectEntry``, ``DataFit``,
``Validation``, ``Calculation``, …).
Returns
-------
str or None
The canonical ``element_type`` string (``"entry"``, ``"data_fit"``,
``"validation"``, or ``"calculation"``), or ``None`` for raw configs
(e.g. dicts) that are not recognised schema objects, so they pass
through unchanged.
"""
cls = type(comp)
mro_names = {c.__name__ for c in cls.__mro__}
if "DirectEntry" in mro_names:
return "entry"
if "DirectEntryFunctionSchema" in mro_names:
return "entry"
if cls.__name__ == "DataFit":
return "data_fit"
if cls.__name__ == "ArrayDataFit":
return "array_data_fit"
if cls.__name__ == "LinearConfidenceInterval":
return "linear_confidence_interval"
if cls.__name__ == "SobolSensitivity":
return "sobol_sensitivity"
if cls.__name__ == "Validation":
return "validation"
if "Calculation" in mro_names:
return "calculation"
return None
def _is_dataframe(value) -> bool:
"""Return True for a pandas or polars DataFrame, duck-typed.
Avoids importing pandas/polars at module load so the schema package stays
usable without them installed.
"""
return hasattr(value, "to_dict") and hasattr(value, "columns")
def _serialize_value(value):
"""Recursively serialize a value, calling to_config on schemas.
A value with no wire form raises rather than passing through: returned
as-is it survives ``to_config`` and fails only once the server reads it.
"""
if value is None:
return None
if isinstance(value, BaseSchema):
return value.to_config()
if isinstance(value, dict):
return {k: _serialize_value(v) for k, v in value.items()}
if isinstance(value, (list, tuple)):
return [_serialize_value(v) for v in value]
# A ParameterValues serialises via the shared helper (drops citations,
# flattens Scalars), not its own ``to_config``/``to_json`` — so it must
# precede the generic ``to_config`` branch below.
if isinstance(value, pybamm.ParameterValues):
return serialize_parameters(value)
# A DataFit parameter may be a pybamm expression linking other parameters,
# which the pipeline parser reads back as serialised symbol-JSON.
if isinstance(value, pybamm.Symbol):
return serialize_parameter_value(value=value)
if hasattr(value, "to_config"):
return value.to_config()
if hasattr(value, "model_dump"):
return value.model_dump(exclude_none=True)
# pandas or polars DataFrame -> dict of lists for serialization
if _is_dataframe(value):
cls = type(value)
mod = getattr(cls, "__module__", "")
if "pandas" in mod:
return value.to_dict(orient="list")
if "polars" in mod:
return value.to_dict(as_series=False)
try:
return value.to_dict(orient="list")
except TypeError:
return value.to_dict(as_series=False)
if isinstance(value, (int, float, str, bool)):
return value
if hasattr(value, "tolist"): # numpy arrays
return value.tolist()
raise ValueError(
f"{type(value).__name__} has no wire form. Convert it to a number, "
f"string, list, dict, or pybamm object before putting it in a schema. "
f"Got {value!r:.200}."
)
[docs]
def reject_runtime_object(value: Any, field_name: str, expected: str) -> NoReturn:
"""Reject a value that is neither a config dict nor a schema instance.
Centralizes the boundary guard shared by the distribution, sampler, prior,
and objective resolvers. The schema contract holds only ``ionworks.pipeline``
objects — a config dict or the matching schema instance — so a live runtime
object must be converted (via ``.to_config()`` or the matching ``ionworks.pipeline``
schema class) before it enters a schema, rather than embedded raw. Callers
invoke this on the fall-through branch after they have already accepted the
dict and schema-instance forms.
Parameters
----------
value : Any
The rejected value, echoed into the error message via ``repr``.
field_name : str
The schema field being resolved (e.g. ``"distribution"``, ``"sampler"``,
``"prior"``, ``"objective"``), used to phrase the message.
expected : str
User-facing description of the accepted non-dict form, phrased with its
article so it reads in the message (e.g. ``"a Distribution instance"``,
``"an ionworks.pipeline objective instance"``). Use the public vocabulary,
not an internal implementation class name.
Raises
------
ValueError
Always. The function never returns normally.
"""
raise ValueError(
f"Invalid {field_name} {value!r}: expected a config dict or {expected}; "
f"got {type(value).__name__}. Live runtime objects are "
f"not accepted inside a schema."
)
[docs]
class BaseSchema(BaseModel):
"""Shared parent of every schema class in this package.
Construct a concrete subclass (``iw.pipeline.Pipeline``, ``iw.pipeline.DataFit``,
…) rather than this. It rejects unknown fields, so a typo fails at
construction, and provides ``.to_config()``, the only supported serializer:
``.model_dump()`` omits the ``type`` discriminator and uses attribute names
(``data_input``) rather than wire names (``data``).
"""
model_config = ConfigDict(
extra="forbid",
validate_assignment=True,
arbitrary_types_allowed=True,
populate_by_name=True,
)
# Fields to exclude from serialization (comes from dict key in parser)
_exclude_fields: set = set()
# When False, to_config omits the "type" discriminator. Set on subclasses
# whose serialised output must be a plain option dict (e.g. algorithm
# options passed to pipeline algorithm classes that reject unknown keys).
_emit_type: bool = True
# When True, to_config emits only the fields the caller explicitly set, for
# subclasses whose output is merged over a consumer's own defaults.
_only_set_fields: bool = False
# When False, a bare DataFrame in the "data" field skips the wrap below.
# Set on payload wrappers (e.g. DataPayloadSpec, MetadataSpec) whose own
# field happens to be named "data" too: they are already the wrapper, so
# re-applying the bare-DataFrame convention would nest it a second time.
_wrap_bare_dataframe: bool = True
[docs]
def to_config(self) -> dict:
"""Build the dict you submit through ``ionworks.api``.
Nested schemas convert recursively, ``None`` fields are dropped, and each
component the server must re-identify gains a ``"type"`` key.
Notes
-----
The only supported serializer. ``.model_dump()`` (even with
``by_alias=True``) omits the ``"type"`` discriminator, so it is not a
payload: ``iw.pipeline.stats.Uniform(lb=0.0, ub=1.0).to_config()`` gives
``{"distribution": "Uniform", "lb": 0.0, "ub": 1.0}``, whereas
``.model_dump(exclude_none=True)`` gives ``{"lb": 0.0, "ub": 1.0}``.
"""
config = {}
# Iteration yields declared fields in declaration order, then any
# extra="allow" keys in the order supplied, so output is stable.
for key, value in self:
if value is None:
continue
# Skip excluded fields
if key in self._exclude_fields:
continue
if self._only_set_fields and key not in self.model_fields_set:
continue
# ``extra="allow"`` keys have no field, so they keep their own name.
field = type(self).model_fields.get(key)
output_key = (
(field.serialization_alias or field.alias or key) if field else key
)
try:
serialized = _serialize_value(value)
except ValueError as e:
raise ValueError(f"{output_key}: {e}") from None
# A bare DataFrame in the `data` field serializes to a column dict;
# the parser expects it wrapped as {"data": <columns>}. Mirror the
# runtime DataInputController, which auto-wraps a bare DataFrame, so
# callers don't have to write {"data": df} by hand. An already-wrapped
# dict or a string path is not a DataFrame and is left untouched.
if (
output_key == "data"
and _is_dataframe(value)
and self._wrap_bare_dataframe
):
serialized = {"data": serialized}
config[output_key] = serialized
# Add type field based on class name (except for top-level containers)
cls_name = self.__class__.__name__
if self._emit_type and cls_name not in (
"Pipeline",
"DataFit",
"ArrayDataFit",
"LinearConfidenceInterval",
"SobolSensitivity",
"Parameter",
"DirectEntry",
"Validation",
"Calculation",
):
config["type"] = cls_name
return config
[docs]
def run(self, *args, **kwargs):
raise NotImplementedError(
"Execution not available. Submit to Ionworks API instead."
)
[docs]
class Pipeline(BaseSchema):
"""A sequence of steps that together produce a parameterised cell model.
Each element pulls parameters from a source (``DirectEntry``), fits them to
data (``DataFit``), derives quantities (``Calculation``) or checks them against
held-out data (``Validation``), and passes its parameters to the next, so order
matters. Submit ``.to_config()`` through ``ionworks.api`` to run it server-side.
Parameters
----------
elements : dict, optional
Mapping of element name to pipeline element (``DataFit``,
``DirectEntry``, ``Validation``, ``Calculation``, …). The name
is used to identify the element in the pipeline report.
``None`` is treated as an empty pipeline by ``to_config``.
output_file : str, optional
Optional file path for persisting the final parameter values
produced by the pipeline. If ``None``, parameters are not
written to disk.
name : str, optional
Human-readable name for the pipeline, used in reports and logs.
description : str, optional
Free-text description of what the pipeline does.
Examples
--------
>>> # known parameters (e.g. ambient temperature)
>>> known = iw.pipeline.direct_entries.DirectEntry(
... parameters={"Ambient temperature [K]": 298.15},
... )
>>> # fit one parameter against an OCP measurement
>>> obj = iw.pipeline.objectives.OCPHalfCell(
... electrode="positive",
... data_input="path/to/ocp.csv",
... )
>>> fit = iw.pipeline.DataFit(
... objectives={"ocp": obj},
... parameters={"Positive electrode capacity [A.h]": iw.pipeline.Parameter(
... "Positive electrode capacity [A.h]", initial_value=3.0, bounds=(2.0, 4.0),
... )},
... )
>>> # validate the fit against held-out cycling data
>>> val_obj = iw.pipeline.objectives.CurrentDriven(
... data_input="path/to/cycle.csv", options={"model": "SPM"}
... )
>>> val = iw.pipeline.Validation(objectives={"cycle": val_obj})
>>> pipeline = iw.pipeline.Pipeline(
... {"known": known, "ocp fit": fit, "validate": val}
... )
>>> config = pipeline.to_config()
>>> # then submit `config` via ionworks.api
"""
elements: dict | None = Field(
default=None,
description=(
"Mapping of element name to pipeline element (``DataFit``, "
"``DirectEntry``, ``Validation``, ``Calculation``, …). The name "
"is used to identify the element in the pipeline report. "
"``None`` is treated as an empty pipeline by ``to_config``."
),
)
output_file: str | None = Field(
default=None,
description=(
"Optional file path for persisting the final parameter values "
"produced by the pipeline. If None, parameters are not written "
"to disk."
),
)
name: str | None = Field(
default=None,
description=("Human-readable name for the pipeline, used in reports and logs."),
)
description: str | None = Field(
default=None,
description="Free-text description of what the pipeline does.",
)
def __init__(
self,
elements: dict | None = None,
output_file: str | None = None,
name: str | None = None,
description: str | None = None,
**extra: Any,
):
# Forward extra kwargs to BaseModel.__init__ so Pydantic enforces
# ``extra='forbid'`` at the field-validation layer (rather than
# raising a Python-level TypeError that bypasses Pydantic).
super().__init__(
elements=elements,
output_file=output_file,
name=name,
description=description,
**extra,
)
@field_validator("elements", mode="before")
@classmethod
def _validate_elements(cls, elements):
"""Resolve and validate each element config against its schema class.
A field validator (not a model validator) so it also runs on
``pipeline.elements = {...}`` reassignment, not just construction. Runs
ahead of ``SimplePipeline``'s expensive-element field validator, which
then classifies the resolved instances. Existing instances (e.g.
``ArrayDataFit``) pass through unchanged.
"""
if isinstance(elements, dict):
from ._element_registry import validate_element
return {name: validate_element(elem) for name, elem in elements.items()}
return elements
[docs]
def to_config(self) -> dict:
"""Build the pipeline payload you submit through ``ionworks.api``.
Walks every element you added to ``elements``, serialises it
(each element's own ``to_config``), and tags it with its kind
(``"entry"``, ``"data_fit"``, ``"validation"``, …) so the
server knows how to dispatch it. Pass the returned dict to the
Ionworks API to run the full pipeline.
"""
config = {"elements": {}}
for name, comp in (self.elements or {}).items():
# Serialize the component
if hasattr(comp, "to_config"):
elem_config = comp.to_config()
elif hasattr(comp, "model_dump"):
elem_config = comp.model_dump(exclude_none=True)
elif isinstance(comp, dict):
elem_config = dict(comp)
else:
elem_config = {}
# Set element_type only for known schema types; raw configs (e.g. dicts) keep their existing element_type
elem_type = _get_element_type(comp)
if elem_type is not None:
elem_config["element_type"] = elem_type
config["elements"][name] = elem_config
# Add other fields
if self.output_file is not None:
config["output_file"] = self.output_file
if self.name is not None:
config["name"] = self.name
if self.description is not None:
config["description"] = self.description
return config
# Keep in sync with the backend's EXPENSIVE_ELEMENT_TYPES
# (backend/src/pydantic_models/simple_pipelines.py).
_EXPENSIVE_TYPES = frozenset({"data_fit", "array_data_fit", "validation"})
[docs]
class SimplePipeline(Pipeline):
"""A pipeline that permits at most one expensive element.
.. warning::
Deprecated, and removed in a future release. Use :class:`Pipeline`,
which accepts the same elements.
``SimplePipeline`` is identical to :class:`Pipeline` except that it
rejects configurations containing more than one expensive element
(``DataFit``, ``ArrayDataFit``, or ``Validation``) at construction time,
matching the constraint the SimplePipeline API endpoint enforces.
Parameters
----------
elements : dict, optional
Mapping of element name to pipeline element (``DataFit``,
``DirectEntry``, ``Validation``, ``Calculation``, …). At most
one element may be a ``DataFit``, ``ArrayDataFit``, or ``Validation``.
output_file : str, optional
Optional file path for persisting the final parameter values
produced by the pipeline. If ``None``, parameters are not
written to disk.
name : str, optional
Human-readable name for the pipeline, used in reports and logs.
description : str, optional
Free-text description of what the pipeline does.
Raises
------
ValueError
If ``elements`` contains more than one ``DataFit``, ``ArrayDataFit``,
or ``Validation`` element.
Examples
--------
>>> entry = iw.pipeline.direct_entries.DirectEntry(
... parameters={"Ambient temperature [K]": 298.15},
... )
>>> obj = iw.pipeline.objectives.OCPHalfCell(
... electrode="positive",
... data_input="path/to/ocp.csv",
... )
>>> fit = iw.pipeline.DataFit(
... objectives={"ocp": obj},
... parameters={"Positive electrode capacity [A.h]": iw.pipeline.Parameter(
... "Positive electrode capacity [A.h]",
... initial_value=3.0,
... bounds=(2.0, 4.0),
... )},
... )
>>> pipeline = iw.pipeline.SimplePipeline(
... elements={"known": entry, "ocp fit": fit},
... name="Single OCP fit",
... )
>>> config = pipeline.to_config()
"""
def __init__(self, *args: Any, **kwargs: Any):
# Not an access-time warning, which would fire on every ``import *``.
warnings.warn(
"iw.pipeline.SimplePipeline is deprecated and will be removed in a future "
"release. Use iw.pipeline.Pipeline, which accepts the same elements, and "
"submit it through the regular pipeline API. Migration guide: "
"https://docs.ionworks.com/build/parameterize/api#migrating-to-pipeline",
DeprecationWarning,
stacklevel=2,
)
super().__init__(*args, **kwargs)
@field_validator("elements")
@classmethod
def _validate_at_most_one_expensive_element(
cls, elements: dict | None
) -> dict | None:
# A field (not model) validator so the constraint also re-runs on
# reassignment (``sp.elements = ...``), since BaseSchema enables
# ``validate_assignment``.
if elements is None:
return elements
# Raw dicts carry their own element_type; schema objects are classified
# via _get_element_type (the same mapping to_config uses).
expensive = [
name
for name, elem in elements.items()
if (
elem.get("element_type")
if isinstance(elem, dict)
else _get_element_type(elem)
)
in _EXPENSIVE_TYPES
]
if len(expensive) > 1:
raise ValueError(
f"SimplePipeline allows at most one DataFit, ArrayDataFit, or Validation "
f"element, but got {len(expensive)}: {expensive}"
)
return elements
