Skip to main content

Agglomerative clustering with TensorFlow data

TensorFlow 2.21 has no native high-level agglomerative-clustering estimator. A custom loop over a dense pairwise matrix is only an illustration and is easy to confuse with a complete implementation. Here TensorFlow creates the feature tensor, while scikit-learn performs flat clustering and SciPy constructs the independently visualized hierarchy.

1. Prepare deterministic tensor data​

import numpy as np
import tensorflow as tf
from sklearn.datasets import make_blobs
from sklearn.preprocessing import StandardScaler

tf.random.set_seed(42)
X, _ = make_blobs(n_samples=150, centers=3, cluster_std=0.60, random_state=42)
X_scaled = StandardScaler().fit_transform(X).astype("float32")
points = tf.constant(X_scaled)
assert points.shape == (150, 2)

2. Inspect pairwise distances​

pairwise = tf.norm(points[:, tf.newaxis, :] - points[tf.newaxis, :, :], axis=2)
assert pairwise.shape == (150, 150)
np.testing.assert_allclose(tf.linalg.diag_part(pairwise).numpy(), 0.0, atol=1e-6)

This quadratic matrix is a teaching aid, not a scalable hierarchy builder. Linkage must define distances between sets of observations and update them after every merge.

3. Fit a maintained estimator​

Ward linkage greedily chooses the merge that causes the smallest increase in within-cluster variance and therefore requires Euclidean geometry.

from sklearn.cluster import AgglomerativeClustering

model = AgglomerativeClustering(n_clusters=3, metric="euclidean", linkage="ward")
labels = model.fit_predict(points.numpy())
assert np.unique(labels).size == 3

4. Build and plot the hierarchy​

SciPy's linkage builds a linkage matrix for the same scaled observations. The labels produced by cutting the tree are arbitrary identifiers, so compare partitions with a label-invariant metric.

import matplotlib.pyplot as plt
from scipy.cluster.hierarchy import dendrogram, fcluster, linkage
from sklearn.metrics import adjusted_rand_score, silhouette_score

linked = linkage(X_scaled, method="ward", metric="euclidean")
cut_labels = fcluster(linked, t=3, criterion="maxclust")
assert adjusted_rand_score(labels, cut_labels) == 1.0
silhouette = silhouette_score(X_scaled, labels)
assert 0.70 < silhouette <= 1.0

fig, ax = plt.subplots(figsize=(9, 4))
dendrogram(linked, truncate_mode="lastp", p=12, show_leaf_counts=True, ax=ax)
ax.set(title="Truncated Ward dendrogram", xlabel="merged node or observation", ylabel="Ward distance")
fig.tight_layout()
plt.close(fig)

Use TensorFlow upstream when a model creates the representation. Once frozen, move the observations to a maintained clustering implementation and record scaling, metric, linkage, and cut rule together.