Predict Battery Degradation with XGBoost

Train XGBoost on NASA battery data and use SHAP to explain capacity fade.

Introduction

30 Second Summary

Every rechargeable battery you own is slowly dying. Your phone lasts a little less each year, and electric vehicle manufacturers need to predict exactly when a battery pack will cross the safety threshold. The ability to forecast that degradation from sensor data is a high-value skill in energy, manufacturing, and IoT.

In this project, you will build a battery capacity predictor using NASA's lithium-ion cycling dataset. You'll engineer features from raw voltage, current, and temperature readings, train an XGBoost regression model, and use SHAP to explain which physical signals indicate a battery is dying.

What You'll Build

By the end, you'll run a single notebook cell, pass in a battery's cycle data, and get back a predicted remaining capacity alongside a visual explanation of exactly why the model thinks that battery is fading.

By the end of this project, you'll have:

  • A trained XGBoost regression model that predicts remaining battery capacity in amp-hours from engineered sensor features.
  • A degradation curve visualization showing predicted vs actual capacity fade over hundreds of charge cycles, confirming your model tracks real-world battery aging.
  • Interactive SHAP waterfall and beeswarm plots that explain individual predictions and reveal which sensor features matter most across all batteries.
  • Secret Mission: Reframe the problem as binary classification to predict imminent battery failure within 20 cycles.

Are there any prerequisites?

You should be comfortable with Python, pandas, and numpy. Familiarity with basic ML concepts like train/test splits and model evaluation will help you focus on the new material (feature engineering and SHAP) instead of the fundamentals.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters for your portfolio.

Set Up Your Environment and Download Data

Your model needs two things before it can learn: the right libraries and real data. XGBoost provides the gradient boosting algorithm, SHAP provides the explainability layer, and NASA's battery cycling dataset provides hundreds of discharge curves to train on.

In this step, you'll install the libraries, download the dataset, and confirm everything is wired up correctly in your Google Colab notebook.

In this step, get ready to:
  • Install XGBoost and SHAP with pinned versions.
  • Download and extract NASA's lithium-ion battery dataset.
  • Import all libraries and verify the setup.
Install XGBoost and SHAP
  • Install the two ML libraries by running this command in a new code cell:
# Cell 1: Setup and Installation
!pip install xgboost==3.3.0 shap==0.52.0 -q

What does this command do?

The ! prefix runs a shell command from inside a Colab cell. pip install downloads and installs the packages.

The -q flag keeps the output minimal. Pinning versions with == ensures your results are reproducible across sessions.

  • Run the cell by pressing Shift+Enter.
  • You should see a brief confirmation that packages were installed successfully.

Seeing a restart session prompt?

Colab sometimes asks you to restart the runtime after installing packages that conflict with pre-loaded libraries. Click Restart session if prompted. Your cell output will clear, but the packages stay installed.

After restarting, re-run Cell 1 to confirm the install is still active, then continue to the next cell.

My Colab session restarted after installing packages. Do I need to re-run the install cell?

Download the NASA Battery Dataset

NASA's Prognostics Center of Excellence publishes lithium-ion battery cycling data from accelerated aging experiments. The dataset contains 4 batteries charged and discharged until end-of-life.

  • Download and extract the dataset by running this in a new cell:
# Cell 2: Download NASA Battery Dataset
!wget -q https://phm-datasets.s3.amazonaws.com/NASA/5.+Battery+Data+Set.zip
!unzip -q -o "5.+Battery+Data+Set.zip" -d battery_data

What's happening here?

wget -q downloads the zip file from NASA's S3 bucket without showing a progress bar.

unzip -q -o extracts the contents into a folder called battery_data/. The -o flag overwrites existing files if you run the cell again.

  • Run the cell. It takes a few seconds to download the ~35 MB file.
  • Confirm the battery files extracted correctly by running this in a new cell:
