Computation-Aware State-Space Model (CASSM)

Latent variable dynamical modeling is a classical problem in neuroscience.

Typically, we assume neural population acitivity lies in some low-dimensional space. In the figure above, the synthetic dynamics are generated by the 3D Lorenz system.

A readout map defines the interaction between neurons as we move from latent space to neural space, and then we sample from that trajectory to obtain spikes.

Latent variable dynamical models attempt to learn the latent dynamics, readout maps, and noise models for this problem.

CASSM provides PyTorch implementations of computation-aware filtering and smoothing models for high-dimensional neural time series. Unlike the original Computation-Aware Kalman Filter (CAKF), here we also learn the models themselves. Just bring your data, and we'll learn everything else.

In effect, we are solving Gaussian Process regression problems efficiently using Bayesian Filtering and Smoothing: no more cubic-time algorithms in the sequence length. After CASSM, Gaussian Process Factor Analysis (GPFA) should only be used for very small recordings. Any larger recordings with long sequence lengths and many neurons will want to use this repository.

Please see the paper for more details.

Installation

Install from PyPI:

pip install cassm

Install the package in editable mode from the repository root:

pip install -e .

For development and tests, install the optional development dependencies:

pip install -e ".[dev]"

Quick Start

import torch
from torch.utils.data import DataLoader, TensorDataset

from cassm.datasets.synthetic_data import LorenzData
from cassm.models import CASSM
from cassm.utils.preprocessing import smooth_firing_rate

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

train_data, valid_data, _, valid_truth, _, _ = LorenzData(
    num_inits=4,
    neurons=120,
    num_trials=4,
    device="cpu",
    seed=2,
)

train_data = smooth_firing_rate(train_data.numpy()).to(device)
valid_data = smooth_firing_rate(valid_data.numpy()).to(device)
valid_truth = valid_truth.to(device)

train_loader = DataLoader(train_data, batch_size=4, shuffle=True)
test_loader = DataLoader(TensorDataset(valid_data, valid_truth), batch_size=4)

model = CASSM(
    projection_dim=10,
    nneurons=train_data.shape[-1],
    timesteps=train_data.shape[1],
    dt=0.01,
    device=device,
    save_model=False,
).to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=5e-2)
model.train_model(
    epochs=3,
    optimizer=optimizer,
    train_loader=train_loader,
    test_loader=test_loader,
)

# filtering output
filtered_state, filtered_noise = model.filter(valid_data, return_type="prediction")

# trial averaged predictions
predicted_rate, predicted_noise = model.predict_rate(valid_data)

See tutorials/cassm_lorenz_tutorial.ipynb for a fuller walkthrough.

Repository Layout

• src/cassm: package source code
• tests: package tests
• tutorials: user-facing notebooks
• pyproject.toml: package metadata and tooling configuration

Testing

pytest

License

CASSM is released under the MIT License. See LICENSE.