Feature Extraction and Waveform Clustering

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

url = "https://uni-bonn.sciebo.de/s/aFQ1gcUbOHRDEtP"
fname = "spike_waveforms.npy"
response = requests.get(f"{url}/download")
print("Downloading Data ...")
with open(fname, "wb") as file:
    file.write(response.content)
print("Done!")

Downloading Data ...
Done!

In principal, we could run spike sorting on the raw spike waveforms. However, because there are too many features (i.e., time points) in the data and they are highly correlated, the sorting process would be very inefficient. Thus, we’ll have to extract features from the waveforms and use those to sort the spikes. Principal Component Analysis (PCA) is a method that allows us to extract these features in an automated and unsupervised manner. PCA will give us a number of uncorrelated components that provide a concise description of our data. This way, we can drastically reduce the number of dimensions while keeping the most relevant information. In this section, we will learn how to apply PCA to the raw spike waveforms.

Exercise: Load the array of of spike waveforms stored in "spike_waveforms.npy", assign it to a variable waveforms print its .shape. How many spikes are stored in the array?

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

(19482, 30, 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)

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=40?

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

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

File ~/repositories/new-learning-platform/notebooks/spike_analysis/.pixi/envs/default/lib/python3.12/site-packages/sklearn/utils/_set_output.py:316, in _wrap_method_output.<locals>.wrapped(self, X, *args, **kwargs)
    314 @wraps(f)
    315 def wrapped(self, X, *args, **kwargs):
--> 316     data_to_wrap = f(self, X, *args, **kwargs)
    317     if isinstance(data_to_wrap, tuple):
    318         # only wrap the first output for cross decomposition
    319         return_tuple = (
    320             _wrap_data_with_container(method, data_to_wrap[0], X, self),
    321             *data_to_wrap[1:],
    322         )

File ~/repositories/new-learning-platform/notebooks/spike_analysis/.pixi/envs/default/lib/python3.12/site-packages/sklearn/base.py:1365, in _fit_context.<locals>.decorator.<locals>.wrapper(estimator, *args, **kwargs)
   1358     estimator._validate_params()
   1360 with config_context(
   1361     skip_parameter_validation=(
   1362         prefer_skip_nested_validation or global_skip_validation
   1363     )
   1364 ):
-> 1365     return fit_method(estimator, *args, **kwargs)

File ~/repositories/new-learning-platform/notebooks/spike_analysis/.pixi/envs/default/lib/python3.12/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 ~/repositories/new-learning-platform/notebooks/spike_analysis/.pixi/envs/default/lib/python3.12/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 ~/repositories/new-learning-platform/notebooks/spike_analysis/.pixi/envs/default/lib/python3.12/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

Exercise: Fit a PCA with n_components=15 to the waveforms recorded at the second channel, then compute and plot the PCA component’s cumulative ratio of explained variance.

Solution

X = waveforms[:, :, 1]
pca = PCA(n_components=15)
X_transformed = pca.fit_transform(X)
plt.plot(np.cumsum(pca.explained_variance_ratio_))

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. We will also use the inverse PCA transformation to reconstruct waveforms from a limited number of components, which demonstrates how much information is retained. Finally, we will plot the transformed data in the new low-dimensional “PCA space,” which should already begin to reveal the distinct clusters of spikes we aim to identify.mprehend. In this section, we are going to visualize the extracted components and compare them to the original waveform to understand how they relate to each other. We’ll also apply PCA to transform the original waveforms and plot the transformed waveforms in the new PCA-space. This visualization should already reveal some of the clusters that we want to identify later.

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.subplot(2, 1, 2)
plt.plot(pca.components_[0])

Exercise: Plot the second and third principle component 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.subplot(3, 1, 2)
plt.plot(pca.components_[1])
plt.subplot(3, 1, 3)
plt.plot(pca.components_[2])

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

Solution

X = waveforms[:, :, 2]
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.subplot(2, 1, 2)
plt.plot(pca.components_[0])

Example: Compute a PCA with n_components=1 for the waveforms recorded at the first channel and transform the data. Then, compute the .inverse_transfom() 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])
plt.plot(X_inverse[0])

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

Solution

plt.plot(X[6])
plt.plot(X_inverse[6])

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 so the original data X and X_inverse become indistinguishable?

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])
plt.plot(X_inverse[6])

Exercise: Complete the for loop below by computing X_transformed for the spike waveforms X at each channel i.

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 = # your code here 
    waveforms_transformed[:, :, i] = X_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 amplitude of the 1st PCA feature at the channel 1 against the amplitude at the 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 Amplitude [a.u.]")
plt.ylabel("Ch 2 Amplitude [a.u.]")

Exercise: Plot the amplitude of the 1st PCA feature at channel 1 against the amplitude 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 Amplitude [a.u.]")
plt.ylabel("Ch 3 Amplitude [a.u.]")

Exercise: Plot the amplitude of the 1st PCA feature at channel 3 against the amplitude 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 Amplitude [a.u.]")
plt.ylabel("Ch 4 Amplitude [a.u.]")

Now that we have extracted low-dimensional features from our waveforms using PCA, we can perform clustering to group spikes into putative single units. We will use a Gaussian Mixture Model (GMM), a probabilistic model that assumes the data is generated from a mixture of a finite number of Gaussian distributions with unknown parameters. A key challenge in clustering is choosing the right number of clusters. To address this, we will use the Bayesian Information Criterion (BIC), a model selection tool that balances model fit with model complexity. A lower BIC score indicates a better model, helping us avoid overfitting by penalizing models with too many clusters.

