README.md

cd-dynamax source code description

We provide the following modifications of the dynamax codebase, to accommodate continuous-discrete models, i.e., those where observations are not assumed to be regularly sampled.

Continuous-time State Space Models with Emissions at Specified Discrete Times

• A modified version of dynamax's ssm.py that incorporates non-regular emission time instants: i.e., the t_emissions array
  • t_emissions is an input argument
    • We use t0 and t1 refer to $t_k$ and $t_{k+1}$, not necessarily regularly sampled
  • t_emissions is a matrix of size $[\textrm{num observations} \times 1]$
    • it should facilitate batching
    • For lax.scan() operations, we recast them in vector shape (i.e., remove final dimension)

Continuous-Discrete Linear Gaussian State Space Models

• We define a ContDiscreteLinearGaussianSSM model

  • We do not currently provide a ContDiscreteLinearGaussianConjugateSSM model implementation, as CD parameter conjugate priors are non-trivial

  • The CD-LGSSM model is based on

    • A continuous-time push-forward operation that computes and returns matrices A and Q
      • based on Equation (3.135) in [1] Särkkä, Simo. Recursive Bayesian inference on stochastic differential equations. Helsinki University of Technology, 2006.

• Continuous-Discrete Kalman filtering and smoothing algorithms are implemented

• Parameter (point)-estimation is possible via stochastic gradient descent based MLE

  • where the marginal log-likelihood is computed based on the CD-Kalman filter

Continuous-Discrete Nonlinear Gaussian State Space Models

• We define a ContDiscreteNonlinearGaussianSSM model

  • The CD-NLGSSM model is based on a continuous-time push-forward operation that solves an SDE forward over the mean $x$ and covariance $P$ of the latent state
    • the parameters of the SDE function are provided in the ParamsCDNLGSSM object, which contains
      • The initial state's prior parameters in ParamsLGSSMInitial, as defined by dynamax
      • The dynamics function in ParamsCDNLGSSMDynamics
      • The emissions function in ParamsCDNLGSSMEmissions
        • These two latter are learnable functions

• Different filtering and smoothing algorithms are implemented

• Parameter (point)-estimation is possible via stochastic gradient descent based MLE

  • the marginal log-likelihood can be computed according to different implemented filtering methods (EKF, UKF, EnKF)

utils

• diffrax_utils.py

  • implements a diffrax based, autodifferentiable ODEsolver

• test_utils.py

• plotting_utils.py

• debug_utils.py

  • Debugging in jax can be difficult---pre-compilation speedups cause typical usage of in-line python debuggers to fail. To make debugging easier, we implemented a wrapper for lax.scan which, with debug=True, runs a (slow, but in-line debuggable!) for loop instead of lax.scan.
  • To use this in a particular piece of code, simply add from utils.debug_utils import lax_scan and replace an existing lax.scan call you wish to debug with lax_scan(..., debug=True).
  • This is an experimental feature, so please report any issues that arise from using this tool---we hope it helps ease the transition into using jax!

Tests

• Establishes functionality of linear and non-linear filters/smoothers, as well as parameter fitting via SGD.
• Checks that non-linear algorithms applied to linear problems return similar results as linear algorithms.

Notebooks

• Example notebooks, with filtering/smoothing of linear and nonlinear continuous-discrete dynamic models.