PCA for Feature Extraction in Spike Sorting

Import Libraries

Import the modules required for this notebook

import numpy as np
from matplotlib import pyplot as plt
from sklearn.decomposition import PCA
from sklearn.mixture import GaussianMixture

Download Data

Download the data required for this session

import requests
from pathlib import Path

url = "https://uni-bonn.sciebo.de/s/aFQ1gcUbOHRDEtP"
fname = Path("spike_waveforms.npy")
tmp_name = fname.with_suffix(".download")

print("Downloading Data ...")
response = requests.get(f"{url}/download", timeout=60)
response.raise_for_status()
tmp_name.write_bytes(response.content)

# Refuse to save server error pages or other non-NPY content as the dataset.
with open(tmp_name, "rb") as file:
    magic = file.read(6)

if magic != b"\x93NUMPY":
    tmp_name.unlink(missing_ok=True)
    raise ValueError("Downloaded file is not a valid .npy file. The server may have returned an error page instead.")

tmp_name.replace(fname)
print("Done!")

Downloading Data ...
Done!

This section will introduce you to PCA using the implementation from the Scikit-learn package. When we apply PCA, we have to decide the number of components we want to keep. Each component explains a certain amount of variance in the data and usually a few components are enough to explain most of the variance. We can plot the variance explained by each component and their sum to get a quick idea about how many components are required to represent the data accurately.

Code Reference

Exercises

Run the cell below to load the data. The data contains 19482 spike waveforms recorded at 4 channels. Each waveform is 30 samples long.

waveforms = np.load("spike_waveforms.npy")
waveforms.shape

(19482, 30, 4)

Run the code below to plot the first 10 waveforms for every channel. As you can see, the latency and amplitude of the spikes differs between channels. This information can be exploited by spike-sorting algorithms.

X = waveforms[:10, :, :]
fig, ax = plt.subplots(2, 2, sharex=True, sharey=True)

ax[0, 0].plot(X[:, :, 0].T);
ax[0, 1].plot(X[:, :, 1].T);
ax[1, 0].plot(X[:, :, 2].T);
ax[1, 1].plot(X[:, :, 3].T);

ax[0, 0].set(ylabel="Amplitude [$\\mu$V]", title="Channel 1")
ax[0, 1].set(title="Channel 2")
ax[1, 0].set(ylabel="Amplitude [$\\mu$V]", xlabel="Sample", title="Channel 3")
ax[1, 1].set(xlabel="Sample", title="Channel 4");

Example: Use PCA with n_components=5 to transform the spike waveforms recorded at the first channel.

X = waveforms[:, :, 0]
pca = PCA(n_components=5)
X_transformed = pca.fit_transform(X)
X.shape

(19482, 30)

Exercise: Use PCA with n_components=7 to transform the spike waveforms recorded at the first channel. Compare the .shape of the original and the transformed data.

X = waveforms[:, :, 0]
pca = PCA(n_components=7)
X_transformed = pca.fit_transform(X)
print(X.shape, X_transformed.shape)

(19482, 30) (19482, 7)

Exercise: Use PCA with n_components=15 to transform the spike waveforms recorded at the first channel. Compare the .shape of the original and the transformed data.

X = waveforms[:, :, 0]
pca = PCA(n_components=15)
X_transformed = pca.fit_transform(X)
print(X.shape, X_transformed.shape)

(19482, 30) (19482, 15)

Exercise: What error message do you observe when you try to apply PCA with n_components=41?

pca = PCA(n_components=41)
X_transformed = pca.fit_transform(X)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[62], line 2
      1 pca = PCA(n_components=41)
----> 2 X_transformed = pca.fit_transform(X)

File ~/Projects/ANDANI/anda-course/03_dimensionality_reduction/.pixi/envs/default/lib/python3.13/site-packages/sklearn/utils/_set_output.py:319, in _wrap_method_output.<locals>.wrapped(self, X, *args, **kwargs)
    317 @wraps(f)
    318 def wrapped(self, X, *args, **kwargs):
--> 319     data_to_wrap = f(self, X, *args, **kwargs)
    320     if isinstance(data_to_wrap, tuple):
    321         # only wrap the first output for cross decomposition
    322         return_tuple = (
    323             _wrap_data_with_container(method, data_to_wrap[0], X, self),
    324             *data_to_wrap[1:],
    325         )