Exercise: Execute the cell below to .reshape() the 3-dimensional matrix waveforms_transformed into a 2-dimensional matrix X by concatenating the components and channels dimension. What is the new .shape of X?

X = waveforms_transformed.reshape(-1, n_components * n_channels, order="F")

X.shape

(19482, 12)

Example: Apply a GaussianMixture model with n_components=5 to the first 1000 spikes in waveforms and compute the models .score() (i.e. the log likelihood).

gmm = GaussianMixture(n_components=5)
gmm.fit(X[:1000])
gmm.score(X[:1000])

np.float64(-69.7095752453037)

Exercise: Apply a GaussianMixture model with n_components=10 to the first 1000 spikes in waveforms and compute the models .score(). Is this model more accurate than the model with n_components=5?

gmm = GaussianMixture(n_components=10)
gmm.fit(X[:1000])
gmm.score(X[:1000])

np.float64(-69.21182996115286)

Exercise: Apply a GaussianMixture model with n_components=100 to the first 1000 spikes in waveforms and compute the models .score(). Will the model’s accuracy always increase with the number of components? What is the largest number of components possible for this model?

gmm = GaussianMixture(n_components=100)
gmm.fit(X[:1000])
gmm.score(X[:1000])

np.float64(-44.97703577694651)

The largest number of components is equal to the number of data points, i.e. 1000

Exercise: Apply a GaussianMixture model with n_components=10 to the first 1000 spikes in waveforms and compute the Bayesian Information Criterion .bic().

gmm = GaussianMixture(n_components=10)
gmm.fit(X[:1000])
gmm.bic(X[:1000])

np.float64(144153.30510299263)

Exercise: Apply a GaussianMixture model with n_components=5 to the first 1000 spikes in waveforms and compute the Bayesian Information Criterion .bic(). Does this model perform better or worse than the one with n_components=10 according to the BIC (Hint: a lower BIC indicates a better model)?

gmm = GaussianMixture(n_components=5)
gmm.fit(X[:1000])
gmm.bic(X[:1000])

np.float64(142563.24204389623)

Exercise: Increase the number of spikes the GaussianMixture model is fit on to 8000. Which model performs better now, the one with 5 or 10 components?

gmm = GaussianMixture(n_components=5)
gmm.fit(X[:8000])
print(gmm.bic(X[:8000]))
gmm = GaussianMixture(n_components=10)
gmm.fit(X[:8000])
print(gmm.bic(X[:8000]))

1120764.9084488382
1117851.4819057845

Exercise: Fit a GaussianMixture model with n_components=13 (which is the optimal number of components according to BIC) to the whole data.

gmm = GaussianMixture(
    n_components=13,
)
gmm.fit(X)

With an optimal GMM fitted to our data, we can now assign each spike to a specific cluster. These cluster assignments are our predicted single units. To validate the quality of our clustering, we will inspect the results both quantitatively and visually. We will examine the probability of each spike’s cluster assignment to gauge the model’s confidence. Finally, we will visualize the clusters in the PCA feature space, coloring each spike by its assigned label to see if the clusters are well-separated, which is a key indicator of successful spike sorting.

Exercise: Use gmm.predict to predict the spike_labels for the data X and print them.

spike_labels = gmm.predict(X)
spike_labels

array([ 3,  2, 11, ...,  3,  3,  1], shape=(19482,))

Exercise: How many unique spike_labels are there?

print(len(np.unique(spike_labels)))

13

Exercise: Use gmm.predict_proba to predict the cluster_probs for X (i.e. the probabilities that each element in X belongs to a given cluster) and print their shape.

cluster_probs = gmm.predict_proba(X)
cluster_probs.shape

(19482, 13)

Exercise: What is the spike_label assigned to the first spike? How many spikes have been assigned to the same label?

print(spike_labels[0])
sum(spike_labels==spike_labels[0])

3

np.int64(3372)

Exercise: Plot the probability of cluster membership cluster_prob for the first spike as a histogram.

Solution

plt.hist(cluster_probs[0])

Exercise: What is the sum of all cluster_probs for a the first spike?

sum(cluster_probs[0])

np.float64(0.9999999999999953)

Exercise: Get the largest cluster membership probability for each spike (i.e. cluster_probs.max(axis=1) and plot them in a histogram.

Solution

plt.hist(cluster_probs.max(axis=1))

Example: Plot the amplitude of PCA feature at channel 1 (X[:, 0, 0]) against the amplitude at channel 2 (X[:, 0, 1]) for the spikes from clusters=[0, 1, 2, 3]

clusters = [0, 1, 2, 3]
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)
    plt.xlabel("Amplitude Ch 1 [a.u.]")
    plt.ylabel("Amplitude Ch 2 [a.u.]")

Exercise: Plot the amplitude of PCA feature at channel 1 (X[:, 0, 0]) against the amplitude at channel 3 (X[:, 0, 2]) for the spikes from clusters=[0, 1, 2, 3].

Solution

clusters = [0, 1, 2, 3]
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("Amplitude Ch 2 [a.u.]")
    plt.ylabel("Amplitude Ch 3 [a.u.]")
plt.legend()