ITS: add Bayesian Structural Timeseries model

[drbenvincent]

• This would lead to feature parity with CausalImpact I believe.
• There is a jax implementation here https://github.com/probml/sts-jax
• There is a tensorflow implementation here https://github.com/WillianFuks/tfcausalimpact
• See the issue I created (Bring structural time series functionality to PyMC pymc-devs/pymc#6284) to bring STS functionality to PyMC. This is probably a blocker for progressing this issue.
• An short term alternative would be to call the jax implementation from within PyMC.

[NathanielF]

What were you thinking for this? MVP using jax? Or waiting for a PYMC implementation?

[drbenvincent]

See #116

A pymc implementation would be great, but unsure of the timescale on that. So JAX wrapper seems like the quickest thing. Not sure if @juanitorduz had plans to work on this. But I think we'd be very happy if you wanted to work on it

[NathanielF]

That'd be cool! If there is no objection from. @juanitorduz I'll look into your contribution guide this weekend. Not sure how quick i can be on this, but from a cursory look allot of the heavy lifting seems done already.

[juanitorduz]

Hey! I'm just starting exploring the sts-jax library so feel free to start working on it! I guess we are open regarding the library and framework for a first iteration. I will try to catch up! 🙂 !

[juanitorduz]

Also, check out this project: https://github.com/jamalsenouci/causalimpact uses both statsmodels and pymc

[drbenvincent]

That'd be cool! If there is no objection from. @juanitorduz I'll look into your contribution guide this weekend. Not sure how quick i can be on this, but from a cursory look allot of the heavy lifting seems done already.

Sounds good. Just let me know if you need any pointers. There is a guide there, but we aren't open for general contributors yet, so it's not battle tested.

[NathanielF]

Sorry haven't found the time to look at this properly. Trying to get more into jax following Ravin's book club on Dynamax... but honestly not sure when i'll get to this.

[drbenvincent]

Just dropping in a link to the developments on state space models by @jessegrabowski here https://discourse.pymc.io/t/state-space-models-in-pymc/9146

[drbenvincent]

Tagging @cetagostini

[jessegrabowski]

How do people do causality in STS framework? I'm interested in bringing more tools related to this into what we've already got in the statespace module. I know about granger test and rotated (cholesky decomposed) IRFs, but is there more? Also does the granger test make sense in a Bayesian framework? It always struck me as a glorified t-test.

[drbenvincent]

It would be great to get your input @jessegrabowski. The causal side of things is rich, but possibly the most direct way of working out how time series approaches are used in causality would be to look at Interrupted Time Series:

• PyMC Example - excess deaths: https://www.pymc.io/projects/examples/en/latest/causal_inference/excess_deaths.html
• Excess deaths example in CausalPy docs: https://causalpy.readthedocs.io/en/latest/notebooks/its_covid.html

These examples take the simplistic "linear regression as time series" approach. The next step (which CausalImpact uses) is to Bayesian structural time-series models.

It should be relatively trivial to plug that into CausalPy. We'd just need to adapt the pymc_models.LinearRegression class so that the build_model factory function returns a time series model. Ideally we'd retain the same kind of Wilkinson notation model formula to provide predictor variables, but it's not the end of the world if a slightly different API is needed.

Could be wrong, but I think it's the kind of thing that we could get done in 1 (or just a few) pair coding sessions.