C04 · Customer churn as survival — Bayesian Modelling in Practice

The brief

The retention team of a telecom company has budget for a retention offer and two questions:

How long do our customers last?
Which of our current customers should get the offer?

A colleague has already answered the first one on a slide: "churned customers stayed 18 months on average". You suspect that number is badly wrong, because three quarters of the customers in the file have not churned at all - and they are not missing data, they are the customers who last longest.

How this notebook works

Each task states what to deliver, not how. Write your code in the YOUR CODE HERE cells.
Stuck? h.hint("task2") reveals hints one level at a time: nudge → approach → code skeleton. Try to get by on nudges.
h.check("task2", k_shared=...) compares your numbers with the reference solution.
A full worked solution lives in notebooks/solutions/. Open it only when you are done (or truly stuck).
Work through example E06 first: this challenge assumes you know what pm.Censored does.

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
from scipy import stats

from pymc_challenges import Hints, data

RANDOM_SEED = 7043
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")

h = Hints("C04")
h.tasks()

C04 - Customer churn: most of your customers have not left yet

task0 - Frame churn as a time-to-event problem (3 hints, 2 check(s))
task1 - Put a number on the colleague's mistake (3 hints, 3 check(s))
task2 - A censored Weibull regression (3 hints, 2 check(s))
task3 - Design a check the model can fail (3 hints, 1 check(s))
task4 - Diagnose, fix, compare (3 hints, 3 check(s))
task5 - Conditional churn risk for active customers (3 hints, 2 check(s))
task6 - Who gets the offer? (3 hints, 2 check(s))

The data

One row per customer, taken at a single snapshot date. tenure is the number of months the customer has been (or was) with the company, Churn says whether they have left. The other columns describe the contract, the services booked and how the customer pays.

data.describe("telco")
telco = data.load("telco")
telco[["customerID", "tenure", "Churn", "Contract", "InternetService", "PaymentMethod",
       "MonthlyCharges", "TotalCharges"]].head()

telco: Tenure (months), churn flag and contract details for 7043 customers.
  source: IBM sample data, Telco customer churn.
  url:    https://raw.githubusercontent.com/IBM/telco-customer-churn-on-icp4d/master/data/Telco-Customer-Churn.csv

	customerID	tenure	Churn	Contract	InternetService	PaymentMethod	MonthlyCharges	TotalCharges
0	7590-VHVEG	1	No	Month-to-month	DSL	Electronic check	29.85	29.85
1	5575-GNVDE	34	No	One year	DSL	Mailed check	56.95	1889.5
2	3668-QPYBK	2	Yes	Month-to-month	DSL	Mailed check	53.85	108.15
3	7795-CFOCW	45	No	One year	DSL	Bank transfer (automatic)	42.30	1840.75
4	9237-HQITU	2	Yes	Month-to-month	Fiber optic	Electronic check	70.70	151.65

Task 0 · Frame churn as a time-to-event problem

Deliver

A clean analysis table. Two things need a decision you can defend in one sentence each: TotalCharges does not parse as a number, and some customers have tenure == 0.
The framing: what is the "event", what is the "time", which rows are censored and at what value? What share of customers is censored - overall and per contract type?
Kaplan-Meier survival curves per contract type in one figure. What do they tell you beyond "customers on long contracts churn less"? How far out can each curve be trusted?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Look at the rows where TotalCharges is blank. What else do they have in common? Then ask what a customer who has been observed for zero months can tell a survival model.

For the framing, think of E06: a patient who was alive when last seen is not a missing value and not a death. Who plays that role here, and what is their "follow-up time"?

2 Approach

pd.to_numeric(..., errors="coerce") turns the blanks into NaN so you can inspect them. A row censored at t = 0 contributes log S(0) = 0 to the likelihood - decide accordingly.

Event = Churn == "Yes", time = tenure; customers with Churn == "No" are right-censored at their current tenure.

Kaplan-Meier: at each distinct event time t_j, multiply survival by 1 - d_j / n_j, where d_j is the number of churns at t_j and n_j the number of customers with tenure >= t_j (churned or not). Compute it per contract type, plot with ax.step(..., where="post"), and also count how many customers are still at risk late on.

3 Skeleton

def kaplan_meier(time, event):
    t_event = np.unique(time[event])
    at_risk = np.array([(time >= t).sum() for t in t_event])
    events = np.array([((time == t) & ___).sum() for t in t_event])
    return np.r_[0, t_event], np.r_[1, np.cumprod(___)]

