Basic Training Tutorial#

This tutorial builds a complete ModularML workflow for battery state-of-health (SOH) estimation using real battery pulse data. The workflow starts with paired charge/discharge voltage traces, packages them as a FeatureSet, trains a small neural network through a ModelGraph, evaluates the workflow with 3-fold cross-validation, and plots SOH predictions.

This notebook serves as an end-to-end example of how the main ModularML pieces fit together. Focused notebooks in How-To Guides explain individual components in more detail.

Scenario#

Battery pulse tests measure voltage response during controlled charge and discharge events. Those short voltage traces contain information about degradation, so they are a useful input for rapid SOH estimation.

Each sample in this tutorial contains an aligned pair of charge and discharge voltage traces. The model predicts soh, expressed as a percentage. The model is a small, two-layer multi-layer perceptron (MLP) with 8 hidden units per layer.

%matplotlib inline

from pathlib import Path
import os

# If running from the project root, move into this notebook's directory
if (Path.cwd() / "docs" / "tutorials").exists():
    os.chdir(Path.cwd() / "docs" / "tutorials")

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import gaussian_kde
from sklearn.preprocessing import MinMaxScaler

from modularml import (
    AppliedLoss,
    CVBinding,
    CrossValidation,
    EvalPhase,
    Experiment,
    FeatureSet,
    Loss,
    ModelGraph,
    ModelNode,
    Optimizer,
    TrainPhase,
)
from modularml.models.torch import SequentialMLP
from modularml.samplers import SimpleSampler

from utils.battery_pulse_data import get_dataset

Setup#

We will use real battery aging data, as presented in Fine-tuning for rapid capacity estimation of lithium-ion batteries. The helper below downloads and caches the raw NMC pulse data if needed, then returns a cleaned dictionary ready for FeatureSet.from_dict().

dataset = get_dataset(
    chemistry="NMC",
    save_dir=Path("data/01_basic_training"),
)
print("Dataset keys:")
for key, value in dataset.items():
    print(f"  {key:<22} shape={np.shape(value)}")

exp = Experiment(label="basic_training", registration_policy="overwrite")

print(f"Samples:              {len(dataset['soh'])}")
print(f"Charge voltage shape: {dataset['chg_voltage'].shape}")
print(f"Discharge shape:      {dataset['dchg_voltage'].shape}")
print(f"SOH range:            [{dataset['soh'].min():.1f}, {dataset['soh'].max():.1f}] %")

The cleaned dataset includes voltage traces, current profiles, SOH targets, and metadata. For this basic tutorial, we keep the model inputs narrow:

  • features: chg_voltage and dchg_voltage

  • target: soh

  • tags: cell identity, cycling group, RPT/cycle metadata, and SOC metadata

FeatureSet stores these feature/target/tag associations, providing unique identifieres to each data sample. Combined with historical records of data splitting and data scaling, the FeatureSet class provides full serailization and traceability of loaded data to model outputs.

feature_keys = ["chg_voltage", "dchg_voltage"]
target_key = "soh"
tag_keys = [
    "cell_id",
    "group_id",
    "rpt",
    "num_cycles",
    "num_cycles_cumulative",
    "expected_soc",
    "true_soc",
]

fs = FeatureSet.from_dict(
    label="BatteryPulseSOHData",
    data=dataset,
    feature_keys=feature_keys,
    target_keys=target_key,
    tag_keys=tag_keys,
)

print(fs)
print(f"Feature shapes: {fs.get_feature_shapes()}")
print(f"Target shapes:  {fs.get_target_shapes()}")

We split by cell_id so samples from the same physical cell do not leak across train, validation, and test sets.

The split happens in two stages. First, we reserve a held-out test split. Then we split the remaining source data into initial train and val splits. Cross-validation will later rotate the pooled train/validation samples while leaving the test split fixed for every fold.

fs.clear_splits()

fs.split_random(
    ratios={"source": 0.8, "test": 0.2},
    group_by="cell_id",
    seed=13,
)

