probml / dynamax

State Space Models library in JAX
https://probml.github.io/dynamax/
MIT License
667 stars 76 forks source link

Add support for diagonal emissions noise in LGSSM models #331

Closed calebweinreb closed 1 year ago

calebweinreb commented 1 year ago

For LGSSMs with a high-dimensional observation space, inference can be much more efficient by assuming diagonal covariance R. Specifically, it should be possible to speed up the calculation of Kalman gain and observation log-likelihoods, and also to reduce memory requirements when the noise is dynamic by storing R as a (ntime x obs_dim) array rather than (ntime x obs_dim x obs_dim).