for name, grp in df.groupby("Contract"):
    ax.step(*kaplan_meier(grp.tenure.values, grp.churned.values), where="post", label=name)

To read the curve at a given month: surv[np.searchsorted(knots, t, side="right") - 1].

Self-checks in the notebook (h.check("task0", …)):

censored_share — share of customers (after cleaning) who have NOT churned, about 0.73
km_m2m_12 — Kaplan-Meier survival of month-to-month customers at 12 months, about 0.70

Task 1 · Put a number on the colleague's mistake

"Churned customers stayed 18 months on average, so a customer lasts about a year and a half."

Deliver

Reproduce the number. Then give the approach its best shot: fit an intercept-only Weibull model to the churned customers alone and report the median customer lifetime and the probability that a new customer is still there after 24 months.
The same model and the same two numbers, fitted to all customers with the censoring handled properly.
Both survival curves on top of the overall Kaplan-Meier estimate. By what factor is the naive median off? And of your two "correct" numbers, which one rests on data and which is an extrapolation?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Which customers are not in the colleague's average, and are their lifetimes longer or shorter than 18 months? "Fit only the churners" is an ordinary, uncensored Weibull fit. The proper version differs by one ingredient from E06.

2 Approach

Parameterise the Weibull scale as exp(a) and keep the shape k positive (LogNormal). Naive fit: observed = tenure of churners, no bounds. Proper fit: all customers, pm.Censored(..., lower=None, upper=upper) with upper = tenure for active customers and np.inf for churners.

Weibull facts: S(t) = exp(-(t / scale) ** k), median = scale * log(2) ** (1 / k). Compute both per posterior draw. Then compare the median with the length of the observation window before you put it on a slide.

3 Skeleton

upper = np.where(churned, np.inf, tenure)

def fit(time, upper):
    with pm.Model():
        a = pm.Normal("a", 4.0, 1.5)
        k = pm.LogNormal("k", 0, ___)
        pm.Censored("tenure", pm.Weibull.dist(alpha=k, beta=pm.math.exp(a)),
                    lower=None, upper=___, observed=time)
        return pm.sample(random_seed=RANDOM_SEED)

naive = fit(tenure[churned], np.full(churned.sum(), np.inf))
proper = fit(tenure, upper)
post = az.extract(proper)
scale, k = np.exp(post["a"].values), post["k"].values
median = scale * np.log(2) ** (1 / k)
s24 = np.exp(-(___) ** k)

Self-checks in the notebook (h.check("task1", …)):

naive_median — posterior mean of the median lifetime (months) from the churners-only Weibull, about 10
naive_s24 — posterior mean of S(24) from the churners-only Weibull, about 0.25
censored_s24 — posterior mean of S(24) from the censored intercept-only Weibull, about 0.79 (close to Kaplan-Meier)

Task 2 · A censored Weibull regression

Build an accelerated-failure-time model for tenure: a censored Weibull whose log-scale is linear in contract type, internet service, payment method and monthly charges, with a single shape parameter.

MonthlyCharges is largely determined by the internet service (check this). Enter it as the deviation from the average charge for the customer's internet service type, in units of $20, so that its coefficient reads "paying more than others with the same service".

Deliver

The model with named dimensions and index variables (no dummy matrices). Think about how to keep three categorical predictors identified.
Priors you can defend, with a prior predictive check on the scale that matters: survival curves.
Clean diagnostics and a table of time ratios with 94% intervals: one-year and two-year contracts relative to month-to-month, plus the other covariates.
One sentence: do you believe the two-year time ratio?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

In an AFT model the covariates act on the log of the time scale, so a coefficient b means "lifetimes are multiplied by exp(b)". With three categorical predictors, ask yourself: if every one of them had a free level per category, could the model tell a shift in one from the opposite shift in another?

2 Approach

pd.factorize(df.Contract, sort=True) gives integer codes plus the labels for coords. Give the contract intercept one free level per type (dims="contract") and make the other two factors pm.ZeroSumNormal(..., dims=...): effects relative to the average category, no reference levels to choose, and the intercept stays identified.

Linear predictor: `mu = a[contract_idx] + b_internet[internet_idx] + b_payment[payment_idx]

b_charges * charges_dev; likelihood pm.Censoredaroundpm.Weibull.dist(alpha=k, beta=pm.math.exp(mu))with the sameupper` as in Task 1.

Prior predictive on the right scale: draw a and k from the prior and plot exp(-(t / exp(a)) ** k) over 0-72 months. Time ratios are exp of differences of posterior draws, e.g. a[Two year] - a[Month-to-month].

3 Skeleton

contract_idx, contracts = pd.factorize(df.Contract, sort=True)
charges_dev = (df.MonthlyCharges - df.groupby("InternetService").MonthlyCharges.transform("mean")) / 20
coords = {"contract": contracts, "internet": internets, "payment": payments, "customer": df.customerID}

with pm.Model(coords=coords) as shared_model:
    a = pm.Normal("a", 4.0, 1.5, dims="contract")
    b_internet = pm.ZeroSumNormal("b_internet", 1.0, dims=___)
    b_payment = pm.ZeroSumNormal("b_payment", 1.0, dims=___)
    b_charges = pm.Normal("b_charges", 0, 0.5)
    k = pm.LogNormal("k", 0, 0.75)
    mu = a[contract_idx] + ___
    pm.Censored("tenure", pm.Weibull.dist(alpha=k, beta=pm.math.exp(mu)),
                lower=None, upper=upper, observed=tenure, dims="customer")

post = shared_idata.posterior
ratio = np.exp(post["a"].sel(contract="Two year") - post["a"].sel(contract=___))

Self-checks in the notebook (h.check("task2", …)):

k_shared — posterior mean of the single Weibull shape, about 0.86 (a hazard that falls with tenure)
time_ratio_two_year — posterior mean time ratio of two-year vs month-to-month contracts in the SHARED-shape model - implausibly large, about 30

Task 3 · Design a check the model can fail

az.plot_ppc_dist on tenure is close to meaningless here (why?). Design a posterior predictive check that respects the censoring and looks at the model by contract type.

Deliver

The figure, with posterior uncertainty.
A verdict per contract type: where, in which direction and by how much does the model miss?
The model's predicted 12-month survival for the customers on one-year contracts, next to the model-free estimate.

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

The observed tenure of an active customer is not a lifetime, it is "time on the books so far". What would a histogram of replicated lifetimes be compared with? Which model-free estimate from Task 0 does respect censoring - and what is the model's counterpart of it for a whole group of customers with different covariates?

2 Approach

Kaplan-Meier estimates the survival function of the customer mix within a contract type. The model's counterpart: compute S_i(t) for every customer in the group from their own linear predictor, and average over customers - separately for each posterior draw, so you get a band. (The curve of an "average customer" is not the same thing.)

Broadcasting does the work: post["a"].values[contract_idx[rows]] has shape (customers, draws). A few hundred draws from az.extract(idata, num_samples=...) and a grid of 30-40 months are plenty. Give each contract type its own panel and y-range, or the two-year curve will hide at the top of the plot.

3 Skeleton

def customer_survival(idata, rows, t, n_draws=400):
    post = az.extract(idata, num_samples=n_draws, random_seed=RANDOM_SEED)
    mu = (post["a"].values[contract_idx[rows]] + post["b_internet"].values[internet_idx[rows]]
          + ___ + charges_dev[rows, None] * post["b_charges"].values)      # (customers, draws)
    return np.exp(-((t / np.exp(mu)) ** post["k"].values))

rows = np.flatnonzero(df.Contract == "One year")
curves = np.stack([customer_survival(shared_idata, rows, t).mean(axis=___) for t in grid], axis=1)
lo, hi = np.quantile(curves, [0.03, 0.97], axis=0)       # band over draws
# overlay ax.step(*kaplan_meier(...)) for the same customers

Self-checks in the notebook (h.check("task3", …)):

s12_one_year — shared-shape model, 12-month survival averaged over one-year-contract customers, about 0.96 (Kaplan-Meier says 0.99)

Task 4 · Diagnose, fix, compare

Deliver

A diagnosis of the misfit in terms of the hazard: what is the model structurally unable to express? (Task 0 already showed you the answer.)
A fix to the model - not to the sampler, and not more covariates - and the refit with clean diagnostics.
At least one other survival family with the same fix applied.
A PSIS-LOO comparison of all your models, and your calibration check from Task 3 for the contenders. Do the two agree about which model is best?
What happened to the contract time ratios from Task 2, and why?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Go back to your Kaplan-Meier figure. For which customers does the risk of leaving fall as tenure grows, and for which does it rise? Now look at the Weibull hazard, h(t) ~ t ** (k - 1): what does a single k force all customers to share? Which group has most of the events, and therefore decides k?

2 Approach

AFT covariates can only stretch the time axis; they cannot turn a falling hazard into a rising one. Let the shape vary by contract: k = pm.LogNormal("k", 0, 0.75, dims="contract") and use k[contract_idx] in the likelihood. Nothing else changes.

Second family: pm.LogNormal.dist(mu=mu, sigma=sigma[contract_idx]) inside the same pm.Censored; its survival function is scipy.stats.norm.sf((log(t) - mu) / sigma).

LOO: pm.compute_log_likelihood(idata) inside each model context, then az.compare({...}); az.loo(idata).elpd gives a single model's value. A censored row's pointwise log-likelihood is simply log S(t_i), so nothing special is needed. LOO rewards the density at the observed event times, and most events are early month-to-month churn - keep that in mind when LOO and your calibration plot disagree.

Memory, if you are on a laptop: with 7032 customers and 4000 draws each pointwise log-likelihood is 225 MB, and az.loo(..., pointwise=True) keeps an equally large weight matrix on its result. Three models held that way cost almost 3 GB. Compute loo = az.loo(idata, pointwise=True) once per model, then loo.log_weights = None and del idata["log_likelihood"], and give az.compare the dict of LOO results.

3 Skeleton

with pm.Model(coords=coords) as group_model:
    ...                                                   # as in Task 2
    k = pm.LogNormal("k", 0, 0.75, dims=___)
    pm.Censored("tenure", pm.Weibull.dist(alpha=k[___], beta=pm.math.exp(mu)),
                lower=None, upper=upper, observed=tenure, dims="customer")
    group_idata = pm.sample(random_seed=RANDOM_SEED)
    pm.compute_log_likelihood(group_idata)

az.compare({"shared": shared_idata, "by contract": group_idata, "log-normal": lognormal_idata})
elpd_gain = az.loo(group_idata).elpd - az.loo(shared_idata).elpd

Self-checks in the notebook (h.check("task4", …)):

k_month_to_month — posterior mean Weibull shape for month-to-month customers, about 0.8 (falling hazard)
k_one_year — posterior mean Weibull shape for one-year contracts, about 1.9 (rising hazard)
elpd_gain — elpd_loo of the Weibull with contract-specific shapes minus that of the shared-shape Weibull, roughly +90

Task 5 · From lifetimes to "who is at risk now?"

The retention team cannot act on customers who already left. For every active customer they need the probability of churning within the next 6 months, given that the customer has stayed until today.

Deliver

The formula, and its posterior for every active customer under your final model.
The average 6-month churn probability per contract type.
A figure of this risk against current tenure, by contract type, for the final model and for the shared-shape model of Task 2. What did the fix change about who looks risky?
A sanity check of this very quantity against the data: Kaplan-Meier can estimate "churn within 6 months given survival to $t$" too. Use it to decide which of your Task 4 contenders to carry into the decision.

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

A customer with 40 months of tenure is not a new customer. Write P(T <= t + 6 | T > t) using nothing but the survival function. When would the answer not depend on t at all?

2 Approach

P(churn within 6 months | active at t) = 1 - S(t + 6) / S(t), evaluated per customer at their own current tenure and per posterior draw - both survival terms from the same draw, so the uncertainty in the ratio is right. Only Churn == "No" customers.

The shape parameter decides whether this risk falls (k < 1) or rises (k > 1) with tenure - compare the shared-shape and the fixed model for contract customers.

Model-free counterpart for a contract type: 1 - KM(t + 6) / KM(t). Compare it with the same ratio computed from the group-averaged model survival curves of Task 3, at several t. This is calibration on the scale the decision will use.

3 Skeleton

active = np.flatnonzero(~churned)
s_now = customer_survival(idata, active, tenure[active][:, None], n_draws=1000)     # (customers, draws)
s_later = customer_survival(idata, active, tenure[active][:, None] + 6, n_draws=1000)
p6 = 1 - ___ / ___
pd.Series(p6.mean(axis=1)).groupby(df.Contract.values[active]).mean()

km_cond = 1 - km_at(grp.tenure, grp.churned, t0 + 6) / km_at(grp.tenure, grp.churned, t0)

Self-checks in the notebook (h.check("task5", …)):