fs.get_split("source").split_random(
    ratios={"train": 2 / 3, "val": 1 / 3},
    group_by="cell_id",
    seed=13,
)

for split_name in ["train", "val", "test"]:
    view = fs.get_split(split_name)
    cells = np.unique(view.get_tags(fmt="numpy", tags="cell_id"))
    print(f"{split_name:>5}: {len(view):>3} samples from {len(cells):>2} cells")

fs.visualize()

Preprocessing is recorded directly on the FeatureSet.

Both voltage features are normalized with MinMaxScaler. Because chg_voltage and dchg_voltage form a combined (samples, 2, 101) block, merged_axes=(1, 2) tells ModularML how to flatten the feature block for scaler fitting and then restore the original feature shapes. The SOH target is also min-max scaled for training.

During cross-validation, ModularML replays these scaler records inside each fold context, fitting them to that fold’s training split.

fs.undo_all_transforms()

fs.fit_transform(
    scaler=MinMaxScaler(),
    domain="features",      # apply to all features in FeatureSet
    keys=feature_keys,
    fit_to_split="train",   # fit scaler only to train features
    merged_axes=(1, 2),
)
fs.fit_transform(
    scaler=MinMaxScaler(),
    domain="targets",       # apply to all targets in FeatureSet
    keys=target_key,
    fit_to_split="train",   # fit scaler only to train targets
)

soh_raw = fs["train"].get_targets(fmt="numpy", targets=target_key, rep="raw")
soh_scaled = fs["train"].get_targets(
    fmt="numpy",
    targets=target_key,
    rep="transformed",
)
print(f"Raw train SOH:    [{soh_raw.min():.2f}, {soh_raw.max():.2f}] %")
print(f"Scaled train SOH: [{soh_scaled.min():.3f}, {soh_scaled.max():.3f}]")

Model Graph#

The model graph contains a single trainable node. The MLP receives the transformed charge/discharge voltage features and predicts the transformed SOH target.

Even for this simple model, using a ModelGraph is useful: the graph owns the model topology, resolves the FeatureSet reference, validates shapes during build(), and provides a visual representation of the computation path. Any changes to used features does not require changing the MLP definition, the input shape is auto-inferred from the upstream FeatureSet reference.

fs_ref = fs.reference(features=feature_keys, targets=target_key, rep="transformed")

mlp_node = ModelNode(
    label="MLP",
    model=SequentialMLP(output_shape=(1, 1), n_layers=2, hidden_dim=8),
    upstream_ref=fs_ref,
)

graph = ModelGraph(
    label="SOHGraph",
    nodes=[mlp_node],
    optimizer=Optimizer("adam", opt_kwargs={"lr": 1e-3}, backend="torch"),
)

graph.build()
graph.visualize(show_features=True, show_targets=True, show_splits=True)

Training#

A TrainPhase declares the training split, sampler, loss, and number of epochs. Here we use a single training phase with mean-squared error (MSE) on scaled SOH.

We additionally create an EvaluationPhase on each of the train, validation, and test phases. These evaluation results are recorded for each fold, providing automatic tracking and grouping of the model’s performance on each data split.

mse_loss = AppliedLoss(
    loss=Loss("mse", backend="torch"),
    on="MLP",
    inputs=["outputs", "targets"],
)

train_phase = TrainPhase.from_split(
    label="train",
    split="train",
    sampler=SimpleSampler(batch_size=32, shuffle=True, seed=42),
    losses=[mse_loss],
    n_epochs=25,
)

train_phase.visualize()

The experiment execution plan is the exact phase sequence that will run inside each CV fold.

After the train phase, three evaluation phases collect predictions on the fold-specific training split, fold-specific validation split, and fixed held-out test split.

eval_train = EvalPhase.from_split(label="eval_train", split="train", losses=[mse_loss])
eval_val = EvalPhase.from_split(label="eval_val", split="val", losses=[mse_loss])
eval_test = EvalPhase.from_split(label="eval_test", split="test", losses=[mse_loss])