!ls battery_data/*.mat
  • You should see B0005.mat, B0006.mat, B0007.mat, and B0018.mat listed.

Not seeing the .mat files?

If the ls command shows no results, the download may have failed silently. Re-run Cell 2 and check that the URL matches exactly.

If you see other files but not the .mat files, they may be nested in a subfolder. Run !find battery_data -name "*.mat" to locate them, then adjust the path in later cells.

The NASA battery dataset download isn't working in my Colab notebook. How do I fix this?

Import your libraries and verify

With the packages installed and data downloaded, you need to import the libraries you'll use throughout the notebook.

  • Create a new cell and add your imports:
# Cell 3: Imports
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import scipy.io
import xgboost as xgb
import shap
from sklearn.metrics import mean_squared_error, r2_score

What does each import do?

  • numpy and pandas handle numerical arrays and tabular data.
  • matplotlib.pyplot creates plots and charts.
  • scipy.io reads MATLAB .mat files, which is NASA's data format.
  • xgboost provides the gradient boosting model you'll train.
  • shap computes feature importance explanations for predictions.
  • mean_squared_error and r2_score from sklearn measure how accurate your model is.
  • Run the cell. It should complete without any output or errors.

Before you run the final check, predict: what version numbers do you expect to see for XGBoost and SHAP based on what you installed?

  • Verify your setup by running this in a new cell:
import xgboost; import shap; print(xgboost.__version__, shap.__version__)
  • You should see 3.3.0 0.52.0 printed. That confirms both packages are installed at the correct versions.

Seeing an ImportError or wrong version?

If you see ModuleNotFoundError, your session may have restarted since the install. Re-run Cell 1 (the pip install cell) first, then try the import again.

If you see a different version number, re-run Cell 1 with the exact pinned versions: xgboost==3.3.0 and shap==0.52.0.

I'm getting an import error for XGBoost or SHAP in Colab even after running pip install. How do I fix this?

✔️ Awesome, I've got everything!

Great. Make sure all three cells are saved in your notebook before moving on.

ⓧ I'd like to double check the full code

Here is everything your notebook should contain so far:

# Cell 1: Setup and Installation
!pip install xgboost==3.3.0 shap==0.52.0 -q

# Cell 2: Download NASA Battery Dataset
!wget -q https://phm-datasets.s3.amazonaws.com/NASA/5.+Battery+Data+Set.zip
!unzip -q -o "5.+Battery+Data+Set.zip" -d battery_data

# Cell 3: Imports
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import scipy.io
import xgboost as xgb
import shap
from sklearn.metrics import mean_squared_error, r2_score

Your environment is locked in and the raw data is ready. Next up, you'll crack open those .mat files, extract the discharge cycles, and see the degradation pattern for yourself.

Parse the Data and Visualize Degradation

You have the raw .mat files sitting in your Colab environment, but they're locked inside a nested MATLAB structure. Before you can predict degradation, you need to see it.

In this step, you'll write a parser that cracks open those files, extract the discharge cycles, and plot the capacity fade curves. Seeing the degradation visually gives you an intuition for what your model needs to learn.

In this step, get ready to:
  • Parse NASA's MATLAB files to extract discharge cycle data.
  • Load all 4 batteries into a Python data structure.
  • Plot the capacity degradation curves and observe the fade pattern.
Write the .mat file parser

NASA stored their battery data in MATLAB .mat files with deeply nested struct arrays. You need a function that navigates this structure, finds each discharge cycle, and pulls out the voltage, current, temperature, and capacity measurements.

  • Create a new code cell in your Colab notebook and paste this parsing function:
# Cell 4: Parse NASA .mat files
def load_battery_data(filepath):
    """Parse NASA PCoE battery .mat file and extract discharge cycles."""
    mat = scipy.io.loadmat(filepath)
    # Get the battery struct - key is the filename without extension
    key = [k for k in mat.keys() if not k.startswith('__')][0]
    battery = mat[key]

    cycles = []
    num_cycles = battery['cycle'][0][0][0].shape[1]

    for i in range(num_cycles):
        cycle_struct = battery['cycle'][0][0][0][0][i]
        cycle_type = str(cycle_struct['type'][0])

        if cycle_type == 'discharge':
            data = cycle_struct['data'][0][0]
            voltage = data['Voltage_measured'][0].flatten()
            current = data['Current_measured'][0].flatten()
            temperature = data['Temperature_measured'][0].flatten()
            time_vec = data['Time'][0].flatten()
            capacity = data['Capacity'][0].flatten()

            if len(capacity) > 0:
                cycles.append({
                    'voltage': voltage,
                    'current': current,
                    'temperature': temperature,
                    'time': time_vec,
                    'capacity': capacity[-1]
                })

    return cycles

What does this code do?

The NASA .mat files have a complex nested structure. This function navigates through it step by step.

  • scipy.io.loadmat(filepath) reads the MATLAB file into a Python dictionary of numpy arrays.
  • The loop iterates over every cycle and filters for 'discharge' cycles only (ignoring charge and impedance cycles).
  • capacity[-1] grabs the final capacity measurement of each discharge. This is the target your model will predict.
Load all battery data

Now use the parser to load all 4 batteries into a dictionary. Each battery gets a list of discharge cycle dictionaries containing the raw sensor arrays and the final capacity value.

  • In the same cell, below the function definition, add this loading code:
# Load all batteries
battery_ids = ['B0005', 'B0006', 'B0007', 'B0018']
all_batteries = {}

for batt_id in battery_ids:
    filepath = f'battery_data/{batt_id}.mat'
    all_batteries[batt_id] = load_battery_data(filepath)
    print(f"{batt_id}: {len(all_batteries[batt_id])} discharge cycles")

What does this code do?

  • battery_ids lists the 4 batteries in NASA's dataset.
  • all_batteries is a dictionary where each key (like 'B0005') maps to a list of discharge cycle dicts with voltage, current, temperature, time, and capacity data.
  • The print statement confirms how many discharge cycles were extracted from each battery.
  • Run the cell by pressing Shift+Enter.

You should see output listing each battery ID followed by its discharge cycle count. Each battery has around 100-170 discharge cycles.

Seeing a KeyError or empty output?

  • If you see a FileNotFoundError, confirm your unzip command from Step 1 created a battery_data/ folder containing the .mat files directly (not nested inside another subfolder).
  • If a battery shows 0 discharge cycles, make sure your function checks for cycle_type == 'discharge' with a lowercase d.

Help me debug my battery .mat file parsing

Plot the capacity degradation curves

With all 4 batteries loaded, you can now visualize how their capacity degrades over time. This plot is your ground truth for what the model needs to learn.

  • Create a new code cell and paste this plotting code:
# Cell 5: Visualize degradation curves
plt.figure(figsize=(10, 6))

for batt_id, cycles in all_batteries.items():
    capacities = [c['capacity'] for c in cycles]
    plt.plot(range(len(capacities)), capacities, label=batt_id)

plt.xlabel('Discharge Cycle Number')
plt.ylabel('Capacity (Ah)')
plt.title('Battery Capacity Degradation Over Cycles')
plt.legend()
plt.axhline(y=1.4, color='r', linestyle='--', label='End-of-Life Threshold (1.4 Ah)')
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

What does this code do?

  • It loops through all 4 batteries and plots each one's capacity (in amp-hours) against its cycle number.
  • The red dashed line at 1.4 Ah marks the end-of-life threshold. NASA defined this as 70% of the initial 2 Ah capacity.
  • Each battery starts near 2.0 Ah and fades downward. The rate of fade is what your model will learn to predict.

Before you run this cell, what shape do you expect the degradation curves to take? Do you think the capacity loss is linear or does it accelerate at some point?

  • Run the cell with Shift+Enter.

You should see a matplotlib plot with 4 colored lines, each starting near 2.0 Ah and dropping toward 1.4 Ah over 100+ cycles. Notice the characteristic "knee" pattern where degradation accelerates later in life.

What is the knee pattern?

Battery degradation is not linear. Capacity drops slowly at first, then hits an inflection point (the "knee") where the loss accelerates sharply. This is caused by internal chemical changes becoming irreversible. Your model needs temporal context to capture this acceleration.

Plot not showing?

  • If you see an empty plot, check that your Cell 4 ran successfully first. The all_batteries variable must be populated before plotting.
  • If you see a NameError for plt, make sure your imports cell (Cell 3) has been run.

Help me debug my degradation curve plot

✔️ Awesome, I've got everything!

Great. Double check you've saved your notebook. You should have Cell 4 (parser + loading) and Cell 5 (plot) both producing output.

ⓧ I'd like to double check the full code

Your notebook should have these cells in order. Scroll up and confirm each is present:

  • Cell 1: !pip install xgboost==3.3.0 shap==0.52.0 -q
  • Cell 2: !wget and !unzip commands for the NASA dataset.
  • Cell 3: All imports (numpy, pandas, matplotlib, scipy.io, xgboost, shap, sklearn.metrics).
  • Cell 4: load_battery_data() function definition followed by the loading loop that populates all_batteries.
  • Cell 5: Degradation curve plot with plt.figure(), the loop plotting each battery, axis labels, and the 1.4 Ah threshold line.

You now have a clear picture of what degradation looks like. Next up, you'll compress those raw voltage, current, and temperature arrays into tabular features that XGBoost can learn from.

Engineer Features from Discharge Curves

Your degradation curves confirm the data is real and loaded correctly. Now you need to transform those raw sensor arrays into something XGBoost can actually learn from.

XGBoost expects tabular input. One row per observation, one column per feature. Each discharge cycle currently exists as thousands of voltage, current, and temperature readings over time. You need to compress each cycle into a handful of meaningful summary statistics.

In this step, get ready to:
  • Define a function that extracts summary statistics from each discharge cycle.
  • Build a single tabular DataFrame from all 4 batteries.
  • Confirm the feature ranges look reasonable and contain no data corruption.
Define the feature engineering function

Each discharge cycle contains thousands of time-series readings. You need to distill each cycle down to a single row of numbers that capture the cycle's signature.

  • Create a new code cell in your notebook (Cell 6) and paste the first part of the feature engineering function:
# Cell 6: Feature engineering from discharge curves
def engineer_features(cycles, battery_id):
    """Extract per-cycle summary features from raw discharge data."""
    rows = []

    for i, cycle in enumerate(cycles):
        voltage = cycle['voltage']
        current = cycle['current']
        temperature = cycle['temperature']
        time_vec = cycle['time']

        # Skip cycles with insufficient data
        if len(voltage) < 10:
            continue

        # Compute time differences for integration
        duration = time_vec[-1] - time_vec[0] if len(time_vec) > 1 else 0

        # Approximate energy: integral of |voltage * current| over time
        if len(time_vec) > 1:
            dt = np.diff(time_vec)
            power = np.abs(voltage[:-1] * current[:-1])
            energy = np.sum(power * dt)
        else:
            energy = 0

What does this code do?

This function loops through every discharge cycle for a given battery and extracts raw sensor arrays. It skips any cycle with fewer than 10 data points (corrupted or incomplete readings).

The energy calculation approximates total energy discharged by integrating the product of voltage and current over time. This is a proxy for how much work the battery did during that cycle.

  • Still in the same cell, paste the feature dictionary and return statement directly below the energy calculation:
        row = {
            'battery_id': battery_id,
            'cycle_number': i,
            'voltage_min': np.min(voltage),
            'voltage_max': np.max(voltage),
            'voltage_mean': np.mean(voltage),
            'voltage_std': np.std(voltage),
            'temp_max': np.max(temperature),
            'temp_mean': np.mean(temperature),
            'discharge_duration': duration,
            'energy': energy,
            'capacity': cycle['capacity']  # Target variable
        }
        rows.append(row)

    return pd.DataFrame(rows)

What features are we extracting?

  • Voltage stats (min, max, mean, std) capture the discharge profile shape. A battery with low voltage_min is being pushed harder.
  • Temperature stats (max, mean) reflect thermal stress. Higher temperatures accelerate degradation.
  • Discharge duration tells you how long the battery lasted. Shorter durations signal capacity loss.
  • Energy is the total work done during the cycle. It correlates directly with remaining capacity.
  • Capacity is the target variable that your model will learn to predict.
  • Run Cell 6 by pressing Shift+Enter.

The cell defines the function without producing any output. No errors means it parsed correctly.

Seeing an IndentationError?

The row dictionary must be indented inside the for loop (two levels of indentation). The return statement must be at the same level as the for loop (one level of indentation inside the function).

help me fix my indentation in engineer_features

Build and inspect the feature DataFrame

The function handles one battery at a time. Now you need to run it across all 4 batteries and combine the results into a single DataFrame.

  • Create a new code cell and paste this code to build the combined DataFrame:
# Build feature DataFrame for all batteries
dfs = []
for batt_id, cycles in all_batteries.items():
    df_batt = engineer_features(cycles, batt_id)
    dfs.append(df_batt)

df = pd.concat(dfs, ignore_index=True)
print(f"Dataset shape: {df.shape}")
print(df.describe())

What does this code do?

  • The loop calls engineer_features() for each battery, producing a small DataFrame per battery.
  • pd.concat() stacks all 4 DataFrames into one, and ignore_index=True resets the row numbers to a clean 0-based sequence.
  • The two print statements show the DataFrame dimensions and summary statistics for every numeric column.

Before you run this cell, think about what you expect to see. How many rows should there be, given that each battery had roughly 100-170 discharge cycles?

  • Run the cell by pressing Shift+Enter.

You should see a shape of approximately (600, 11). That means roughly 600 discharge cycles across all batteries, each represented by 11 columns (8 features plus battery_id, cycle_number, and capacity).

The .describe() output confirms feature ranges look reasonable. Voltage values should sit between roughly 3.0 and 4.2 V. Temperature should be in the 20-40 C range. No NaN counts should appear.

Seeing 0 rows or unexpected NaN values?

  • If the shape shows 0 rows, check that you ran Cell 4 (which populates all_batteries) before this cell.
  • If you see NaN values in .describe(), confirm your energy calculation handles the else: energy = 0 branch correctly.
  • If you see far fewer than 600 rows, the if len(voltage) < 10 guard may be filtering too aggressively. Double check that the threshold is 10, not a larger number.

help me debug my engineer_features output

✔️ Awesome, I've got everything!

Great. Make sure both cells (the function definition and the DataFrame builder) are saved and showing no errors.

ⓧ I'd like to double check the full code

Here is the complete Cell 6 code. Confirm your notebook matches:

# Cell 6: Feature engineering from discharge curves
def engineer_features(cycles, battery_id):
    """Extract per-cycle summary features from raw discharge data."""
    rows = []

    for i, cycle in enumerate(cycles):
        voltage = cycle['voltage']
        current = cycle['current']
        temperature = cycle['temperature']
        time_vec = cycle['time']

        if len(voltage) < 10:
            continue

        duration = time_vec[-1] - time_vec[0] if len(time_vec) > 1 else 0

        if len(time_vec) > 1:
            dt = np.diff(time_vec)
            power = np.abs(voltage[:-1] * current[:-1])
            energy = np.sum(power * dt)
        else:
            energy = 0
        row = {
            'battery_id': battery_id,
            'cycle_number': i,
            'voltage_min': np.min(voltage),
            'voltage_max': np.max(voltage),
            'voltage_mean': np.mean(voltage),
            'voltage_std': np.std(voltage),
            'temp_max': np.max(temperature),
            'temp_mean': np.mean(temperature),
            'discharge_duration': duration,
            'energy': energy,
            'capacity': cycle['capacity']
        }
        rows.append(row)

    return pd.DataFrame(rows)

# Build feature DataFrame for all batteries
dfs = []
for batt_id, cycles in all_batteries.items():
    df_batt = engineer_features(cycles, batt_id)
    dfs.append(df_batt)

df = pd.concat(dfs, ignore_index=True)
print(f"Dataset shape: {df.shape}")
print(df.describe())

You now have a clean tabular dataset where each row represents one discharge cycle and each column captures a measurable physical property. Next up, you will feed these features into an XGBoost model to see how well it can predict remaining capacity.

Train a Baseline XGBoost Model

You have a clean DataFrame of engineered features from the last step. Now it's time to train your first model and see how well it can predict battery capacity from a single discharge cycle's signature.

This first model only sees voltage, temperature, and duration statistics from one cycle in isolation. It has no idea whether it's looking at cycle 10 or cycle 150. You'll discover this limitation when the evaluation reveals a systematic blind spot.

In this step, get ready to:
  • Split data by battery for a realistic train/test evaluation.
  • Train an XGBRegressor and measure its accuracy with RMSE and R-squared.
  • Plot residuals to uncover a systematic prediction bias.
Split the data by battery

Instead of randomly splitting rows, you'll train on three batteries (B0005, B0006, B0007) and test on one held-out battery (B0018). This simulates a real deployment where you predict degradation for a battery the model has never seen.

  • Create a new code cell in your notebook and add the following code to define your feature columns and split the data:
# Cell 7: Train baseline XGBoost (no temporal features)
feature_cols = ['voltage_min', 'voltage_max', 'voltage_mean', 'voltage_std',
                'temp_max', 'temp_mean', 'discharge_duration', 'energy']

# Split: train on B0005, B0006, B0007; test on B0018
train_df = df[df['battery_id'] != 'B0018'].copy()
test_df = df[df['battery_id'] == 'B0018'].copy()

X_train = train_df[feature_cols]
y_train = train_df['capacity']
X_test = test_df[feature_cols]
y_test = test_df['capacity']

What does this code do?

feature_cols lists the 8 engineered columns you built in the previous step. These are the inputs the model will learn from.

The split uses battery_id to separate train and test sets. B0018 is held out entirely, so the model must generalize to a battery it has never seen during training.

Why split by battery instead of randomly?

A random split would leak information. Adjacent cycles from the same battery have nearly identical features, so the model could "memorize" a battery's trajectory rather than learn general degradation patterns. Splitting by battery forces the model to generalize to a completely new device.

Train the baseline XGBRegressor

XGBoost builds an ensemble of decision trees where each new tree corrects the mistakes of the previous ones. This is called gradient boosting.

  • Add this code to the same cell to train the model and print evaluation metrics:
# Train baseline model
baseline_model = xgb.XGBRegressor(
    n_estimators=100,
    max_depth=4,
    learning_rate=0.1,
    random_state=42
)
baseline_model.fit(X_train, y_train)

# Evaluate
y_pred_baseline = baseline_model.predict(X_test)
rmse_baseline = np.sqrt(mean_squared_error(y_test, y_pred_baseline))
r2_baseline = r2_score(y_test, y_pred_baseline)
print(f"Baseline RMSE: {rmse_baseline:.4f} Ah")
print(f"Baseline R-squared: {r2_baseline:.4f}")

What does this code do?

  • n_estimators=100 means the ensemble uses 100 boosted trees.
  • max_depth=4 limits each tree to 4 levels deep, preventing overfitting.
  • learning_rate=0.1 controls how much each new tree's correction contributes to the final prediction.
  • RMSE measures average prediction error in Ah. R-squared tells you what fraction of capacity variation the model explains (1.0 = perfect).
  • Run the cell.

You should see an R-squared around 0.7 to 0.8. That looks decent on the surface. The model captures most of the capacity variation from discharge features alone.

Seeing an error about feature names or shapes?

  • Make sure Cell 6 (feature engineering) ran successfully before this cell. The df DataFrame must exist in memory.
  • Check that feature_cols matches the column names exactly as they appear in your DataFrame. Run df.columns.tolist() to verify.
  • Help me debug my XGBoost training error.
Plot residuals to reveal the blind spot

An R-squared of 0.7 looks fine as a single number. But aggregate metrics can hide systematic failures. A residual plot shows whether errors are random or patterned.

Before you run the next cell, think about this: do you expect the model's errors to be randomly scattered, or do you think there's a pattern tied to where a battery is in its life?

  • Create a new code cell and add this code to plot the residuals against cycle number:
# Cell 8: Plot residuals to reveal the problem
residuals = y_pred_baseline - y_test.values
plt.figure(figsize=(10, 5))
plt.scatter(test_df['cycle_number'].values, residuals, alpha=0.6)
plt.axhline(y=0, color='r', linestyle='--')
plt.xlabel('Cycle Number')
plt.ylabel('Residual (Predicted - Actual)')
plt.title('Baseline Model Residuals: Systematic Over-Prediction in Late Life')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

What does this code do?

  • residuals = y_pred_baseline - y_test.values computes the error for each prediction. Positive means the model predicted higher than reality.
  • Plotting residuals against cycle number reveals whether errors correlate with battery age. If the model is unbiased, points should scatter evenly above and below the red zero line.
  • Run the cell.

The scatter plot shows a clear upward drift on the right side. For late-life cycles (high cycle numbers), the residuals are consistently positive. The model over-predicts capacity for batteries nearing end-of-life.

This is the intended result. The model has no concept of time. It sees a discharge signature and guesses capacity, but two cycles can have similar voltage curves while being 100 cycles apart in age. Without temporal context, it cannot learn the degradation trajectory.

Why does the model fail on late-life cycles?

A battery at cycle 10 and a battery at cycle 140 can produce similar voltage and temperature profiles during a single discharge. The model has no way to tell them apart. It averages over all training examples with similar discharge signatures, which biases predictions upward for end-of-life batteries where capacity has degraded significantly.

Plot looks flat or shows no pattern?

  • Make sure you're plotting test_df['cycle_number'] on the x-axis, not an index. The cycle number carries the temporal information.
  • Confirm your test set is B0018 only. If train and test batteries are mixed, the pattern won't appear.
  • Help me understand my residuals plot.

✔️ Awesome, I've got everything!

Great. Double check you've saved your notebook.

ⓧ I'd like to double check the full code

Here is Cell 7 (train baseline model):

# Cell 7: Train baseline XGBoost (no temporal features)
feature_cols = ['voltage_min', 'voltage_max', 'voltage_mean', 'voltage_std',
                'temp_max', 'temp_mean', 'discharge_duration', 'energy']

# Split: train on B0005, B0006, B0007; test on B0018
train_df = df[df['battery_id'] != 'B0018'].copy()
test_df = df[df['battery_id'] == 'B0018'].copy()

X_train = train_df[feature_cols]
y_train = train_df['capacity']
X_test = test_df[feature_cols]
y_test = test_df['capacity']
# Train baseline model
baseline_model = xgb.XGBRegressor(
    n_estimators=100,
    max_depth=4,
    learning_rate=0.1,
    random_state=42
)
baseline_model.fit(X_train, y_train)

# Evaluate
y_pred_baseline = baseline_model.predict(X_test)
rmse_baseline = np.sqrt(mean_squared_error(y_test, y_pred_baseline))
r2_baseline = r2_score(y_test, y_pred_baseline)
print(f"Baseline RMSE: {rmse_baseline:.4f} Ah")
print(f"Baseline R-squared: {r2_baseline:.4f}")

And here is Cell 8 (residuals plot):

# Cell 8: Plot residuals to reveal the problem
residuals = y_pred_baseline - y_test.values
plt.figure(figsize=(10, 5))
plt.scatter(test_df['cycle_number'].values, residuals, alpha=0.6)
plt.axhline(y=0, color='r', linestyle='--')
plt.xlabel('Cycle Number')
plt.ylabel('Residual (Predicted - Actual)')
plt.title('Baseline Model Residuals: Systematic Over-Prediction in Late Life')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

You've proven the model's blind spot. It can learn from discharge signatures but it cannot capture the degradation trajectory over time. In the next step, you'll fix this by giving the model temporal context and then use SHAP to explain exactly which features matter most.

Fix with Temporal Features and Explain with SHAP

The residual plot from the previous step told a clear story. Your baseline model consistently over-predicts capacity for late-life cycles because it treats every discharge cycle as an isolated snapshot.

In this step, you'll give the model temporal context by engineering lag and rolling features that encode the degradation trajectory. Once the model works, you'll use SHAP to explain which physical signals actually drive its predictions.

In this step, get ready to:
  • Add temporal features (lag, rolling mean, rate of change) and retrain XGBoost to eliminate the systematic residual bias.
  • Generate a SHAP waterfall plot explaining a single late-life prediction.
  • Create a SHAP beeswarm plot revealing global feature importance across all test predictions.
Add temporal features and retrain the model

The baseline model only sees per-cycle discharge statistics. It has no idea whether a cycle is the 10th or the 150th. Lag features (capacity from previous cycles) and rolling means give the model a window into the degradation trend.

  • Create a new code cell (Cell 9) and define the temporal feature engineering function by pasting this code:
# Cell 9: Add temporal features
def add_temporal_features(df):
    """Add lag and rolling features within each battery."""
    df = df.copy()
    df = df.sort_values(['battery_id', 'cycle_number']).reset_index(drop=True)

    temporal_dfs = []
    for batt_id in df['battery_id'].unique():
        batt_df = df[df['battery_id'] == batt_id].copy()

        # Lag features
        batt_df['capacity_lag_1'] = batt_df['capacity'].shift(1)
        batt_df['capacity_lag_5'] = batt_df['capacity'].shift(5)

        # Rolling features
        batt_df['capacity_rolling_mean_5'] = batt_df['capacity'].shift(1).rolling(5, min_periods=1).mean()
        batt_df['capacity_rolling_mean_10'] = batt_df['capacity'].shift(1).rolling(10, min_periods=1).mean()

        # Rate of change
        batt_df['capacity_diff'] = batt_df['capacity'].diff()

        temporal_dfs.append(batt_df)

    return pd.concat(temporal_dfs, ignore_index=True)

df_temporal = add_temporal_features(df)
df_temporal = df_temporal.dropna()  # Drop rows where lag features are NaN

What does this code do?

This function creates five new features for each discharge cycle, computed within each battery separately so no information leaks across batteries.

  • capacity_lag_1 and capacity_lag_5 store the capacity from 1 and 5 cycles ago. These let the model see where the battery has been.
  • capacity_rolling_mean_5 and capacity_rolling_mean_10 smooth out noise by averaging recent history. The .shift(1) ensures the model never peeks at the current cycle's answer.
  • capacity_diff captures the cycle-over-cycle rate of change. A large negative value means degradation is accelerating.
  • Run the cell.
  • Confirm there are no errors. The .dropna() removes the first few rows of each battery where lag values cannot be computed.

Now you'll define the expanded feature list, re-split the data, and train an improved model. Before you run this, think about what you expect to happen to the RMSE compared to your baseline.

  • In the same cell, below the code you just pasted, add the training code:
# Updated feature set including temporal context
temporal_feature_cols = feature_cols + [
    'cycle_number', 'capacity_lag_1', 'capacity_lag_5',
    'capacity_rolling_mean_5', 'capacity_rolling_mean_10', 'capacity_diff'
]

# Re-split with temporal features
train_temporal = df_temporal[df_temporal['battery_id'] != 'B0018']
test_temporal = df_temporal[df_temporal['battery_id'] == 'B0018']

X_train_t = train_temporal[temporal_feature_cols]
y_train_t = train_temporal['capacity']
X_test_t = test_temporal[temporal_feature_cols]
y_test_t = test_temporal['capacity']

# Train improved model with more trees and regularization
improved_model = xgb.XGBRegressor(
    n_estimators=200,
    max_depth=5,
    learning_rate=0.05,
    subsample=0.8,
    colsample_bytree=0.8,
    random_state=42
)
improved_model.fit(X_train_t, y_train_t)

# Evaluate improvement
y_pred_improved = improved_model.predict(X_test_t)
rmse_improved = np.sqrt(mean_squared_error(y_test_t, y_pred_improved))
r2_improved = r2_score(y_test_t, y_pred_improved)
print(f"Improved RMSE: {rmse_improved:.4f} Ah (was {rmse_baseline:.4f})")
print(f"Improved R-squared: {r2_improved:.4f} (was {r2_baseline:.4f})")

What changed in this model?

The model now sees 14 features instead of 8. The 6 new temporal features encode where the battery is in its degradation arc.

The hyperparameters also shifted: more trees (200 vs 100), a lower learning rate (0.05 vs 0.1), and subsample/colsample_bytree at 0.8 for regularization. This prevents the model from overfitting to the lag features.

  • Run the cell.

You should see the RMSE drop to around 0.02-0.03 Ah (down from the baseline), and R-squared jump close to 0.99. The temporal features gave the model the context it was missing.

RMSE didn't improve?

Make sure your Cell 9 runs after Cell 7 (where feature_cols is defined). The temporal_feature_cols list extends feature_cols, so that variable must already exist.

Also confirm you ran df_temporal = df_temporal.dropna() before the split. Without it, NaN rows from the lag features corrupt training.

Help me debug my temporal feature engineering

The metrics prove the model improved, but you need to verify the systematic bias is gone. Plot the residuals the same way you did for the baseline.

  • Below the training code in the same cell, add the residuals plot:
# Verify residuals are now flat
residuals_improved = y_pred_improved - y_test_t.values
plt.figure(figsize=(10, 5))
plt.scatter(test_temporal['cycle_number'].values, residuals_improved, alpha=0.6)
plt.axhline(y=0, color='r', linestyle='--')
plt.xlabel('Cycle Number')
plt.ylabel('Residual (Predicted - Actual)')
plt.title('Improved Model Residuals: No Systematic Bias')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

What should I see?

The scatter plot should show residuals clustered tightly around the zero line across all cycle numbers. Compare this to the baseline residuals plot from Step 4, where points drifted above zero for late-life cycles. The systematic bias is gone.

  • Run the cell and confirm the residuals are flat with no upward drift in late-life cycles.
Generate SHAP explanations for a single prediction

The model works. Now you need to explain WHY it predicts what it does. SHAP (SHapley Additive exPlanations) assigns each feature a contribution score for every individual prediction. A waterfall plot shows these contributions stacked up, revealing which features pushed a specific prediction higher or lower.

  • Create a new code cell (Cell 10) and generate SHAP values with a waterfall plot:
# Cell 10: SHAP Explainability
explainer = shap.Explainer(improved_model)
shap_values = explainer(X_test_t)

# Waterfall plot for a late-life prediction
late_life_idx = len(X_test_t) - 10  # Near end of battery life
shap.plots.waterfall(shap_values[late_life_idx], max_display=12)

What does this code do?

shap.Explainer(improved_model) creates a TreeExplainer optimized for XGBoost. It uses the tree structure to compute exact Shapley values efficiently.

explainer(X_test_t) computes SHAP values for every test prediction. The waterfall plot then visualizes one specific prediction: a cycle near end-of-life (10 cycles from the end of B0018's test data).

  • Run the cell.

You should see a waterfall chart showing feature contributions stacked from a base value. Features in red pushed the prediction higher. Features in blue pushed it lower. For a late-life cycle, you'll likely see capacity lag features and rolling means pulling the prediction down toward low capacity.

Seeing a JavaScript or rendering error?

SHAP plots sometimes conflict with Colab's rendering backend. If the waterfall doesn't display, add this line before the plot call:

shap.initjs()

Help me fix SHAP plot rendering in Colab

Create a beeswarm plot for global feature importance

The waterfall explained one prediction. A beeswarm plot shows feature importance across ALL test predictions, revealing which features matter most globally and how their values correlate with impact.

  • Create a new code cell (Cell 11) and generate the beeswarm plot:
# Cell 11: Beeswarm plot for global feature importance
shap.plots.beeswarm(shap_values, max_display=14)

How do I read a beeswarm plot?

Each dot is one test prediction. The x-axis is the SHAP value (impact on model output). Color indicates the feature's actual value (red = high, blue = low). Features are sorted by importance from top to bottom.

You'll see cycle_number and the rolling capacity features near the top. This confirms the model relies on temporal degradation context, not just per-cycle discharge shape.

  • Run the cell.
  • Confirm the beeswarm plot renders with features ranked by importance, and that temporal features (cycle_number, rolling capacity means) appear among the top drivers.

Beeswarm plot looks empty or truncated?

Make sure shap_values was computed in the previous cell (Cell 10). If you restarted the runtime, you need to re-run all cells from the top.

Help me debug my SHAP beeswarm plot

✔️ Awesome, I've got everything!

Great. Confirm you've run Cells 9, 10, and 11 and that all three plots rendered (flat residuals, waterfall, beeswarm).

ⓧ I'd like to double check the full code

Scroll through your notebook and confirm these cells are present in order:

  • Cell 9: The add_temporal_features function definition, the call to create df_temporal, the dropna, then the feature list, train/test split, model training, print statements showing improved RMSE and R-squared, and the residuals plot.
  • Cell 10: shap.Explainer(improved_model), shap_values = explainer(X_test_t), late_life_idx, and the waterfall plot call.
  • Cell 11: shap.plots.beeswarm(shap_values, max_display=14) on its own.

Your model now predicts battery capacity with high accuracy and you can explain every prediction down to the feature level. Next up, the Secret Mission challenges you to reframe this as a binary classification problem to predict imminent battery failure.

Secret mission

Predict Imminent Battery Failure

Your regression model predicts capacity, but in the real world, operators need a binary alarm: is this battery about to die? Reframe the problem as classification to predict imminent failure within 20 cycles, handling severe class imbalance along the way.

Clean Up Your Resources

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project runs entirely in Google Colab, so there are no ongoing costs.

Resources you used:

  • Google Colab notebook (persists in your Google Drive).
  • NASA battery dataset files (battery_data/ directory, stored in the ephemeral Colab runtime).
  • Installed Python packages (XGBoost, SHAP) in the Colab runtime.

Keep everything running

No action needed. Choose this if you are still actively building or want to revisit the notebook soon.

  • Your Colab notebook stays in Google Drive and you can reopen it anytime.
  • The Colab runtime will automatically disconnect after a period of inactivity. When it does, downloaded files and installed packages are cleared. You will need to re-run the setup cells next time.

Pause - I'll come back to this later

Disconnect the Colab runtime to free up compute resources while keeping your notebook saved in Google Drive.

  • In your Colab notebook, click Runtime in the top menu bar.
  • Click Disconnect and delete runtime.
  • This clears the downloaded dataset and installed packages from memory. Your notebook code remains saved in Google Drive.
  • When you return, re-run all cells from the top to reinstall packages and re-download the data.

Delete - I don't want to use this again

Remove the notebook from Google Drive and disconnect the runtime entirely.

  • In your Colab notebook, click Runtime in the top menu, then click Disconnect and delete runtime. This removes all downloaded files and installed packages.
  • Open Google Drive at drive.google.com.
  • Navigate to the Colab Notebooks folder (or wherever you saved the notebook).
  • Right-click your battery degradation notebook and select Move to trash.
  • Empty your Google Drive trash to permanently delete the file.

Nice Work!

Nice Work!

You just built a battery capacity degradation predictor using NASA's lithium-ion cycling dataset, trained an XGBoost model that learns degradation trajectories, and explained its predictions with SHAP.

You've learned how to:

  • Engineer tabular features from raw time-series sensor data (voltage, current, temperature) and diagnose model limitations through residual analysis.
  • Design temporal features (lag, rolling mean, rate of change) that let XGBoost learn degradation trends a snapshot model misses.
  • Generate SHAP waterfall and beeswarm plots to explain which physical signals drive battery degradation predictions in a safety-critical domain.
  • Secret Mission: Reframed the problem as binary classification with scale_pos_weight for class imbalance to predict imminent battery failure.

Ready to quiz yourself?