5 - Clustering with K-Means: discovering groups of patients¶
In chapter 5, Lucy asks the team a new kind of question: can we identify groups of patients at risk before they become ill? This time there are no labels. Nobody knows which groups exist, so the models of the previous chapters, which learn from examples with a known answer, cannot be used. Martha proposes unsupervised learning: let the algorithm discover the groups that naturally exist in the data, based on how similar the patients are. As she warns, the algorithm does not understand what a patient is; it only compares numbers with a similarity rule that we choose.
In this notebook:
- Import the libraries.
- K-Means by hand, step by step (Martha's notes, section 5.4).
- Local optima and the initialization: K-Means++.
- The Minermont patient population.
- Why the features must be scaled.
- Choosing $K$: the elbow method.
- The five groups and their clinical interpretation.
- The variables we include decide the groups.
- Assigning a new patient to a group.
Step 1: Import the libraries¶
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import StandardScaler
from sklearn.cluster import KMeans
from sklearn.metrics import adjusted_rand_score
Step 2: K-Means by hand¶
K-Means divides $n$ points into $K$ groups (clusters), each represented by its centroid $\mu_k$. It alternates two steps:
- Assignment: each point goes to the cluster whose centroid is closest (squared Euclidean distance): $$C(i) = \arg\min_k \|x_i - \mu_k\|^2$$
- Update: each centroid is recalculated as the mean of the points assigned to it: $$\mu_k = \frac{1}{|S_k|} \sum_{x_i \in S_k} x_i$$
and repeats them until the centroids stop moving. As Martha points out, the centroids "don't move like objects": we simply recalculate their position after each reassignment. The quantity that K-Means reduces is the sum of squared errors (SSE), or inertia, the sum of the squared distances from each point to its centroid:
$$\text{SSE} = \sum_{k=1}^{K} \sum_{x_i \in S_k} \|x_i - \mu_k\|^2$$
We try it on 150 points in two dimensions, generated around three centers, starting from three arbitrary centroids:
rng = np.random.default_rng(3)
X_toy = np.vstack([rng.normal(center, 0.8, (50, 2)) for center in [[2, 2], [7, 3], [4.5, 7]]])
def sse(X, labels, centroids):
return sum(np.sum((X[labels == k] - c) ** 2) for k, c in enumerate(centroids))
def kmeans(X, centroids, max_iter=50):
"""K-Means by hand. Returns the final centroids, the labels, the SSE of each iteration and the history of centroids."""
history, sses = [centroids.copy()], []
for _ in range(max_iter):
distances = np.sum((X[:, None, :] - centroids[None, :, :]) ** 2, axis=2) # points x centroids
labels = distances.argmin(axis=1) # 1. assignment
sses.append(sse(X, labels, centroids))
new_centroids = np.array([X[labels == k].mean(axis=0) for k in range(len(centroids))]) # 2. update
if np.allclose(new_centroids, centroids):
break
centroids = new_centroids
history.append(centroids.copy())
return centroids, labels, sses, history
start = np.array([[2.0, 7.0], [3.0, 6.0], [4.0, 5.0]])
centroids, labels, sses, history = kmeans(X_toy, start)
print(f"Iterations: {len(sses)}")
print("SSE at each iteration:", [round(float(s), 1) for s in sses])
Iterations: 10 SSE at each iteration: [1633.3, 909.3, 842.9, 827.8, 818.9, 806.2, 737.2, 396.2, 194.5, 194.0]
colors = np.array(['tab:blue', 'tab:orange', 'tab:green'])
steps = [0, 1, 3, len(history) - 1]
fig, axes = plt.subplots(1, 4, figsize=(18, 4.5))
for ax, t in zip(axes, steps):
c = history[t]
lab = np.sum((X_toy[:, None, :] - c[None]) ** 2, axis=2).argmin(axis=1)
ax.scatter(*X_toy.T, c=colors[lab], s=15, alpha=0.6)
path = np.array(history[:t + 1])
for k in range(3):
ax.plot(path[:, k, 0], path[:, k, 1], 'k--', linewidth=0.8)
ax.scatter(*c.T, c='k', marker='X', s=200)
ax.set_title(f'Iteration {t}: SSE {sse(X_toy, lab, c):.1f}')
plt.suptitle('K-Means: the centroids (X) are recalculated after each assignment', y=1.02)
plt.show()
The SSE never goes up: it falls from 1633.3 to 194.0 in 10 iterations, and the algorithm stops when the centroids no longer move. In the plots, the three centroids start on a diagonal and each one travels, iteration after iteration, towards the centre of one of the three groups.
Step 3: Local optima and the initialization¶
Martha's notes add an important warning: K-Means always converges, but it can converge to a local optimum, a stable grouping that is not the best possible one. It depends on where the centroids start. Let's start from three centroids placed side by side at the top of the plot:
bad_start = np.array([[1.0, 7.0], [2.0, 7.0], [3.0, 7.0]])
bad_centroids, bad_labels, bad_sses, _ = kmeans(X_toy, bad_start)
print(f"Good start: final SSE {sses[-1]:.1f}")
print(f"Bad start: final SSE {bad_sses[-1]:.1f}")
fig, axes = plt.subplots(1, 2, figsize=(11, 4.5))
for ax, lab, c, title in [(axes[0], labels, centroids, f'Best solution: SSE {sses[-1]:.1f}'),
(axes[1], bad_labels, bad_centroids, f'Local optimum: SSE {bad_sses[-1]:.1f}')]:
ax.scatter(*X_toy.T, c=colors[lab], s=15, alpha=0.6)
ax.scatter(*c.T, c='k', marker='X', s=200)
ax.set_title(title)
plt.show()
Good start: final SSE 194.0 Bad start: final SSE 680.1
From this start the algorithm also converges, but to a grouping with an SSE of 680.1, three and a half times worse than the best one: two centroids share the group at the bottom left and a single centroid covers the other two groups. From there, no reassignment lowers the SSE, so the algorithm stops.
Two remedies are common. The first is to run the algorithm several times from different random starts and keep the solution with the lowest SSE. The second is K-Means++, which chooses better starting centroids: the first one is a random point, and each of the next ones is chosen at random with probability proportional to $D(x)^2$, the squared distance from each point to the nearest centroid already chosen. Distant points are therefore more likely to be chosen, and the centroids start spread out. We compare 200 random starts of each kind:
def random_start(X, k, rng):
return X[rng.choice(len(X), k, replace=False)]
def kmeans_plus_plus_start(X, k, rng):
chosen = [X[rng.integers(len(X))]]
for _ in range(1, k):
d2 = np.min(np.sum((X[:, None, :] - np.array(chosen)[None]) ** 2, axis=2), axis=1) # D(x)^2
chosen.append(X[rng.choice(len(X), p=d2 / d2.sum())])
return np.array(chosen)
best_sse = sses[-1]
for name, starter, offset in [('random points', random_start, 0), ('K-Means++', kmeans_plus_plus_start, 1000)]:
finals = [kmeans(X_toy, starter(X_toy, 3, np.random.default_rng(offset + s)))[2][-1] for s in range(200)]
stuck = np.sum(np.array(finals) > best_sse + 1)
print(f"Start with {name:<14}: {stuck} of 200 runs end in a local optimum")
sk = KMeans(n_clusters=3, init='k-means++', n_init=10, random_state=0).fit(X_toy)
print(f"\nscikit-learn (K-Means++, best of 10 starts): SSE {sk.inertia_:.1f}")
Start with random points : 19 of 200 runs end in a local optimum Start with K-Means++ : 9 of 200 runs end in a local optimum scikit-learn (K-Means++, best of 10 starts): SSE 194.0
With random starting points, 19 of the 200 runs get stuck; with K-Means++, 9. It does not remove the problem, but it halves it. That is why scikit-learn uses K-Means++ by default and, with n_init=10, keeps the best of 10 runs; its result has the same SSE as our best solution.
Step 4: The Minermont patient population¶
Noah prepares the anonymized data of the adult patients: age, visits, basic lab results and family history. Real records are confidential, so we generate a synthetic population of 2000 patients. The generator creates five hidden profiles; the algorithm will never see which profile each patient came from. We keep that information apart only to check, at the end, something we could never check with real patients.
| Column | Meaning |
|---|---|
age |
Age in years |
visits_last_year |
Visits to the hospital in the last year |
glucose |
Fasting glucose (mg/dL) |
systolic_bp |
Systolic blood pressure (mmHg) |
cholesterol |
Total cholesterol (mg/dL) |
chronic_conditions |
Number of diagnosed chronic conditions |
family_history |
1 if there is a family history of heart disease or type 2 diabetes |
def generate_population(n, seed):
"""Synthetic adult patients of Minermont. Returns the data and, separately, the hidden profile of each patient."""
rng = np.random.default_rng(seed)
# five hidden profiles: (mean, sd) of age, glucose, blood pressure and cholesterol; average visits and chronic conditions;
# probability of family history; share of the population
profiles = [
dict(age=(30, 6), visits=1.0, glucose=(84, 6), sbp=(113, 7), chol=(170, 20), chronic=0.05, family=0.12, share=0.30),
dict(age=(72, 6), visits=7.0, glucose=(130, 15), sbp=(140, 10), chol=(205, 25), chronic=2.5, family=0.35, share=0.22),
dict(age=(45, 13), visits=8.0, glucose=(90, 8), sbp=(118, 9), chol=(185, 22), chronic=0.3, family=0.2, share=0.25),
dict(age=(68, 10), visits=20.0, glucose=(150, 25), sbp=(146, 15), chol=(215, 35), chronic=5.0, family=0.4, share=0.08),
dict(age=(52, 4), visits=1.5, glucose=(106, 4), sbp=(135, 4), chol=(225, 15), chronic=0.1, family=0.75, share=0.15),
]
hidden = rng.choice(len(profiles), n, p=[p['share'] for p in profiles])
cols = {c: np.empty(n) for c in ['age', 'visits', 'glucose', 'sbp', 'chol', 'chronic', 'family']}
for k, p in enumerate(profiles):
m = hidden == k
c = m.sum()
cols['age'][m] = rng.normal(*p['age'], c)
cols['visits'][m] = rng.poisson(p['visits'], c)
cols['glucose'][m] = rng.normal(*p['glucose'], c)
cols['sbp'][m] = rng.normal(*p['sbp'], c)
cols['chol'][m] = rng.normal(*p['chol'], c)
cols['chronic'][m] = rng.poisson(p['chronic'], c)
cols['family'][m] = rng.binomial(1, p['family'], c)
df = pd.DataFrame({
'age': np.clip(np.round(cols['age']), 18, 95).astype(int), 'visits_last_year': cols['visits'].astype(int),
'glucose': np.round(cols['glucose']).astype(int), 'systolic_bp': np.round(cols['sbp']).astype(int),
'cholesterol': np.round(cols['chol']).astype(int), 'chronic_conditions': cols['chronic'].astype(int),
'family_history': cols['family'].astype(int),
})
return df, hidden
patients, hidden_profile = generate_population(2000, seed=42)
features = ['age', 'visits_last_year', 'glucose', 'systolic_bp', 'cholesterol', 'chronic_conditions']
print(patients.describe().round(1).T[['mean', 'std', 'min', 'max']].to_string())
mean std min max age 49.7 18.1 18.0 92.0 visits_last_year 5.7 6.0 0.0 32.0 glucose 105.2 25.5 68.0 219.0 systolic_bp 126.1 15.1 85.0 175.0 cholesterol 194.4 29.8 113.0 309.0 chronic_conditions 1.1 1.8 0.0 12.0 family_history 0.3 0.5 0.0 1.0
We cluster with the six numerical measurements. We leave the family history (a yes/no column) out of the clustering and use it afterwards to describe the groups; in Step 8 we will see why.
Step 5: Why the features must be scaled¶
K-Means uses the squared Euclidean distance, so, as the book warns, it is sensitive to the scale of the features. Cholesterol varies by tens of mg/dL, while the number of chronic conditions varies by one or two units: without scaling, cholesterol would decide almost everything. We standardize every feature (mean 0, standard deviation 1).
Since our data are synthetic, we can measure how well each clustering recovers the hidden profiles with the adjusted Rand index (ARI): 1 means the same grouping, 0 is what we would expect by chance. With real patients this check does not exist, which is precisely why unsupervised learning needs clinical interpretation.
X_raw = patients[features].to_numpy(dtype=float)
X = StandardScaler().fit_transform(X_raw)
for name, data in [('without scaling', X_raw), ('standardized', X)]:
km = KMeans(n_clusters=5, n_init=10, random_state=0).fit(data)
print(f"K-Means with K = 5, {name:<16} agreement with the hidden profiles (ARI) = {adjusted_rand_score(hidden_profile, km.labels_):.3f}")
K-Means with K = 5, without scaling agreement with the hidden profiles (ARI) = 0.408
K-Means with K = 5, standardized agreement with the hidden profiles (ARI) = 0.881
Without scaling the agreement is only 0.408: cholesterol and glucose, the features with the largest numbers, decide the groups almost on their own. With standardized features, K-Means recovers the hidden profiles with an ARI of 0.881.
Step 6: Choosing $K$, the elbow method¶
In real life we do not know how many groups there are. As in the book, we run K-Means for $K$ from 1 to 10 and plot the SSE. It always decreases as $K$ grows (with $K = n$ it would be zero), so the question is where adding clusters stops being worth it: the elbow of the curve. To make the decision less subjective, we also compute a simple automatic rule: the point of the curve that is farthest from the straight line joining its first and last points.
ks = list(range(1, 11))
sse_by_k = [KMeans(n_clusters=k, n_init=10, random_state=0).fit(X).inertia_ for k in ks]
rows = [{'K': k, 'SSE': s, 'reduction vs K-1': (sse_by_k[i - 1] - s) / sse_by_k[i - 1] if i else np.nan}
for i, (k, s) in enumerate(zip(ks, sse_by_k))]
print(pd.DataFrame(rows).set_index('K').round({'SSE': 0, 'reduction vs K-1': 3}).to_string())
# automatic elbow: distance from each point to the line joining the first and the last point
p1, p2 = np.array([ks[0], sse_by_k[0]]), np.array([ks[-1], sse_by_k[-1]])
d = p2 - p1
dist = [abs(d[0] * (p1[1] - s) - (p1[0] - k) * d[1]) / np.linalg.norm(d) for k, s in zip(ks, sse_by_k)]
elbow = ks[int(np.argmax(dist))]
print(f"\nAutomatic elbow: K = {elbow}")
plt.figure(figsize=(8, 4))
plt.plot(ks, sse_by_k, 'o-')
plt.plot([ks[0], ks[-1]], [sse_by_k[0], sse_by_k[-1]], 'k:', linewidth=1)
plt.axvline(elbow, color='gray', linestyle='--')
plt.xlabel('Number of clusters K')
plt.ylabel('SSE (inertia)')
plt.title('The elbow method')
plt.xticks(ks)
plt.grid(True)
plt.show()
SSE reduction vs K-1 K 1 12000.0 NaN 2 6645.0 0.446 3 4968.0 0.252 4 4099.0 0.175 5 3504.0 0.145 6 3303.0 0.057 7 3134.0 0.051 8 2947.0 0.060 9 2808.0 0.047 10 2648.0 0.057 Automatic elbow: K = 3
The SSE falls quickly at first and then more slowly. As in the book, beyond $K = 5$ the improvements become much smaller: each extra cluster only lowers the SSE by 5 or 6%, compared with 14.5% for the fifth one. The automatic rule, however, points to $K = 3$. The elbow is a heuristic, and two reasonable ways of reading the same curve disagree. As Martha proposes, we decide on a reasonable $K$ and then check whether the groups make clinical sense. Let's compare the groups for several values of $K$:
short = ['age', 'visits_last_year', 'glucose', 'systolic_bp', 'cholesterol', 'chronic_conditions']
for k in [3, 4, 5, 6]:
km = KMeans(n_clusters=k, n_init=10, random_state=0).fit(X)
groups = patients.groupby(km.labels_)[short].mean().round(0).astype(int)
groups['patients'] = np.bincount(km.labels_)
print(f"K = {k}")
print(groups.sort_values('age').to_string(index=False), end='\n\n')
K = 3 age visits_last_year glucose systolic_bp cholesterol chronic_conditions patients 35 4 86 114 175 0 984 60 4 111 135 217 1 622 70 13 142 142 207 4 394 K = 4 age visits_last_year glucose systolic_bp cholesterol chronic_conditions patients 35 4 86 114 174 0 955 54 3 103 132 220 0 450 67 20 152 145 216 5 169 72 7 132 139 204 2 426 K = 5 age visits_last_year glucose systolic_bp cholesterol chronic_conditions patients 30 1 84 113 169 0 625 45 8 90 118 188 0 428 53 2 106 135 224 0 355 67 20 152 145 216 5 169 72 7 131 139 204 2 423 K = 6 age visits_last_year glucose systolic_bp cholesterol chronic_conditions patients 30 1 84 113 169 0 625 45 8 90 118 188 0 428 53 2 106 135 224 0 344 67 21 152 146 219 5 155 72 8 133 137 188 3 208 73 7 132 141 218 2 240
With $K = 3$ the silent-risk patients disappear inside a large group of 622 patients aged 60 on average. With $K = 4$ they form the core of a group of 450, but mixed with other patients, which dilutes its averages (3 visits a year, glucose of 103 mg/dL). With $K = 5$ the group appears on its own, with 355 patients. With $K = 6$ nothing new appears: the stable chronic patients are simply split in two by their cholesterol. So $K = 5$ is the choice that is both supported by the curve and clinically meaningful.
Step 7: The five groups¶
We run K-Means with $K = 5$. The cluster numbers that the algorithm returns are arbitrary, so we renumber the clusters by increasing mean age. Then we describe each group with the mean of each feature in the original units, its size and the percentage of patients with a family history, the column we left out of the clustering.
scaler = StandardScaler().fit(X_raw)
kmeans5 = KMeans(n_clusters=5, n_init=10, random_state=0).fit(scaler.transform(X_raw))
order = np.argsort(scaler.inverse_transform(kmeans5.cluster_centers_)[:, 0]) # by mean age
renumber = {old: new + 1 for new, old in enumerate(order)}
patients['cluster'] = [renumber[c] for c in kmeans5.labels_]
profile = patients.groupby('cluster')[features].mean().round(1)
profile['patients'] = patients['cluster'].value_counts().sort_index()
profile['% family history'] = (100 * patients.groupby('cluster')['family_history'].mean()).round(0)
print(profile.to_string())
age visits_last_year glucose systolic_bp cholesterol chronic_conditions patients % family history cluster 1 30.3 1.2 84.2 112.7 169.3 0.1 625 12.0 2 45.4 8.2 90.3 118.1 188.3 0.3 428 21.0 3 53.3 1.8 106.4 135.0 224.3 0.1 355 67.0 4 67.3 20.4 152.0 145.4 216.0 5.0 169 39.0 5 72.4 7.3 131.4 139.0 203.7 2.5 423 40.0
The five groups match the ones the team finds in the book:
- Cluster 1 (625 patients): young (30 years on average), about one visit a year, lab results in optimal ranges. "Healthy young people", as Owen labels them.
- Cluster 2 (428): patients of various ages with frequent visits (8 a year), but normal lab results and almost no chronic conditions: the diverse active cases.
- Cluster 3 (355): about 53 years old, fewer than two visits a year, glucose of 106 mg/dL and blood pressure of 135 mmHg, both at the top of the normal range, and 67% with a family history. This is the silent-risk group that surprises the team.
- Cluster 4 (169): the smallest group, with 20 visits a year and 5 chronic conditions: high-complexity patients.
- Cluster 5 (423): the oldest (72 years), with 2.5 chronic conditions and regular visits: stable chronic patients.
None of these names came from the algorithm, which only produced five centroids: naming the groups is the clinical team's job. Notice also that the family history, which K-Means never saw, is highest in cluster 3, as in the book.
We plot the patients in two of the dimensions, coloured by cluster. The groups overlap in any single pair of features: K-Means separates them using the six features at once.
fig, axes = plt.subplots(1, 2, figsize=(13, 5))
for ax, (fx, fy) in zip(axes, [('age', 'glucose'), ('glucose', 'systolic_bp')]):
for c in range(1, 6):
sub = patients[patients['cluster'] == c]
ax.scatter(sub[fx], sub[fy], s=8, alpha=0.5, label=f'cluster {c}')
ax.set_xlabel(fx)
ax.set_ylabel(fy)
ax.grid(True)
axes[0].legend(markerscale=3)
plt.suptitle('The five clusters seen in two dimensions')
plt.show()
print("Hidden profile (rows) against cluster found by K-Means (columns):")
print(pd.crosstab(hidden_profile, patients['cluster'], rownames=['hidden profile'], colnames=['cluster']).to_string())
print(f"ARI = {adjusted_rand_score(hidden_profile, patients['cluster']):.3f}")
Hidden profile (rows) against cluster found by K-Means (columns): cluster 1 2 3 4 5 hidden profile 0 591 7 2 0 0 1 0 1 23 5 414 2 34 419 17 0 1 3 0 1 0 164 8 4 0 0 313 0 0 ARI = 0.881
Each hidden profile ends up mostly in a single cluster, and all 313 patients of the silent-risk profile are in cluster 3. The errors are at the edges between profiles that overlap, for example 34 active cases younger than usual that end up with the healthy young. With real patients this table would not exist: it is only possible because we generated the data ourselves.
Step 8: The variables we include decide the groups¶
Martha's warning, "it only compares numbers using a similarity rule we choose, based on the variables we include", can be tested. We repeat the clustering adding the family history, the yes/no column we had left out:
with_family = StandardScaler().fit_transform(patients[features + ['family_history']])
km_family = KMeans(n_clusters=5, n_init=10, random_state=0).fit(with_family)
print(f"ARI with the family history included: {adjusted_rand_score(hidden_profile, km_family.labels_):.3f}\n")
print("Share of patients with a family history in each new cluster:")
print(pd.Series(patients.groupby(km_family.labels_)['family_history'].mean().round(2).to_numpy(),
index=[f'cluster {c + 1}' for c in range(5)]).to_string())
ARI with the family history included: 0.661 Share of patients with a family history in each new cluster: cluster 1 0.12 cluster 2 1.00 cluster 3 0.40 cluster 4 0.34 cluster 5 0.00
The agreement falls from 0.881 to 0.661, and the table explains why: K-Means has spent two of its five clusters separating patients with a family history (cluster 2, all of them) from patients without one (cluster 5, none of them). Once standardized, a yes/no column forms two distant clouds of points, and the squared distance rewards separating them. The algorithm has not made a mistake: it has done exactly what the distance we gave it asks for. Choosing the variables, and therefore what "similar" means, is a clinical decision, not only a technical one.
Step 9: Assigning a new patient¶
Once the groups exist, a new patient is assigned to the cluster with the nearest centroid, after applying the same scaling. Liam's example in the book is a patient with a glucose of 105 mg/dL and a blood pressure of 135 mmHg who feels well: in a routine visit, nothing would raise an alarm.
new_patients = pd.DataFrame([
{'age': 54, 'visits_last_year': 1, 'glucose': 105, 'systolic_bp': 135, 'cholesterol': 220, 'chronic_conditions': 0},
{'age': 29, 'visits_last_year': 1, 'glucose': 88, 'systolic_bp': 115, 'cholesterol': 175, 'chronic_conditions': 0},
{'age': 75, 'visits_last_year': 9, 'glucose': 135, 'systolic_bp': 142, 'cholesterol': 200, 'chronic_conditions': 3},
])
new_patients['cluster'] = [renumber[c] for c in kmeans5.predict(scaler.transform(new_patients[features].to_numpy(dtype=float)))]
print(new_patients.to_string(index=False))
age visits_last_year glucose systolic_bp cholesterol chronic_conditions cluster 54 1 105 135 220 0 3 29 1 88 115 175 0 1 75 9 135 142 200 3 5
The first patient, with the values of Liam's example, falls in cluster 3, the silent-risk group, and could be invited to the prevention program that Minermont launches in the book. The other two are assigned to the healthy young (cluster 1) and to the stable chronic patients (cluster 5).
Conclusion¶
- Unsupervised learning looks for structure in data without labels. K-Means groups the points around $K$ centroids by alternating two steps, assignment to the nearest centroid and recalculation of each centroid as a mean, and each iteration lowers or keeps the SSE.
- It always converges, but possibly to a local optimum: several starts and K-Means++ make that much less likely.
- Distances depend on units, so the features must be scaled, and the variables we include decide which groups can appear.
- The elbow method helps to choose $K$, but it is a heuristic: the final decision combines the curve with the clinical meaning of the groups.
- The algorithm only returns numbers and centroids. Naming the groups, deciding whether they make clinical sense and acting on them (like Minermont's prevention program for the silent-risk group) is the work of the medical team.