exp.execution_plan.add_phase(train_phase)
exp.execution_plan.add_phase(eval_train)
exp.execution_plan.add_phase(eval_val)
exp.execution_plan.add_phase(eval_test)

print("Execution plan:")
for phase in exp.execution_plan.all:
    print(f"  - {phase.label}")

Cross-Validation#

CrossValidation pools the existing train and val splits, creates three grouped folds, and runs the same experiment plan in each fold. The group_by="cell_id" setting keeps every cell’s samples together within a fold.

The held-out test split is copied unchanged into each fold, so the test panel in every parity plot reflects the same external test set evaluated after a different train/validation fold.

cv = CrossValidation(
    bindings=CVBinding(
        fs=fs,
        source_splits=["train", "val"],
        group_by="cell_id",
    ),
    n_folds=3,
    seed=13,
    experiment=exp,
)

cv_res = cv.run()
print(cv_res)
print(f"Fold labels: {cv_res.fold_labels}")

Evaluation#

The final evaluation results store both model outputs and targets. These built-in result containers provide convenient access of the model predictions. We can easily inverse-transform predictions and targets back to SOH percent using unscale=True.

Each fold produces one parity plot with three panels: training, validation, and test. Points are colored by local KDE density, and each panel reports RMSE in SOH percentage points.

def prediction_arrays(eval_res, node="MLP"):
    true = eval_res.stacked_tensors(
        node=node,
        domain="targets",
        fmt="np",
        unscale=True,
    ).reshape(-1)
    pred = eval_res.stacked_tensors(
        node=node,
        domain="outputs",
        fmt="np",
        unscale=True,
    ).reshape(-1)
    return true, pred

def density_scatter(ax, true, pred):
    try:
        xy = np.vstack([true, pred])
        density = gaussian_kde(xy)(xy)
        order = np.argsort(density)
        sc = ax.scatter(
            true[order],
            pred[order],
            c=density[order],
            s=18,
            cmap="viridis",
            alpha=0.85,
            edgecolors="none",
        )
    except Exception:
        sc = ax.scatter(true, pred, s=18, alpha=0.65, edgecolors="none")
    return sc

def plot_fold_parity(fold_label, fold_results):
    panels = [
        ("Training", fold_results.get_eval_result("eval_train")),
        ("Validation", fold_results.get_eval_result("eval_val")),
        ("Test", fold_results.get_eval_result("eval_test")),
    ]
    arrays = [(name, *prediction_arrays(eval_res)) for name, eval_res in panels]

    all_values = np.concatenate([arr for _, true, pred in arrays for arr in [true, pred]])
    pad = 0.05 * (all_values.max() - all_values.min())
    limits = (all_values.min() - pad, all_values.max() + pad)

    fig, axes = plt.subplots(ncols=3, figsize=(10, 3.4), sharex=True, sharey=True)
    for ax, (name, true, pred) in zip(axes, arrays, strict=True):
        density_scatter(ax, true, pred)
        rmse = np.sqrt(np.mean((pred - true) ** 2))
        ax.plot(limits, limits, "k--", lw=1, alpha=0.7)
        ax.set_title(f"{name}\nRMSE = {rmse:.2f} %", fontsize=10)
        ax.set_xlim(limits)
        ax.set_ylim(limits)
        ax.set_aspect("equal", adjustable="box")
        ax.grid(alpha=0.2)
        ax.set_xlabel("True SOH (%)")
    axes[0].set_ylabel("Predicted SOH (%)")
    fig.suptitle(f"SOH parity: {fold_label}", y=1.04, fontsize=12)
    fig.tight_layout()
    return fig, axes

for fold_label in cv_res.fold_labels:
    fold = cv_res.get_fold(fold_label)
    fig, axes = plot_fold_parity(fold_label, fold)
    plt.show()

What’s Next#

This notebook provides a simple workflow for using ModularML to train and evaluate an MLP for battery SOH estimation. It shows the path through data loading to feature normalization and splitting, model graph construction, experiment training, cross-validation, and post-run visualization.

Additional details on each major component of ModularML can be found in the following how-to notebooks: