JAX MoSeq

jax-moseq is a library for fitting state-space models with Gibbs sampling. If you would like to apply these models to keypoint tracking data, we recommend using the keypoint MoSeq front-end, which has its own documentaton and installation instructions. The code is written in a functional style using jax. We also provide light-weight classes for model fitting. We currently support the following models, all of which use a sticky hierarchical dirichlet process (HDP) prior by default.

• Vector autoregressive (VAR)

• Switching Linear Dynamical System (SLDS)

• SLDS with keypoint emissions

Installation

Install JAX as described in the [Keypoint MoSeq docs](https://keypoint-moseq.readthedocs.io/en/latest/), then install jax-moseq

pip install -U jax-moseq

Getting started

Checkout the example notebooks

API documentation

• AR-HMM
  • Initialization
  • Gibbs sampling
  • Log probability
• SLDS
  • Initialization
  • Gibbs sampling
  • Log probability
• Keypoint-SLDS
  • Initialization
  • Gibbs sampling
  • Log probability
  • Geometric transforms
• Utilities
  • Distributions
  • Transitions
    • count_transitions()
    • sample_betas()
    • sample_pi()
    • sample_hdp_transitions()
    • resample_hdp_transitions()
    • init_hdp_transitions()
    • sample_dir_transitions()
    • resample_dir_transitions()
  • Kalman smoothing
  • Autoregression
  • Odds and ends
    • concatenate_stateseqs()
    • get_durations()
    • get_frequencies()
    • symmetrize()
    • psd_solve()
    • psd_inv()
    • jax_io()
    • pad_affine()
    • fit_pca()
    • wrap_angle()
    • pad_along_axis()
    • unbatch()
    • batch()
    • get_mixed_map_iters()
    • set_mixed_map_iters()
    • get_mixed_map_gpus()
    • set_mixed_map_gpus()
    • mixed_map()