StateSpaceDynamics.jl: A Julia package for probabilistic state space models (SSMs)

Description

StateSpaceDynamics.jl is a comprehensive and self-contained Julia package for working with probabilistic state space models (SSMs). It implements a wide range of state-space models, taking inspiration from the SSM package written in Python by the Linderman Lab. This package is designed to be fast, flexible, and all-encompassing, leveraging Julia's speed and expressiveness to provide researchers and data scientists with a powerful toolkit for state-space modeling.

This package is geared towards applications in neuroscience, so the models incorparate a certain neuroscience flavor (e.g., many of our models are trialized as common in experiemntal paradigms). However, the models are general enough to be used in other fields such as finance, robotics, and many other domains involving sequential data analysis.

We are continuously working to expand our model offerings. If you have suggestions for additional models or features, please open an issue on our GitHub repository.

Installation

You can install StateSpaceDynamics.jl using Julia's package manager. From the Julia REPL, type ] to enter the Pkg REPL mode and run:

pkg> add StateSpaceDynamics

Usage

StateSpaceDynamics.jl is designed to be user friendly with intuitive syntax. Below is an example of how to fit a Poisson Linear Dynamical System (PLDS) to data.

using StateSpaceDynamics
using LinearAlgebra

# create a toy system

# initial conditions
x0 = [1.0, -1.0] # initial state
P0 = Matrix(Diagonal([0.1, 0.1])) # initial state covariance

# state model parameters
A = [cos(0.1) -sin(0.1); sin(0.1) cos(0.1)] # transition matrix
Q = Matrix(Diagonal([0.01, 0.01])) # process noise

# observation model parameters
C = [1.2 1.2; 1.2 1.2; 1.2 1.2] # observation matrix
log_d = log.([0.1, 0.1, 0.1]) # log of the natural parameters of the Poisson distribution

# generate data
tSteps = 100
trials = 10

true_plds = PoissonLDS(;A=A, Q=Q, C=C, log_d=log_d, x0=x0, P0=P0, obs_dim=3, latent_dim=2)
latents, observations = sample(true_plds, tSteps, trials)

# fit the model 
plds = PoissonLDS(;obs_dim=3, latent_dim=2)
fit!(plds, observations)

Available Models

• Mixture Models
  • Gaussian Mixture Models
  • Poisson Mixture Models
  • Binomial Mixture Models
  • Negative Binomial Mixture Models
  • Student's t Mixture Models
• Hidden Markov Models
  • Gaussian HMMs
  • Poisson HMMs
  • Binomial HMMs
  • Negative Binomial HMMs
  • Autoregressive HMMs (ARHMM)
• Linear Dynamical Systems
  • Gaussian Linear Dynamical Systems (Kalman Filter)
  • Poisson Linear Dynamical Systems (PLDS)
  • PFLDS
  • Switching Linear Dynamical Systems (SLDS)
  • Recurrent Switching Linear Dynamical Systems (rSLDS)
• Generalized Linear Models (Not state space models but needed for HMM-GLMs)
  • Gaussian GLMs
  • Poisson GLMs
  • Binomial GLMs
  • Negative Binomial GLMs
• HMM-GLM's
  • Gaussian HMM-GLMs
  • Poisson HMM-GLMs
  • Bernoulli HMM-GLMs
  • Negative Binomial HMM-GLMs
  • Multinomial HMM-GLMs

Related Packages

• HiddenMarkovModels.jl: A Julia package for Hidden Markov Models.

• StateSpaceLearning.jl : A Julia package for time series analysis using state space models.

• ssm : A python package for state space models.

References

1. Bishop, C. M. (2006). Pattern recognition and machine learning. Springer.

2. Murphy, K. P. (2012). Machine learning: a probabilistic perspective. MIT press.

3. Sarrka S., Svensson, L. (2023). Bayesian Filtering and Smoothing. Cambridge University Press.

4. Paninski, L. et al. (2010). A new look at state-space models for neural data. Journal of computational neuroscience, 29(1-2), 107-126.

5. Macke, J. H. et al. (2011). Empirical models of spiking in neural populations. Advances in neural information processing systems, 24, 1350-1358.