Implement Bayesian Structural Time Series (BSTS)

[cetagostini]

PR Description:

This PR introduces Bayesian Structural Time Series (BSTS) modeling to CausalPy, enabling advanced time series analysis, forecasting, and causal impact estimation. Inspired by principles from structural time series modeling basic with only trend/seasonal components (e.g., TensorFlow Probability's STS first example), this work decomposes time series into trend, seasonality, and optional regressor components within a Bayesian framework.

Key Changes:

1. BayesianStructuralTimeSeries Model (causalpy/pymc_models.py):

   • New PyMC model for BSTS. Defaults to pymc-marketing components (LinearTrend, YearlyFourier), supports custom components, and optional exogenous regressors (X).
   • Adapted _data_setter, predict, score, and fit for time-dependent forecasting and ensuring mu (sum of components) is sampled.

2. StructuredTimeSeries Experiment (causalpy/experiments/structured_time_series.py):

   • New experiment class for causal inference with BSTS.
   • Handles DatetimeIndex data, treatment_time, and patsy formulas for exogenous regressors.
     • Intercept Logic: Correctly manages formulas like y ~ 0 or y ~ 1 by ensuring the BSTS model's trend component provides the baseline, avoiding redundant patsy-generated intercepts for exogenous regressors.
   • Provides summary(), plot() (3-panel: fit/counterfactual, pointwise impact, cumulative impact), and get_plot_data().

3. Testing & Integration:

   • Added comprehensive integration tests (test_bayesian_structural_time_series) covering various scenarios, including custom components and error handling.
   • Fixed a plotting issue in plot_xY (causalpy/plot_utils.py).

4. API & Docs:

   • StructuredTimeSeries available as causalpy.StructuredTimeSeries.
   • Added backward-compatible wrapper in causalpy.pymc_experiments.py.
   • Updated relevant docstrings.

---

📚 Documentation preview 📚: https://causalpy--473.org.readthedocs.build/en/473/

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[cetagostini]

Thinking to bring HSGP to here, and be able to replicate the second example.

[drbenvincent]

WOAH!

Ultra quick review / questions until I've got some more time:

• Is this actually BSTS? It doesn't seem the use the BSTS stuff from @jessegrabowski in pymc-extras. If it's leaning heavily on linear trend and Fourier bases (from pymc-marketing, is this more like a prophet approach? No autoregressive components as far as I can tell? Still would be very very useful and an advancement on what we have so far.
• Not able to take a deep dive at this point, but my hope in terms of the API would be that we could use the existing InterruptedTimeSeries experiment class and simply inject a new time series based pymc model class. Will aim to look into this in more detail - if there was a reason for a new experiment class then we can look into whether we can adapt it to make it more general.
• Need to edit docs/source/notebooks/index.md to get the new notebook to render in the docs.
• The new notebook is a great start. Could be good to compare the previous approach "y ~ 1 + t + C(month)" in its_pymc.ipynb to see how things differ.

Sorry for the relatively superficial review at this point - family weekend time. Will enjoy diving into the details soon.

[cetagostini]

Hey @drbenvincent

Still BSTS, the Structural time series (STS) models are a family of probability models for time series that includes and generalizes many standard time-series modeling ideas, including:

• autoregressive processes,
• moving averages,
• local linear trends,
• seasonality, and regression and variable selection on external covariates (other time series potentially related to the series of interest).

Definitely, it's handy and popular to talk about state-space, but my understanding is that BSTS in the broad sense does not require a true state‐space; it merely requires a Bayesian, additive decomposition into trend, seasonality, and optional regressors. Unless you insist on the classic DLM/state‐space definition from a Kalman‐Filter sense.

In the Google blog the first model its very similar to this, the only difference is that the slope from the linear trend is "evolving slowly over time following a random walk".

Thanks to the HSGP class, we can add something similar if we want to.

[cetagostini]

Not able to take a deep dive at this point, but my hope in terms of the API would be that we could use the existing InterruptedTimeSeries experiment class and simply inject a new time series based pymc model class. Will aim to look into this in more detail - if there was a reason for a new experiment class then we can look into whether we can adapt it to make it more general.

Happy to do it, my only opinion is that the API for the new class is quite different and allow for arbitrary components for trend and seasonality. IF they follow pymc-marketing protocols.

I'm doing that because will help me to easily input HSGP and replace this fourier bases approach (just as example), for the same structural time series. Could probably help to just replace by the state-space.

I'll be working on that as well :) I feel should be easy to do.

[cetagostini]

The new notebook is a great start. Could be good to compare the previous approach "y ~ 1 + t + C(month)" in its_pymc.ipynb to see how things differ.

Almost the same. The bsts fits more during training (high R2) and has less variance (lower sd).

[drbenvincent]

Nice. So I am a beginner in time series modelling - which is why I've not given this a go yet. But that's a good clarification BSTS != state space.

Do you think your implementation leaves the door open for extension into autoregressive components etc?

[cetagostini]

Nice. So I am a beginner in time series modelling - which is why I've not given this a go yet. But that's a good clarification BSTS != state space.

Do you think your implementation leaves the door open for extension into autoregressive components etc?

Indeed, I'll try to bring an example tomorrow!

[drbenvincent]

Firstly, this is very cool! Will probably need a few rounds of review to get this merged, but it will happen :)

Let's try to make it fit into the current docs and codebase as seamlessly as possible. So far we've got a nice distinction between "experiments" and "models".

The new model class BayesianStructuralTimeSeries fits in with the current structure very nicely. Given the range of experiments we've got so far the obvious candidate is to work with the interrupted time series experiment, and maybe multi-observation differences in differences. So the natural way to make this work would be to make the InterruptedTimeSeries class accept your new BayesianStructuralTimeSeries class. That would be in an ideal world, but let's see what we can do quickly to move in that direction.

The API differences between InterruptedTimeSeries and StructuredTimeSeries are minimal. Although obviously what is done with the provided data diverges a lot. Trying to resolve this in one go might be a bit much, but at the moment I've got these concrete suggestions:

• Move the new StructuredTimeSeries class into interrupted_time_series.py and delete structured_time_series.py.
• I relatively strongly think we should avoid the duplication of the plotting code. Like I say, we're trying to make the interrupted time series experiment work with a BSTS model, I don't think there is any reason why any custom plot logic is needed. We could extract the plot code into a mixin class (called something like ITSPlotter) which can be injected into both StructuredTimeSeries and InterruptedTimeSeries.

I think this already goes a long way towards keeping the nice distinction between experiment classes and model classes. The PR is already not far off that - the main issue I see is that StructuredTimeSeries is currently treated like a new experiment when in fact it's exactly the same as interrupted time series, but just separate in order to handle the different model input.

The new notebook is clearly applied to the interrupted time series experiment, so I think the notebook contents can be appended to the existing its_pymc.ipynb. That could be quite nice - the notebook would then be a kind of "how to do ITS from simple linear model to proper time series modeling".

[AlexAndorra]

This looks really cool and useful!!
Just dropping in as I see you guys are thinking of making this broader, with HSGP and broader TS decomposition (which I love!): any reason you're not plugging into pymc-extras.statespace? It should already have everything you need, and you can offload the model dev to it

[drbenvincent]

any reason you're not plugging into pymc-extras.statespace

Thanks @AlexAndorra - this is the direction I was hoping we'd go in. I've done some basic experimentation with it, but I'm a time series newbie and am not confident enough to know when I need to use what component, how to deal with non-stationarity etc. There's a lot more lower hanging fruit for me to work on in CausalPy, so I'm happy that this has been dropped on my lap by @cetagostini :)

I have a very crude understanding of how BSTS and state space approaches relate to each other (though @cetagostini explained it a bit above).

Would it make sense to end up having both this kind of BSTS model class and (eventually) the pymc-extras state space stuff? If that is silly then I guess the proposed implementation could be over-written by a state space implementation? Very happy to take guidance on this.

[drbenvincent]

@cetagostini The BayesianStructuralTimeSeries.build_model method seems like a reasonably simple bit of pymc code. Though as far as I can tell, this seems to be build as an out of the box solution. Nothing wrong with that, but the TFP blog post you linked has a very modular API where users can build their own model from components.

I don't think having a fixed out of the box BSTS model is bad, and in some ways it is good because it is what it is and doesn't require a user to meddle much. That said, it could be pretty cool to think about an API where users can modularly build up a pymc model to pass into the experiment class. Though I imagine this would probably be a considerable bit more work?

[AlexAndorra]

Yeah, the statespace module has a submodule unlocking STS. I think it's doing everything you guys are trying to do here. The one thing we're missing (and actively working on; almost done!) is vectorizing the Kalman filter, to allow for batched time series

[cetagostini]

This looks really cool and useful!!
Just dropping in as I see you guys are thinking of making this broader, with HSGP and broader TS decomposition (which I love!): any reason you're not plugging into pymc-extras.statespace? It should already have everything you need, and you can offload the model dev to it

I'll take a look here! @AlexAndorra Thanks for jumping!

@cetagostini The BayesianStructuralTimeSeries.build_model method seems like a reasonably simple bit of pymc code. Though as far as I can tell, this seems to be build as an out of the box solution. Nothing wrong with that, but the TFP blog post you linked has a very modular API where users can build their own model from components.

I'm already thinking no this, thats why I left the parameters trend_component and seasonal_component with a simple wrapper with a similar pymc-marketing signature, then should be easy to add them and replace by a state-space trend or seasonal.

That said, it could be pretty cool to think about an API where users can modularly build up a pymc model to pass into the experiment class

Thats cool, and seeing signatures, I feel could be. But not sure if for a first iteration? My guess, better to make a class that can manage different stuff like non state-spaces, and state spaces type of models, then we can move forward in order PR and say, lets bring any pymc time series model, and thats it. What do you think @drbenvincent ?

[drbenvincent]

Sure @cetagostini I'm happy with an iterative approach. Happy to proceed along the lines in my first review.

[cetagostini]

@drbenvincent

Apply the changes:

• Under experiements, change Interrupted Time Series name to Structural Time Series, more generic and adhoc to get Interrupted time and basis expansion time series. This will help to hold the same structural time series class able to get StateSpaceTimeSeries, BasisExpansionTimeSeries or InterruptedTimeSeries. All under same family.
• Keep interrupted time series signatures to backward compatibility.
• Move the example under the same ITS notebook.
• Avoid duplication of plotting and stuff because we rely on the same class.

[cetagostini]

@drbenvincent Implementing the StateSpace it's taking more effort than expected, will continue during week... But I think, it requires probably a following PR to not make this massive, instead of all in one.. Not sure, take a look and give me feedback. I already add a class draft in the notebook but need more work to be able to be used by the new experiment STS class.

[cetagostini]

@drbenvincent

integration with state space ready, took more than expected because the state-spaces needed a few wrappers.... Can you modify and give me hand with the docs? All in my local looks okay, not sure whats missing here!

Currently:

• Test are done.
• Backward compatible.
• Support for BSTS (state and non state spaces)

Important

We are using y_hat and mu in different places, I think, y_hat should be the one for the plots, right?

[drbenvincent]

Thanks @cetagostini.

I can see that we've got 2 new examples in the its_pymc.ipynb notebook. That looks cool.

Though it's making me think that we should make the synthetic data a bit more complex - at the moment it's just got a simple linear trend and annual seasonality. We should find some classic time series dataset suitable for interrupted time series. That way, the new stuff you've done should be better than the pre-existing y ~ 1 + t + C(month) linear model, which would better sell the whole thing :)

I'm also seeing the state space model fit doesn't look too great. Maybe we could work on that, possibly getting some input from the state space guru himself?

The changes to the experiment classes isn't quite what I was thinking - though it's close. I was hoping to keep the class name InterruptedTimeSeries because it's most obviously the name of what we're doing in terms of a quasi experiment. Can we move the entire contents of structural_time_series.py and put it in interrupted_time_series.py, and use the InterruptedTimeSeries name for the experiment class rather than StructuralTimeSeries? Happy to jump on a call about it.

In terms of plots, it maybe depends. But for causal impact stuff we want mu. A frequentist approach focusses on the model expectation (mu), so in the Bayesian case we just have a posterior distribution over that expectation.

With the state space stuff, I see you've got some checks for the presence of pymc-extras. Maybe we should add some info in the ITS docs notebook to tell the reader to install it, maybe pointing them to the pymc-extras install instructions?

Though this is looking much closer to done now :)

[cetagostini]

@drbenvincent

Sure, I'll make the changes.

The issue with state space it's observe even in the examples in BSTS from pycm-extras.

You can see that the series shift a bit from the data.

I'll keep the name InterruptedTimeSeries, agree with you. Regarding the data, I'll need to build something, how urgent it is?

Finally, why so many things with conflicts? 😅

[drbenvincent]

I'll take a look at the conflicts when I'm back from a work trip.

Just wanted to keep up the momentum but I know you've got other demands on your time.