File ~/Projects/ANDANI/anda-course/03_dimensionality_reduction/.pixi/envs/default/lib/python3.13/site-packages/sklearn/base.py:1403, in _fit_context.<locals>.decorator.<locals>.wrapper(estimator, *args, **kwargs)
   1393     estimator._validate_params()
   1395 with (
   1396     config_context(
   1397         skip_parameter_validation=(
   (...)   1401     callback_management_context(estimator),
   1402 ):
-> 1403     return fit_method(estimator, *args, **kwargs)

File ~/Projects/ANDANI/anda-course/03_dimensionality_reduction/.pixi/envs/default/lib/python3.13/site-packages/sklearn/decomposition/_pca.py:466, in PCA.fit_transform(self, X, y)
    443 @_fit_context(prefer_skip_nested_validation=True)
    444 def fit_transform(self, X, y=None):
    445     """Fit the model with X and apply the dimensionality reduction on X.
    446 
    447     Parameters
   (...)    464     C-ordered array, use 'np.ascontiguousarray'.
    465     """
--> 466     U, S, _, X, x_is_centered, xp = self._fit(X)
    467     if U is not None:
    468         U = U[:, : self.n_components_]

File ~/Projects/ANDANI/anda-course/03_dimensionality_reduction/.pixi/envs/default/lib/python3.13/site-packages/sklearn/decomposition/_pca.py:540, in PCA._fit(self, X)
    538 # Call different fits for either full or truncated SVD
    539 if self._fit_svd_solver in ("full", "covariance_eigh"):
--> 540     return self._fit_full(X, n_components, xp, is_array_api_compliant)
    541 elif self._fit_svd_solver in ["arpack", "randomized"]:
    542     return self._fit_truncated(X, n_components, xp)

File ~/Projects/ANDANI/anda-course/03_dimensionality_reduction/.pixi/envs/default/lib/python3.13/site-packages/sklearn/decomposition/_pca.py:554, in PCA._fit_full(self, X, n_components, xp, is_array_api_compliant)
    550         raise ValueError(
    551             "n_components='mle' is only supported if n_samples >= n_features"
    552         )
    553 elif not 0 <= n_components <= min(n_samples, n_features):
--> 554     raise ValueError(
    555         f"n_components={n_components} must be between 0 and "
    556         f"min(n_samples, n_features)={min(n_samples, n_features)} with "
    557         f"svd_solver={self._fit_svd_solver!r}"
    558     )
    560 self.mean_ = xp.mean(X, axis=0)
    561 # When X is a scipy sparse matrix, self.mean_ is a numpy matrix, so we need
    562 # to transform it to a 1D array. Note that this is not the case when X
    563 # is a scipy sparse array.
    564 # TODO: remove the following two lines when scikit-learn only depends
    565 # on scipy versions that no longer support scipy.sparse matrices.

ValueError: n_components=41 must be between 0 and min(n_samples, n_features)=30 with svd_solver='covariance_eigh'

Exercise: Plot the ratio of variance explained by the PCA components (i.e. pca.explained_variance_ratio_)

pca = PCA(n_components=15)
X_transformed = pca.fit_transform(X)

Solution

plt.plot(pca.explained_variance_ratio_);

Exercise: Compute the cumulative ratio of explained variance using np.cumsum() and plot it

Solution

plt.plot(np.cumsum(pca.explained_variance_ratio_));

Exercise: What ratio of the variance is explained by the first three components (i.e. what is the sum() of the first three elements of .explained_variance_ratio_?) and how many components are required to account for over 99% of variance?

print(sum(pca.explained_variance_ratio_[:3]))
print(sum(pca.explained_variance_ratio_[:8]))

0.820759476628594
0.9926894524591845

PCA is a powerful technique, but its components can be abstract. To build intuition, it’s helpful to visualize what these components represent in relation to the original data. In this section, we will plot the principal components themselves to see what waveform shapes they capture. These components are directions of variation in centered data space and are shown in arbitrary units, not as actual recorded waveforms in microvolts. We will also use the inverse PCA transformation to reconstruct waveforms from a limited number of components, which demonstrates how much information is retained.

Code Reference

Exercises

Example: Fit a PCA with n_components=3 to the waveforms recorded at the first channel and plot the first component, along with the first 100 recorded waveforms.

X = waveforms[:, :, 0]
pca = PCA(n_components=3)
X_transformed = pca.fit_transform(X)

plt.subplot(2, 1, 1)
plt.plot(X[:100].T, linewidth=0.5)
plt.ylabel("Amplitude [$\\mu$V]")

plt.subplot(2, 1, 2)
plt.plot(pca.components_[0])
plt.xlabel("Sample")
plt.ylabel("PC weight [a.u.]")
plt.tight_layout()

Exercise: Plot the second and third principal components along with the first 100 waveforms recorded at the first channel.

Solution

plt.subplot(3, 1, 1)
plt.plot(X[:100].T, linewidth=0.5)
plt.ylabel("Amplitude [$\\mu$V]")

plt.subplot(3, 1, 2)
plt.plot(pca.components_[1])
plt.ylabel("PC weight [a.u.]")

plt.subplot(3, 1, 3)
plt.plot(pca.components_[2])
plt.xlabel("Sample")
plt.ylabel("PC weight [a.u.]")
plt.tight_layout()

Example: Compute a PCA with n_components=1 for the waveforms recorded at the first channel and transform the data. Then, compute the .inverse_transform() of the transformed data and plot the result of the inverse transformation X_inverse along with the original data X for the first spike.

X = waveforms[:, :, 0]
pca = PCA(n_components=1)
X_transformed = pca.fit_transform(X)
X_inverse = pca.inverse_transform(X_transformed)

plt.plot(X[0], label="Original")
plt.plot(X_inverse[0], label="Reconstructed")
plt.xlabel("Sample")
plt.ylabel("Amplitude [$\\mu$V]")
plt.legend();

Exercise: Plot the original data X and the result of the inverse transformation X_inverse for the 7th spike.

Solution

plt.plot(X[6], label="Original")
plt.plot(X_inverse[6], label="Reconstructed")
plt.xlabel("Sample")
plt.ylabel("Amplitude [$\\mu$V]")
plt.legend();

Exercise: Re-compute the PCA and increase the number of n_components, then fit_transform and inverse_transform the data. How many components do you need before the original data X and the reconstruction X_inverse look nearly indistinguishable? How does this compare to the cumulative explained variance from the previous section?

Solution

X = waveforms[:, :, 0]
pca = PCA(n_components=9)
X_transformed = pca.fit_transform(X)
X_inverse = pca.inverse_transform(X_transformed)
plt.plot(X[6], label="Original")
plt.plot(X_inverse[6], label="Reconstructed")
plt.xlabel("Sample")
plt.ylabel("Amplitude [$\\mu$V]")
plt.legend();

In the previous sections, we used PCA to describe each spike waveform with only a few numbers instead of the full waveform across time. We can now use these PCA scores as features for visualization. If spikes from different neurons have distinct waveform shapes across the tetrode channels, they may form separate groups in this low-dimensional PCA feature space.

Here, we compute PCA separately for each channel and keep the first three components per channel. We then compare the first PCA score across different channels to see whether the spikes form visually separable clusters. Each point in the plots represents one detected spike.

Code Reference

Exercises

Run the code below to compute PCA with 3 components on the waveforms recorded on every channel and store the transformed data in a new array waveforms_transformed.

n_components = 3
pca = PCA(n_components=n_components)
n_spikes = waveforms.shape[0]
n_channels = waveforms.shape[-1]
waveforms_transformed = np.zeros((n_spikes, n_components, n_channels))
for i in range(n_channels):
    X = waveforms[:, :, i]
    X_transformed = pca.fit_transform(X)
    waveforms_transformed[:, :, i] = X_transformed

Example: Plot the 1st PCA score at channel 1 against the 1st PCA score at channel 2 for all spikes (each spike is one dot).

plt.scatter(waveforms_transformed[:, 0, 0], waveforms_transformed[:, 0, 1], s=1)
plt.xlabel("Ch 1 PC1 score")
plt.ylabel("Ch 2 PC1 score");

Exercise: Plot the 1st PCA score at channel 1 against the 1st PCA score at channel 3 for all spikes. Can you visually identify clusters that the individual spikes may fall into?

Solution

plt.scatter(waveforms_transformed[:, 0, 0], waveforms_transformed[:, 0, 2], s=1)
plt.xlabel("Ch 1 PC1 score")
plt.ylabel("Ch 3 PC1 score");

Exercise: Plot the 1st PCA score at channel 3 against the 1st PCA score at channel 4 for all spikes. Can you visually identify clusters that the individual spikes may fall into?

Solution

plt.scatter(waveforms_transformed[:, 0, 2], waveforms_transformed[:, 0, 3], s=1)
plt.xlabel("Ch 3 PC1 score")
plt.ylabel("Ch 4 PC1 score")

Exercise: Plot the 1st PCA score at channel 2 against the 1st PCA score at channel 3 for all spikes. Can you visually identify clusters that the individual spikes may fall into?

Solution

plt.scatter(waveforms_transformed[:, 0, 1], waveforms_transformed[:, 0, 2], s=1)
plt.xlabel("Ch 2 PC1 score")
plt.ylabel("Ch 3 PC1 score");

Example: Create a 3D scatterplot with the 1st PCA score at channels 1, 2, and 3.

fig = plt.figure(figsize=(7, 6))
ax = fig.add_subplot(projection="3d")

scatter = ax.scatter(
    waveforms_transformed[:, 0, 0],
    waveforms_transformed[:, 0, 1],
    waveforms_transformed[:, 0, 2],
    s=1,
)

ax.set_xlabel("Ch 1 PC1 score")
ax.set_ylabel("Ch 2 PC1 score")
ax.set_zlabel("Ch 3 PC1 score")
ax.view_init(elev=20, azim=80)

Exercise: Create a 3D scatterplot with the 1st PCA score at channels 1, 3, and 4. Adjust the viewing angle with ax.view_init(elev, azim) so that you can see the clusters.

Solution

fig = plt.figure(figsize=(7, 6))
ax = fig.add_subplot(projection="3d")

scatter = ax.scatter(
    waveforms_transformed[:, 0, 0],
    waveforms_transformed[:, 0, 2],
    waveforms_transformed[:, 0, 3],
    s=1,
)

ax.set_xlabel("Ch 1 PC1 score")
ax.set_ylabel("Ch 3 PC1 score")
ax.set_zlabel("Ch 4 PC1 score")
ax.view_init(elev=60, azim=20)

In the previous section, we could already see clusters by plotting a few PCA scores against each other. However, those plots only show two or three dimensions at a time, while our full feature representation contains several PCA scores across multiple channels. To cluster spikes using all of these features at once, we need an algorithm that can work in higher-dimensional space.

A Gaussian Mixture Model (GMM) is one way to do this. It models the data as a mixture of several Gaussian-shaped clusters and assigns each spike to the cluster that best explains its PCA features. Applying the GMM itself is a demonstration rather than the main teaching objective of the notebook: the goal is to see how the PCA features can be passed into a clustering algorithm and then visualized back in PCA space.

Code Reference

Applying a Gaussian Mixture Model

Run the cell below to reshape waveforms_transformed into a 2D array and then apply a GaussianMixture model with n_components=13 to fit 13 Gaussian clusters to the data.

Note: How to find the optimal number of clusters is outside the scope of this notebook, but for this data, 13 clusters is optimal according to the Bayesian Information Criterion (BIC).

X = waveforms_transformed.reshape(-1, n_components * n_channels, order="F")
gmm = GaussianMixture(
    n_components=13,
    random_state=7
)
gmm.fit(X)

Now we can use the fitted model to assign every spike to a cluster.

spike_labels = gmm.predict(X)
spike_labels

array([ 2, 10,  9, ...,  2,  2, 11], shape=(19482,))

First, inspect how many spikes were assigned to each cluster.

plt.hist(spike_labels)
plt.xlabel("Cluster")
plt.ylabel("Count");

Exercise: Plot the 1st PCA score at channel 1 against the 1st PCA score at channel 2 for all spikes and highlight the spikes that belong to clusters 0 through 5.

Solution

clusters = [0, 1, 2, 3, 4, 5]
X = waveforms_transformed
plt.scatter(X[:, 0, 0], X[:, 0, 1], s=1, color="gray")
for cluster in clusters:
    X = waveforms_transformed[spike_labels == cluster]
    plt.scatter(X[:, 0, 0], X[:, 0, 1], s=5, label=cluster)
    plt.xlabel("Ch 1 PC1 score")
    plt.ylabel("Ch 2 PC1 score")
plt.legend();

Exercise: Plot the 1st PCA score at channel 2 against the 1st PCA score at channel 3 for all spikes and highlight the spikes that belong to clusters 0 through 5.

Solution

clusters = [0, 1, 2, 3, 4, 5]
X = waveforms_transformed
plt.scatter(X[:, 0, 1], X[:, 0, 2], s=1, color="gray")
for cluster in clusters:
    X = waveforms_transformed[spike_labels == cluster]
    plt.scatter(X[:, 0, 1], X[:, 0, 2], s=1, label=cluster)
    plt.xlabel("Ch 2 PC1 score")
    plt.ylabel("Ch 3 PC1 score")
plt.legend();

Bonus: Create a 3D scatterplot with the 1st PCA score at channels 1, 2, and 3 and highlight the points that belong to clusters 0 through 5. Adjust the viewing angle with ax.view_init(elev, azim) so that you can see the clusters.

Solution

clusters = [0, 1, 2, 3, 4, 5]
X = waveforms_transformed
fig = plt.figure(figsize=(7, 6))
ax = fig.add_subplot(projection="3d")

scatter = ax.scatter(X[:, 0, 0], X[:, 0, 1], X[:, 0, 2], s=1, color="gray")

for cluster in clusters:
    X = waveforms_transformed[spike_labels == cluster]
    scatter = ax.scatter(X[:, 0, 0], X[:, 0, 1], X[:, 0, 2], s=3, label=cluster)

ax.set_xlabel("Ch 1 PC1 score")
ax.set_ylabel("Ch 2 PC1 score")
ax.set_zlabel("Ch 3 PC1 score")
ax.legend()
ax.view_init(elev=20, azim=-20)