gdalle / HiddenMarkovModels.jl

A Julia package for simulation, inference and learning of Hidden Markov Models.
https://gdalle.github.io/HiddenMarkovModels.jl/
MIT License
85 stars 6 forks source link

Integration with SSMProblems Ecosystem #104

Open THargreaves opened 3 days ago

THargreaves commented 3 days ago

The Turing team have been developing SSMProblems, a generic interface for defining state space model and their associated filters/smoothers. The goal of this project is to create a flexible and modular ecosystem that can easily be extended with new algorithms, or applied to new models.

As a test case for the ecosystem, we wanted to wrap the implementation of the forward-backward algorithm provided by HiddenMarkovModels.jl so that it is compatible with SSMProblems.jl. This would allow it to integrate seamlessly with implementations we are developing for the interface such as Rao-Blackwellisation and PMCMC methods.

Since HiddenMarkovModels.jl adopts a similar philosophy to SSMProblems.jl when it comes to control variables, wrapping the algorithms is generally quite simple. The only trouble I'm running into is that SSMProblems.jl demands a bit more structure within the filtering loop. Specifically, it breaks the filtering up into an initialisation step before repeated alternating steps of predicting the next state and updating the state given the observations.

To demonstrate this, I've taken the code from src/inference/forward.jl, copied it and shuffled it around into some functions. This is presented in this Gist: https://gist.github.com/THargreaves/db47fdc4bb715eb69474050473f56817

As far as I can see there are no downsides to this change on the HiddenMarkovModels.jl but I wouldn't be interested to see any pointed out if that is not the case.

On the other hand, this dissection allows for a few additional use cases of the same code. E.g.:

This issue is intended to open a discussion as to whether this structure could be adopted by HHMs.jl.

gdalle commented 3 days ago

At first glance I see no major issue with adopting this framework, but it would help me if you could open a PR to HMMs implementing exactly that, without the wrapper and extra shenanigans. Just to see it in context