Skip to main content

Failure prediction with XGBoost

At a glance

Effort: ~1–2 hours · You'll build: an engineered tabular feature set and a gradient-boosted classifier with feature importance · Stack: the SDK for data in/out, plus pandas, numpy and xgboost.

Some questions are best answered from tabular features, not raw sequences: will this asset fail in the next week? XGBoost is the workhorse here — it eats engineered features, handles missing values natively, trains fast, and tells you which features drive the prediction. We'll predict an electric submersible pump (ESP) failure seven days out, but the recipe fits any "label the history, predict the window" problem.

Need data to run this?

These steps read series and failure events that already exist. Generate a sandbox first — section E ingests the labelled failure history this build trains on.

New to machine learning?

No background needed. Skim the gentle primer for the ideas in plain language — model, feature, training, and the algorithm itself — and use Generate sample data for a sandbox to run this against.

1. Engineer one feature row per asset per day​

For each pump on each day, summarise the recent week of its sensors into a feature row. Gaps stay as NaN — XGBoost handles missing values without imputation.

import datahub_sdk, numpy as np, pandas as pd

client = datahub_sdk.DataHubClient.from_env()

def feature_row(pump, day):
    def agg(metric, *aggs):
        rf = datahub_sdk.RetrieveFilter(
            ts=f"{pump}_{metric}", start=day - pd.Timedelta(days=7), end=day,
            aggregates=list(aggs), granularity="7d")
        pts = client.timeseries.retrieve_datapoints(rf)[0].get_datapoints()
        return pts[-1] if pts else None
    ip, mt, vb = agg("intake_pressure_bar", "avg", "min"), \
                 agg("motor_temp_c", "avg", "max"), \
                 agg("vibration_mm_s", "avg", "max")
    return {
        "intake_avg": float(ip.average) if ip else np.nan,
        "intake_min": float(ip.min) if ip else np.nan,
        "motor_avg":  float(mt.average) if mt else np.nan,
        "motor_max":  float(mt.max) if mt else np.nan,
        "vib_avg":    float(vb.average) if vb else np.nan,
        "vib_max":    float(vb.max) if vb else np.nan,
    }

2. Label from the failure history​

A row is positive if the pump failed within the next 7 days. Pull historical failure events and label each (pump, day) accordingly.

# failures: DataFrame of (pump, failed_on), pulled from /events (type = esp_failure)
failures = pd.DataFrame([
    {"pump": e.metadata.get("pump"), "failed_on": pd.to_datetime(e.event_time, utc=True)}
    for e in client.events.filter(datahub_sdk.EventFilter(
        datahub_sdk.BasicEventFilter(type="esp_failure"), limit=1000))])

# one feature row per pump per day across the history you have
training_grid = [(p, d) for p in ["pump_esp_a12"]
                 for d in pd.date_range(end=pd.Timestamp.now(tz="UTC"), periods=260, freq="1d")]

def label(pump, day, failures):
    upcoming = failures[(failures["pump"] == pump) &
                        (failures["failed_on"] > day) &
                        (failures["failed_on"] <= day + pd.Timedelta(days=7))]
    return int(len(upcoming) > 0)

rows = [(feature_row(p, d), label(p, d, failures)) for p, d in training_grid]
X = pd.DataFrame([r for r, _ in rows])
y = np.array([lab for _, lab in rows])

3. Train — with class imbalance and early stopping​

Failures are rare, so weight the positive class and optimise for ranking (aucpr). Early stopping on a validation split avoids over-fitting.

from xgboost import XGBClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import average_precision_score

Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.2, stratify=y, random_state=0)

clf = XGBClassifier(
    n_estimators=600, max_depth=5, learning_rate=0.03,
    subsample=0.8, colsample_bytree=0.8,
    eval_metric="aucpr", early_stopping_rounds=40,
    scale_pos_weight=(y == 0).sum() / max((y == 1).sum(), 1))
clf.fit(Xtr, ytr, eval_set=[(Xte, yte)], verbose=False)

print(f"avg precision: {average_precision_score(yte, clf.predict_proba(Xte)[:, 1]):.3f}")

4. Score live, with the reasons attached​

Score today's row per pump, write the risk back as a series, and raise a failure_predicted event for high-risk pumps — including the feature importances so the planner sees why (rising vibration vs. falling intake pressure point to different fixes).

import_ = dict(sorted(zip(X.columns, clf.feature_importances_),
                      key=lambda kv: kv[1], reverse=True)[:3])

today = pd.DataFrame([feature_row("pump_esp_a12", pd.Timestamp.now(tz="UTC"))])
risk = float(clf.predict_proba(today)[0, 1])

client.timeseries.create([datahub_sdk.TimeSeries(
    external_id="pump_esp_a12_failure_risk", name="ESP A-12 7-day failure risk", unit="score", value_type="float")])
client.timeseries.insert_from_lists(
    timestamps=[pd.Timestamp.now(tz="UTC")], values=[risk], ts="pump_esp_a12_failure_risk")

if risk > 0.6:
    client.events.create([datahub_sdk.Event(
        external_id=f"failure_predicted_a12_{int(pd.Timestamp.now().timestamp())}",
        type="failure_predicted", status="open",
        event_time=pd.Timestamp.now(tz="UTC"),
        metadata={"pump": "pump_esp_a12", "risk": f"{risk:.2f}",
                  "top_drivers": ", ".join(import_)})])

Where to take it further​

  • Real explainability. Use SHAP for per-prediction attributions, not just global importance — the planner sees what drove this pump's score.
  • Monotonic constraints. Tell XGBoost that higher vibration only ever raises risk (monotone_constraints) for a model that matches engineering intuition.
  • Compare. Benchmark against Random Forest and the sklearn gradient booster in fraud classification.

Further reading​

See also​