variational Latent Gaussian Process with JAX

Introduction

This repo contains a JAX implementation of variational Latent Gaussian Process (vLGP) (arXiv) (video) by Yuan Zhao (yuan.zhao@stonybrook.edu) and Il Memming Park (memming.park@stonybrook.edu). It has been developed with the goal of recovering low-dimensional dynamics from neural population recordings.

Installation

pip install git+https://github.com/yuanz271/vlgpax.git

Get started

Learn how to use it in the example.

Data structure

• Trial A single trial
  • y spike train
  • x design matrix of regressors, see pyneuroglm for building design matrix
  • t timing of bins, useful for unevenly spaced timing and nonstationary kernels
  • z factor posterior mean
  • v factor posterior variances
• Session A container of trials
  • trials list of Trial
  • binsize binwidth if evenly spaced
• Params Parameters
  • n_factors number of latent factors
  • C loading matrix, (n_factors + n_regressors, n_neurons)
  • K kernel matrices
  • args algorithm settings

Training

from vlgpax import Session, vi
from vlgpax.kernel import RBF

binsize = 1e-2
session = Session(binsize)
# Add trials to the session
session.add_trial(tid=1, y=y)  # ID, spike train, ...
# ...

# Kernel function k(x, y) = scale * exp(-0.5 ||x/lengthscale - y/lengthsale||^2)
kernel = RBF(scale=1., lengthscale=100.)  # lengthscale has the same unit as that of binsize

session, params = vi.fit(session, n_factors=2, kernel=kernel)

Citation

@Article{Zhao2017,
  author    = {Yuan Zhao and Il Memming Park},
  title     = {Variational Latent Gaussian Process for Recovering Single-Trial Dynamics from Population Spike Trains},
  journal   = {Neural Computation},
  year      = {2017},
  volume    = {29},
  number    = {5},
  pages     = {1293--1316},
  month     = {may},
  doi       = {10.1162/neco_a_00953},
  publisher = {{MIT} Press - Journals},
}