Skip to content

Plating-Based Semantics for Hierarhical Inference#63

Draft
DanWaxman wants to merge 33 commits intomainfrom
dw-add-plating
Draft

Plating-Based Semantics for Hierarhical Inference#63
DanWaxman wants to merge 33 commits intomainfrom
dw-add-plating

Conversation

@DanWaxman
Copy link
Collaborator

WIP.

Resolves #33. Resolves #43.

The `CDNLGSSM` class of cd_dynamax does not have a vmap-able path sample method. For uniformity, we want to change to calling diffrax directly.
This provides a more uniform interface for simulating SDEs: it is similar to the ODE, is directly vmap-able, and avoids some unecessary dependence on cd_dynamax.

This is partially WIP, since some sampling in the SDE needs to be resolved

Fixes #43.
This fixes some linting, and also adds back in a log-factor.
Partially LLM-generated, with lots of direction for making things composable with vmap, automatically inferring the batch shape from _PYRO_STACK, and rewriting things in terms of eqx modules for this to happen automatically. Will be checking carefully and removing verbosity next.
This adds cuthbert support using the Taylor submodule, which is just doing an EnKF I think. Subsequent commits will work with additional inference methods.
We add in particle filters, add defaults, and properly separate continuous and discrete filtering methods.

This also adds a small patch to cuthbert to ensure differentiability.
They accepted my PR :)
Fixes #58)

Types are kinda messed up right now, so this is mostly a lot of type ignores. These should be resolved, but potentially as a separate PR.
This adds new integration subpackages. For now, these live in `dynestyx/inference/<integration>/<...>.py`, and include continuous and discrete time filters. This makes the base `filters.py` much more wieldy.
This introduces stricter linting rules, useful for things like unused imports or performance issues.
This streamlines use a bit and properly separates cd_dynamax utils from the rest of the library.
Continuous-time features which would have used an EnKF previously now use all the Gaussian filters.

This may lead to some instabilities for certain methods (EKF bad...).
EKF too bad performance. Make sure unapcking shapes are right.
Idea: ensure appropriate batch dimensions before passing forward.
@DanWaxman
Copy link
Collaborator Author

Closing this as superseded by #173

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.

Support (hierarchical) inference over multiple trajectories Change SDESimulator to Call Diffrax Directly

2 participants