LGSSM with diagonal noise

[calebweinreb]

Changes

This PR addresses issue https://github.com/probml/dynamax/issues/331 by allowing the emissions covariance to be stored as a static or time-varying 1D array containing just the diagonal entries and optimizing certain computations when this is the case. Changes include:

• Modified the type signature for ParamsLGSSMEmissions
• Rewrote _get_params in inference.py and parallel_inference.py because we can longer rely on the dimension of params to know if they are time-dependent. In the new version, a param is assumed to be time-dependent if X.shape[0]==num_timepoints.
• When emissions cov is diagonal, use Woodbury identity to compute Kalman gain in the _condition_on step of serial Kalman filtering, and for computing the conditional emissions precision matrix used in parallel Kalman filtering.
• Use MultivariateNormalDiagPlusLowRankCovariance from TFP to compute emissions log-likelihoods when emissions cov is diagonal (for both serial and parallel inference)
• Modified all tests so that they are repeated once for full-shape emissions cov and once for reduced-shape.

Benchmark

The optimizations seem to speed things up when the emissions dim surpasses ~50.

• Here's a comparison of run times for _condition_on with full vs. diagonal emissions cov

  

• Here are run times for getting obs log likelihood using full vs. low-rank MVN normal from tensorflow probability

  

• Here are run times for parallel LGSSM filtering with full vs. diagonal emissions cov, where both the Woodbury identity and the low-rank MVN are used in the diagonal case.

  

Benchmarking code: https://gist.github.com/calebweinreb/35bcb53f25d83d0f002b731ccca3b91a

[slinderman]

Thanks @calebweinreb! This is a really nice feature. I requested a few minor changes above.

[calebweinreb]

The requested changes are implemented here https://github.com/probml/dynamax/commit/9c202d853f247125cfc4196376b4fe70517128cb

I'm a little confused about the tests failing since it seems to be a dependency issue... In any casepytest dynamax passes when I run it locally.

[slinderman]

I think the tests were failing because they were still using python 3.7. I just updated our tests to use 3.10.6 (current Colab version) and they seem to run. Can you merge the latest changes from main and push again? That should trigger another run of the tests.

[calebweinreb]

Hey! Looks like all the tests pass now.

[bantin]

Thanks for all your work on this @calebweinreb. Sorry to revive an old thread -- is there an easy way of fitting an LDS with diagonal emissions covariance? As best I can tell, the EM code in models.py assums a full covariance.