Integration with SSMProblems Ecosystem

[THargreaves]

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.:

• You can forward simulate from a model by repeatedly calling predict
• In data association settings (e.g. target tracking) you need to predict then associate measurements with targets, then update conditioned on the associations

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

[gdalle]

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

[gdalle]

Related issue:

• Clarify forecast / simulate / sample, especially with controls #96

[THargreaves]

Related issue:

• Clarify forecast / simulate / sample, especially with controls #96

Ah nice. Yes, the SSMProblems.jl interface is designed to handle forward simulation and future prediction—not entirely sure what they mean by sampling though but I think what they're after is covered by these first two.

[gdalle]

The new release v0.6 dropped with the hidden functions you wanted and a bunch of performance improvements for single-sequence workflows. Let me know if you need more stuff!