cassm 0.1.0

Computation-aware State-Space Models

CASSM

Computation-aware state-space models for neural data.

CASSM provides PyTorch implementations of computation-aware filtering and smoothing models for high-dimensional neural time series. The package includes synthetic data utilities, model implementations, metrics, and a short tutorial notebook for training on Lorenz-generated spike trains.

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,
    valid_truth=valid_truth,
    clip_value=300,
)

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.