SPARKS ← All demos Documentation Download this notebook (GitHub)

Oscillatory sequences in the medial entorhinal cortex

Can SPARKS reveal structure in neural data without being told what to look for? This notebook trains it, unsupervised, on a calcium-imaging recording of about 500 neurons in the medial entorhinal cortex (MEC) of a mouse, from Gonzalo Cogno et al., Nature, 2024. That study found that this population fires in slow sequences, which repeat every few minutes.

You will need the demos extra: pip install "sparks-ai[demos]". The recording is downloaded from the paper's repository on first use.

Training takes about 30 minutes on a recent GPU.

import os
import urllib.request

import matplotlib.pyplot as plt
import numpy as np
import torch
from tqdm import tqdm

import sparks
from sparks.data.mec import make_mec_data
from sparks.models import LinearDecoder
from sparks.training import build_param_groups, get_lr_multiplier

Data

The recording holds the neurons' deconvolved, binarised spikes, sampled at 7.73 Hz. We bin them 4 samples at a time, into bins of about 0.5 s. Since the model learns without labels, it is trained and evaluated on the same 50 minutes of recording, from 600 s to 3600 s.

data_path = 'calcium_activity_matrix_60584_session17.mat'
if not os.path.exists(data_path):
    urllib.request.urlretrieve('https://github.com/soledadgcogno/Ultraslow-oscillatory-sequences/raw/main/'
                               'Selection%20of%20codes/calcium_activity_matrix_60584_session17.mat', data_path)

window = np.array([[600, 3600]])  # start and stop times, in seconds
data = make_mec_data(data_path, train_times=window, test_times=window, downsampling_factor=4)

spikes = data.test_dataset.get_spikes(0).numpy()  # [time, neurons]
dt = data.test_dataset.dt
print(f'{spikes.shape[1]} neurons, {spikes.shape[0]} time bins of {dt:.2f} s')

fig, ax = plt.subplots(figsize=(14, 4))
ax.imshow(spikes.T, aspect='auto', cmap='Greys', interpolation='none', extent=[0, len(spikes) * dt, spikes.shape[1], 0])
ax.set_xlabel('Time (s)')
ax.set_ylabel('Neuron')
plt.show()

484 neurons, 5798 time bins of 0.52 s

In this order, the neurons' spikes look random.

Model

The Hebbian attention layer's STDP time constant (tau_s) is long, 500 s, to match the slow timescale of the dynamics. The model compresses the population into a 2-dimensional latent at each time bin, from which a linear decoder reconstructs the spikes of every neuron.

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # on Apple silicon, you can try 'mps'
torch.manual_seed(0)

latent_dim = 2
n_neurons = data.train_dataset.n_neurons

model = sparks.SPARKS(n_neurons_per_session=n_neurons,
                      embed_dim=32,
                      latent_dim=latent_dim,
                      bottleneck_dim=8,
                      hebbian_config=sparks.HebbianAttentionConfig(tau_s=500., dt=dt),
                      decoder=LinearDecoder(in_dim=latent_dim, output_dim_per_session=n_neurons),
                      device=device)

Training

• unsupervised=True trains the model to predict its inputs: it is a variational autoencoder.
• The latents are regularised toward a learned autoregressive prior (kl_type='sequential'), with a weight beta.
• The recording is processed in chunks of 1000 time bins (chunk_length), carrying the model's state from one chunk to the next.

n_epochs = 2000
beta = 1e-4

param_groups, transition = build_param_groups(model, weight_decay=0.1, kl_type='sequential')
optimizer = torch.optim.AdamW(param_groups, lr=0.01)
total_steps = n_epochs * len(data.train_loader)
scheduler = torch.optim.lr_scheduler.LambdaLR(
    optimizer, lambda step: get_lr_multiplier(step, warmup_steps=int(0.05 * total_steps), total_steps=total_steps, min_lr=1e-4))
loss_fn = torch.nn.BCEWithLogitsLoss()

pbar = tqdm(range(n_epochs))
for epoch in pbar:
    sparks.train(model, [data.train_loader], loss_fn, optimizer, scheduler=scheduler, unsupervised=True,
                 beta=beta, kl_type='sequential', transition=transition, chunk_length=1000)

    if (epoch + 1) % 50 == 0:
        test_loss, _, _ = sparks.evaluate(model, [data.test_loader], loss_fn, unsupervised=True, chunk_length=1000)
        pbar.set_description(f'test loss: {test_loss.item():.4f}')

test loss: 0.1919: 100%|██████████| 2000/2000 [27:53<00:00,  1.20it/s]

A ring in the latent space

Over time, the latents go round and round a ring in the 2-dimensional latent space.

_, latents, _ = sparks.evaluate(model, [data.test_loader], unsupervised=True, chunk_length=1000)
latents = latents[0][0].numpy()  # the only session and trial: [time, 2]
time = np.arange(len(latents)) * dt

fig, ax = plt.subplots(figsize=(6, 5))
points = ax.scatter(latents[:, 0], latents[:, 1], c=time, cmap='magma', s=3)
fig.colorbar(points, label='Time (s)')
ax.set_xlabel('Neural axis 1')
ax.set_ylabel('Neural axis 2')
plt.show()

The position around the ring is the phase of an oscillation: its angle around the ring's centre.

centre = latents.mean(axis=0)
phase = np.arctan2(latents[:, 1] - centre[1], latents[:, 0] - centre[0])

fig, ax = plt.subplots(figsize=(14, 3))
ax.plot(time, phase)
ax.set_xlabel('Time (s)')
ax.set_ylabel('Phase (rad)')
plt.show()

Sequences in the neurons' activity

Each neuron's preferred phase is the circular mean of the phases at which it spikes. Sorting the neurons by their preferred phase reveals the sequences: the population fires in the same order on every cycle.

preferred_phase = np.angle((spikes * np.exp(1j * phase)[:, None]).sum(axis=0))
order = np.argsort(preferred_phase)

fig, axs = plt.subplots(2, 1, figsize=(14, 8), sharex=True)
for ax, neurons, title in zip(axs, [np.arange(spikes.shape[1]), order], ['Recorded order', 'Sorted by preferred phase']):
    ax.imshow(spikes[:, neurons].T, aspect='auto', cmap='Greys', interpolation='none',
              extent=[0, len(spikes) * dt, spikes.shape[1], 0])
    ax.set_title(title)
    ax.set_ylabel('Neuron')
axs[1].set_xlabel('Time (s)')
plt.show()