← 4. Feature extractionmodule indexsolutionsall modules6. Template matching →

Notebook 5 — Clustering into units

SWC ENC 2026 · ephys-pop module

In the feature space, each neuron is a cloud of points. Clustering turns those clouds into labelled units, with no ground-truth labels to guide it — this is the unsupervised heart of spike sorting. Then we average each cluster's snippets into a template: the unit's canonical spike, the thing matching pursuit will hunt for in Notebook 6.

In this notebook you will:

1. Decide how many units there are, using the BIC.
2. Cluster the spikes with a Gaussian mixture.
3. Average each cluster into a template and check it against the truth.

import numpy as np
import matplotlib.pyplot as plt
import picosort as ps

rec = ps.make_recording(n_units=6, duration_s=20.0, seed=0)
whitened, filtered, W = ps.preprocess(rec)
times, peak_channels = ps.detect_spikes(whitened, rec.probe, rec.fs, threshold=5.0)
snippets, times = ps.extract_snippets(filtered, times, peak_channels=peak_channels)
features, parts = ps.spike_features(snippets, rec.probe)

# Drop tiny threshold crossings: below ~100 µV they're noise and collision fragments,
# not real spikes. This keeps the clustering clean.
keep = parts["amplitude"] > 100
features, snippets = features[keep], snippets[keep]
times, depth, amplitude = times[keep], parts["depth"][keep], parts["amplitude"][keep]
print(f"{len(times)} spikes kept for clustering")

1. Clustering: fitting blobs to points

Clustering means: given a pile of points with no labels, discover the groups. Our tool is the Gaussian mixture. It models each group as a Gaussian — a blob of points scattered around a centre, densest in the middle and thinning out with distance (a 2-D bell curve). A mixture is simply several such blobs. Fitting one to the data finds each blob's centre and shape, and for every point the blob it most likely came from. Here it is on a toy with three blobs:

from sklearn.mixture import GaussianMixture

rng = np.random.default_rng(0)
centres = np.array([[0, 0], [5, 4], [8, -1]])
toy = np.vstack([rng.normal(c, 0.8, size=(120, 2)) for c in centres])
assignment = GaussianMixture(n_components=3, random_state=0).fit_predict(toy)

plt.figure(figsize=(5, 4))
for g in range(3):
    m = assignment == g
    plt.scatter(toy[m, 0], toy[m, 1], s=12, color=ps.plotting.unit_color(g))
plt.title("a Gaussian mixture recovers the three blobs")
plt.xlabel("feature 1"); plt.ylabel("feature 2"); plt.show()

The mixture coloured the three groups correctly, with no labels to learn from. Our spikes are just such a pile of points — in the depth–amplitude feature space from Notebook 4 — and each neuron is a blob. Two questions remain: how many blobs are there, and where does each spike belong?

2. How many units? The BIC

Clustering needs to know how many clusters to look for, and we don't — that's part of what we're trying to discover. First, see what goes wrong when you guess:

fig, ax = plt.subplots(1, 3, figsize=(13, 4), sharex=True, sharey=True)
for a, k in zip(ax, [3, 6, 9]):
    lab_k = GaussianMixture(k, n_init=5, random_state=0).fit_predict(features)
    for g in range(k):
        a.scatter(depth[lab_k == g], amplitude[lab_k == g], s=10, color=ps.plotting.unit_color(g))
    a.set_title(f"k = {k}  ({'too few' if k == 3 else 'just right' if k == 6 else 'too many'})")
    a.set_xlabel("depth (µm)")
ax[0].set_ylabel("amplitude (µV)")
plt.tight_layout(); plt.show()

With k = 3 (left) the mixture is forced to lump distinct units into one blob — a poor fit, with points stranded far from their cluster's centre. With k = 9 (right) it splits a single real unit across two colours — which explains the data no better than one cluster did, since those points were already well accounted for. k = 6 (middle) is the sweet spot: one cluster per real group, none wasted.

