Skip to content

Add some state space model examples ported from dynamax#267

Open
theorashid wants to merge 1 commit into
BasisResearch:dsx-infer-stagingfrom
theorashid:examples-ssm-from-dynamax
Open

Add some state space model examples ported from dynamax#267
theorashid wants to merge 1 commit into
BasisResearch:dsx-infer-stagingfrom
theorashid:examples-ssm-from-dynamax

Conversation

@theorashid

Copy link
Copy Markdown

Adds a "State Space Models" section adapted from some dynamax examples and reimplemented with the dynestyx API. Some changes from dynamax:

  • Filters/smoothers use the cuthbert backend; ekf_mlp uses cd_dynamax for its Jacobian-based EKF (cuthbert's Taylor EKF doesn't fit the MLP)
  • Dropped EM in lgssm_learning (dynestyx has no EM); Bayesian route uses NUTS
  • ekf_spiral is EKF-only
  • MLP rewritten with flax.nnx
  • mkdocs.yml: new nav section + MathJax switched to SVG (fixes matrix delimiter rendering site-wide)

ekf_mlp emits a harmless cd_dynamax UserWarning

UserWarning: cd_dynamax discrete-time filters ignore absolute time arguments in GaussianStateEvolution/GaussianObservation. For genuinely time-varying discrete-time models, use a filter/backend that preserves absolute time semantics, such as EnKFConfig(filter_source='cuthbert').
  params_nl = gaussian_to_nlgssm_params(dynamics)

Don't think this is relevant here?

cc @juanitorduz: this should help with #264 to see how to do basic SSMs. We should port over an example for a basic ETS model (or similar structural model). At some point, I'd like to do an example like the online MLP where we get a new data point (in time) and we online update both the filtered forecast and the hyperparameter estimate. I know they're looking into a similar (non-forecasting) example in cuthbert for this too.

Five worked examples under docs/tutorials/state_space_models/, adapted from
the dynamax notebooks and reimplemented with the dynestyx API
@DanWaxman

Copy link
Copy Markdown
Collaborator

Thanks Theo!! This looks great! I'll take a closer look after #266 merges (and try to take care of any merge conflicts that introduces).

UserWarning: cd_dynamax discrete-time filters ignore absolute time arguments in GaussianStateEvolution/GaussianObservation. For genuinely time-varying discrete-time models, use a filter/backend that preserves absolute time semantics, such as EnKFConfig(filter_source='cuthbert'). params_nl = gaussian_to_nlgssm_params(dynamics)

Don't think this is relevant here?

I agree

@juanitorduz

Copy link
Copy Markdown

Thanks @theorashid ! Let me take a detailed look 🙏 ! Exiting!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants