Cebra Decoding with no training Example

The following example show how to use the Cebra decoding without training.

import os

# load example_cebra_decoding.html
with open(os.path.join("..", "examples", "example_cebra_decoding.html"), "rt") as fh:
    html_data = fh.read()

tmp_dir = os.path.join("..", "docs", "source", "auto_examples")
if os.path.exists(tmp_dir):
    # building the docs with sphinx-gallery
    with open(os.path.join(tmp_dir, "out.html"), "wt") as fh:
        fh.write(html_data)
# set example path for thumbnail
# sphinx_gallery_thumbnail_path = '_static/CEBRA_embedding.png'

CEBRA example

Decoding without patient-individual training¶

This example demonstrates movement decoding without patient-individual training, as shown in the publication "Invasive neurophysiology and whole brain connectomics for neural decoding in patients with brain implants" (Merk et al, 2024). This notebook requires installation of the cebra package.

In [1]:

import cebra
from cebra import CEBRA
import pandas as pd
import numpy as np
import os
from scipy.ndimage import gaussian_filter1d
from sklearn import linear_model
from matplotlib import pyplot as plt

In this notebook feature computation is not shown. We will read a npy saved array that contains py_neuromodulation computed FFT features with binary movement labels.

In addition, optimal functional connectivity correlation values for each channel were computed and saved in a table. The movement decoding channel will thus be selected based on highest correlation with the published optimal connectivity profile. An openly available connectome for selection of respective ROI regions can be obtained: DOI: 10.5281/zenodo.10805915.

In [2]:

PATH_features = r"C:\Users\ICN_admin\Documents\Paper Decoding Toolbox\AcrossCohortMovementDecoding\features_out_fft"
PATH_CONN_CORR = r"C:\Users\ICN_admin\Documents\Paper Decoding Toolbox\AcrossCohortMovementDecoding\across patient running\RMAP\df_best_func_rmap_ch.csv"

ch_all = np.load(os.path.join(PATH_features,"channel_all.npy"), allow_pickle="TRUE").item()
df_best_rmap = pd.read_csv(PATH_CONN_CORR)

We now concatenate features from all patients of the Berlin cohort and leave out a single one as a test patient.

In [3]:

def get_sub_data(cohort, sub):

    ch_best = df_best_rmap.query("sub == @sub and cohort == @cohort").iloc[0]["ch"]
    runs = list(ch_all[cohort][sub][ch_best].keys())
    if len(runs) > 1:
        dat_concat = np.concatenate(
            [ch_all[cohort][sub][ch_best][run]["data"] for run in runs],
            axis=0,
        )
        lab_concat = np.concatenate(
            [ch_all[cohort][sub][ch_best][run]["label"] for run in runs],
            axis=0,
        )
    else:
        dat_concat = ch_all[cohort][sub][ch_best][runs[0]]["data"]
        lab_concat = ch_all[cohort][sub][ch_best][runs[0]]["label"]
    return dat_concat, lab_concat

cohort = "Berlin"
sub_test = list(ch_all[cohort].keys())[5]

dat_test, lab_test = get_sub_data(cohort, sub_test)
lab_test = gaussian_filter1d(np.array(lab_test, dtype=float), sigma=1.5)
dat_train = []
lab_train = []
for sub_train in ch_all[cohort].keys():

    if sub_test == sub_train:
        continue

    dat_concat, lab_concat = get_sub_data(cohort, sub_train)

    label_cont = gaussian_filter1d(np.array(lab_concat, dtype=float), sigma=1.5)
    dat_train.append(dat_concat)
    lab_train.append(label_cont)

dat_train = np.concatenate(dat_train, axis=0)
lab_train = np.concatenate(lab_train, axis=0)

Next, we define the CEBRA model architecture. We use here the default offset10-model and a 3D embedding dimension and fit the model.

In [4]:

embedding_model = CEBRA(
        model_architecture="offset10-model",
        batch_size=100,
        temperature_mode="auto",
        learning_rate=0.005,
        max_iterations=2000,
        time_offsets=1,
        output_dimension=3,
        device="cpu",
        verbose=True,
    )