The BIC (Bayesian Information Criterion) turns this trade-off into one number to minimise: $\mathrm{BIC} = \text{misfit} + \text{penalty}$, where misfit measures how badly the model explains the data (smaller is better) and penalty charges a fixed cost for every cluster's parameters. Why does it have an elbow? Each cluster you add up to the true number lets the model wrap a Gaussian around a real group it was missing, so the misfit drops steeply — easily worth the small penalty. Each cluster beyond the true number only carves up a group that was already well explained, so the misfit barely moves while you keep paying the penalty. A steep drop, then a flat plateau — and the bend between them, the elbow, is the honest unit count.

▸ Go deeper: Gaussian mixtures, and where the BIC comes from (optional)

The model. A Gaussian mixture says each spike's feature vector $x$ was drawn from one of $k$ Gaussians — one per unit — with unknown mixing weights $\pi_j$, means $\mu_j$, and covariances $\Sigma_j$:

$$p(x \mid \theta) = \sum_{j=1}^{k} \pi_j\, \mathcal{N}(x;\, \mu_j, \Sigma_j).$$

We fit $\theta$ by maximising the log-likelihood $\log \hat L = \sum_i \log p(x_i\mid\theta)$, using the EM algorithm: alternate between the E-step (compute each spike's responsibility — the posterior probability it belongs to unit $j$) and the M-step (re-estimate each $\pi_j, \mu_j, \Sigma_j$ as responsibility-weighted averages). This is a soft version of the assign-then-average loop you already know. (k-means is the special case with equal spherical $\Sigma_j = \sigma^2 I$ and hard responsibilities — which is exactly why it merges close clusters and splits elongated ones, and why the full mixture does better here.)

The penalty. More clusters always fit better, so raw likelihood can't choose $k$. The BIC adds a complexity penalty:

$$\mathrm{BIC}(k) = -2\log \hat L \;+\; m_k \log n,$$

with $m_k$ the number of free parameters and $n$ the number of spikes. It is a large-$n$ approximation (via Laplace's method) to $-2\log p(\text{data}\mid k)$, the model evidence — the probability the data would arise under a $k$-unit model, with the parameters averaged out rather than fitted. Extra clusters beyond the truth buy almost no likelihood but keep paying $\log n$ per parameter, so the curve flattens: the elbow is the honest unit count.

Exercise 1 (~5 min). Compute the BIC curve over k = 2..10 and read off the elbow.

Check / unstuck. The curve should drop hard through k = 6, then level off. Stuck? Use ps.bic_curve(features).

from sklearn.mixture import GaussianMixture

def bic_curve(features, k_values=range(2, 11)):
    # YOUR CODE HERE: for each k, fit GaussianMixture(k, n_init=5, random_state=0)
    # to features and record .bic(features). Return {k: bic}.
    raise NotImplementedError

bics = bic_curve(features)
ks = list(bics); vals = list(bics.values())
plt.figure(figsize=(5, 3.4))
plt.plot(ks, vals, "o-"); plt.axvline(6, color="tab:red", ls="--", label="elbow -> 6 units")
plt.xlabel("number of clusters k"); plt.ylabel("BIC (lower = better)")
plt.legend(); plt.title("BIC has an elbow at the true unit count"); plt.show()

3. Cluster with a Gaussian mixture

With $k = 6$, fit a Gaussian mixture: it models each unit as an ellipse in feature space and assigns every spike to the most likely one. (A Gaussian mixture handles clusters of different size and shape, where simpler k-means would merge two close units and split a third.)

Exercise 2 (~4 min · easy). Fit the mixture and assign labels.

Check / unstuck. You should get 6 clusters that match the clouds by eye. Stuck? Use ps.cluster_spikes(features, 6).

labels = GaussianMixture(6, n_init=10, random_state=0).fit_predict(features)
ps.plotting.plot_feature_space(depth, amplitude, labels, title="clusters in depth–amplitude")
plt.show()

The six clusters match the six blobs by eye. But we clustered on just two features (depth, amplitude) — how do we know these are really six distinct neurons, and not (say) one neuron artificially split? A good check is to look at the clusters in the full waveform space: every spike's complete multichannel snippet (~2000 numbers), where two genuinely different neurons should differ in many ways at once.

We can't plot 2000 dimensions, so we squash them to 2-D — and which squashing matters. The obvious tool is PCA, but PCA is linear: it keeps the two directions of largest variance, and distinct clusters can still land on top of each other in that projection (you saw exactly this in Notebook 4, where PC1 just re-encoded amplitude). t-SNE is a nonlinear embedding built for this job: it tries to keep every point near its true neighbours and push unrelated points apart, so genuinely separate groups fall into cleanly separated islands. Compare the two on our snippets:

from sklearn.decomposition import PCA
X = snippets.reshape(len(snippets), -1)          # every spike as its full multichannel snippet
pca2 = PCA(n_components=2).fit_transform(X)       # linear 2-D projection
emb = ps.tsne_embedding(X)                         # nonlinear 2-D embedding (sklearn's TSNE)

fig, axes = plt.subplots(1, 2, figsize=(11, 4.5))
for u in np.unique(labels):
    m = labels == u
    axes[0].scatter(pca2[m, 0], pca2[m, 1], s=10, color=ps.plotting.unit_color(int(u)))
    axes[1].scatter(emb[m, 0], emb[m, 1], s=10, color=ps.plotting.unit_color(int(u)))
axes[0].set_title("PCA (linear) — some clusters overlap")
axes[1].set_title("t-SNE (nonlinear) — clean islands")
for a in axes:
    a.set_xticks([]); a.set_yticks([])
plt.tight_layout(); plt.show()

In the PCA projection several clusters pile up in the middle — the linear map can't pull them apart. In the t-SNE map all six sit in clean, separate islands: strong confirmation that our six clusters really are six distinct units. That's why t-SNE is a staple of spike-sorting QC.

One caution. t-SNE is for looking, not measuring — the distances between islands and their sizes are not meaningful (it deliberately warps space to separate neighbourhoods), and on structureless data it can even invent apparent clusters. Use it to eyeball separation, never as proof on its own.

4. From clusters to templates

A template is the average of a cluster's snippets — the same noise-averages-away trick from Notebook 1, but now with clusters we discovered instead of labels we were handed. Average each cluster's multichannel snippets to get one template per unit.

Exercise 3 (~5 min). Complete templates_from_labels: for each label, average its snippets into a (n_channels, n_samples) template.

Check / unstuck. You should get 6 templates, each a clean footprint at a distinct depth. Stuck? Use ps.templates_from_labels(snippets, labels).

def templates_from_labels(snippets, labels):
    ids = np.array(sorted(u for u in np.unique(labels) if u >= 0))
    # YOUR CODE HERE: for each id, average snippets[labels==id] over spikes (axis 0).
    # Stack into (n_units, n_channels, n_samples). Return templates, ids.
    raise NotImplementedError

templates, ids = templates_from_labels(snippets, labels)
ps.plotting.plot_templates(templates, rec.probe)
plt.show()

5. Did we recover the real units?

Because this is synthetic data, we can check. Match each learned template to the true template it most resembles and compare their shapes — they should be nearly identical, up to noise.

gt = rec.ground_truth
corr = np.zeros((len(templates), gt.n_units))
for a in range(len(templates)):
    for b in range(gt.n_units):
        corr[a, b] = np.corrcoef(templates[a].ravel(), gt.templates[b].ravel())[0, 1]
best = corr.max(axis=1)
print("each learned template's correlation to its best-matching true template:")
print(np.round(best, 3))
print(f"worst match: {best.min():.3f}  ->  all six units recovered" if best.min() > 0.8
      else f"worst match: {best.min():.3f}")

Wrap-up

You found the number of units from the BIC elbow, clustered the spikes with a Gaussian mixture, and averaged each cluster into a template that matches a real neuron. We now have a template per unit — but clustering assigned each spike to exactly one unit, which breaks down when two spikes overlap.

Next (Notebook 6 — matching pursuit): use the templates to pull apart overlapping spikes, the collisions clustering can't handle.

← 4. Feature extractionmodule indexsolutionsall modules6. Template matching →