p6_month_to_month — average over ACTIVE month-to-month customers of the posterior mean 6-month churn probability, about 0.12
p6_one_year — the same for active one-year-contract customers, about 0.02-0.03 (contract-specific shape model)

Task 6 · Who gets the offer?

The offer on the table, with the retention team's working assumptions:

OFFER_COST = 40.0        # dollars per customer who receives the offer
OFFER_EFFECT = 0.30      # assumed: the offer cuts the 6-month churn probability by 30% (relative)
VALUE_MONTHS = 12        # a customer saved from churning is worth 12 months of their MonthlyCharges

Deliver

The posterior of the expected net value of making the offer, for every active customer.
A ranked top-10 list with the uncertainty a manager needs to see.
A targeting rule, the number of customers it selects, and the total expected profit with an interval. Compare it with the blanket policy "offer it to every month-to-month customer".
Robustness: how much does the target list change under the runner-up model from Task 4? And which assumption would you most want to test before spending real money?
Three sentences for the retention team.

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

What does the company gain if the offer goes to a customer who would have stayed anyway? And to one who leaves regardless? Write the expected net value of one offer as a function of that customer's churn probability and monthly charges. Which customers can never be worth it, whatever their risk?

2 Approach

net_i = OFFER_EFFECT * p6_i * VALUE_MONTHS * MonthlyCharges_i - OFFER_COST, computed per posterior draw, so each customer has a posterior of net value. Rank by its mean; report an interval and P(net > 0) per customer. A risk-neutral rule targets everyone with positive posterior mean.

Total profit of a list is the sum over its customers within each draw (the customers share parameters, so their uncertainties are correlated - do not add up intervals). A cumulative-profit curve over the ranked list shows the optimum and what over-targeting costs. For robustness, recompute the list with the runner-up model's p6.

3 Skeleton

charges = df.MonthlyCharges.values[active][:, None]
value = OFFER_EFFECT * p6 * VALUE_MONTHS * charges - ___          # (customers, draws)
order = np.argsort(-value.mean(axis=1))
target = value.mean(axis=1) > 0
profit = value[target].sum(axis=___)                               # one total per draw
np.quantile(profit, [0.03, 0.97])
cumulative = value[order].cumsum(axis=0)                           # profit of targeting the top n

Self-checks in the notebook (h.check("task6", …)):

n_target — number of active customers whose posterior mean net value is positive - a few hundred, not thousands
profit — posterior mean total net value (dollars) of targeting exactly those customers, roughly 5,000

Going further

Month-to-month churn is a bathtub: a spike in month one, a long quiet period, and a rise again after four years. Neither family here can express that. Try a piecewise-constant hazard (a piecewise exponential model is a Poisson regression in disguise) and rerun the Task 5 check.
tenure is recorded in whole months, so a "churn at 1" is really a churn somewhere in (0, 1]. Treat events as interval-censored - pm.Censored will not do it, but a pm.Potential or pm.CustomDist on $S(t-1) - S(t)$ will. Does the first-month spike look different?
The offer's effect was assumed. Put a prior on it (say, a Beta centred on 0.3), push it through the decision, and work out how large a randomised pilot would have to be before the target list stops changing.
The file is a snapshot taken on one date, not a cohort followed from sign-up. What would you need to know about how customers entered the file to rule out left-truncation, and how would it enter the likelihood if you knew each customer's sign-up date?

h.progress()

C04: 0/21 hints used, 15/15 checks passed.

Reference solution

A complete, executed solution with commentary on every task — including the ways the obvious model fails. Try the hint ladder first.

Open the worked solution →

#The brief

#How this notebook works

#The data

#Task 0 · Frame churn as a time-to-event problem

#Task 1 · Put a number on the colleague's mistake

#Task 2 · A censored Weibull regression

#Task 3 · Design a check the model can fail

#Task 4 · Diagnose, fix, compare

#Task 5 · From lifetimes to "who is at risk now?"

#Task 6 · Who gets the offer?

#Going further

Reference solution

The brief

How this notebook works

The data

Task 0 · Frame churn as a time-to-event problem

Task 1 · Put a number on the colleague's mistake

Task 2 · A censored Weibull regression

Task 3 · Design a check the model can fail

Task 4 · Diagnose, fix, compare

Task 5 · From lifetimes to "who is at risk now?"

Task 6 · Who gets the offer?

Going further