embedding_model.fit(dat_train, lab_train)

pos: -9.3236 neg:  13.6278 total:  4.3041 temperature:  0.1000: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2000/2000 [05:20<00:00,  6.24it/s]

Out[4]:

CEBRA(batch_size=100, device='cpu', learning_rate=0.005, max_iterations=2000,
      model_architecture='offset10-model', output_dimension=3,
      temperature_mode='auto', verbose=True)

In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook.
On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.

The model architecture consists of several 1D convolutional layers with skip connections:

In [5]:

embedding_model.model_

Out[5]:

Offset10Model(
  (net): Sequential(
    (0): Conv1d(7, 32, kernel_size=(2,), stride=(1,))
    (1): GELU(approximate='none')
    (2): _Skip(
      (module): Sequential(
        (0): Conv1d(32, 32, kernel_size=(3,), stride=(1,))
        (1): GELU(approximate='none')
      )
    )
    (3): _Skip(
      (module): Sequential(
        (0): Conv1d(32, 32, kernel_size=(3,), stride=(1,))
        (1): GELU(approximate='none')
      )
    )
    (4): _Skip(
      (module): Sequential(
        (0): Conv1d(32, 32, kernel_size=(3,), stride=(1,))
        (1): GELU(approximate='none')
      )
    )
    (5): Conv1d(32, 3, kernel_size=(3,), stride=(1,))
    (6): _Norm()
    (7): Squeeze()
  )
)

We can now also investigate if the model training converged by plotting the loss function.

In [6]:

cebra.plot_loss(embedding_model)
plt.title("Model training loss")

Out[6]:

Text(0.5, 1.0, 'Model training loss')

The training and test data can then be transformed through the embedding model and used for fitting a simple linear regression model for prediction of the test subject.

In [7]:

embedding_train = embedding_model.transform(dat_train)
embedding_test = embedding_model.transform(dat_test)

ML_model = linear_model.LinearRegression()
ML_model.fit(embedding_train, lab_train)
pred = ML_model.predict(embedding_test)

The training and test embeddings can be visualized and color-coded by the movement labels. Note that we used a gaussian filter to approximate movement kinematics to use the continuous CEBRA fit.

In [12]:

fig = plt.figure(figsize=(4, 2), dpi=300)
for i in range(2):

    if i == 0:
        emb_ = embedding_train[::10]
        label_ = lab_train[::10]
        ax = plt.subplot(121, projection="3d")
        plt.title("Train embedding", fontsize=8)
    else:
        emb_ = embedding_test
        label_ = lab_test
        ax = plt.subplot(122, projection="3d")
        plt.title("Test embedding",  fontsize=8)
    
    x = ax.scatter(
        emb_[:, 0],
        emb_[:, 1],
        emb_[:, 2],
        c=label_,
        cmap="viridis",
        s=0.02,
        vmin=0,
        vmax=1,
        alpha=0.8,
    )
    ax.view_init(elev=-80, azim=250)
    plt.axis("off")

fig.subplots_adjust(right=0.8)
cbar_ax = fig.add_axes([0.85, 0.4, 0.02, 0.2])
cbar = fig.colorbar(x, cax=cbar_ax, shrink=0.1)
cbar.set_label("Movement")

Finally, the test set predictions are plotted against the true movements:

In [13]:

plt.figure(figsize=(5, 3), dpi=300)
time_ = np.arange(0, 1000, 1) * 0.1

plt.plot(
    time_,
    lab_test[500:1500],
    label="True",
)
plt.plot(time_, pred[500:1500], label="Predicted")
plt.legend()
plt.xlabel("Time [s]")
plt.ylabel("Amplitude [a.u.]")
plt.title(
    f"Movement prediction\nwithout individual training\n"
    f"corr={np.round(np.corrcoef(lab_test, pred)[0, 1], 2)}"
)
plt.tight_layout()