Multivariate Structural Statespace Components

[jessegrabowski]

This PR lifts the requirement that models built with the structural sub-module of PyMC be univariate. It's a chonky PR, so I split it into commits. Most of the files changes are changed by the first commit, which is just reorganization of files. It is safe to ignore that one.

Here are the steps I followed:

0. The structural module was getting pretty unweildly, so I broke it into a bunch of sub-files. This makes the code easier to find and extend. This is handled in the Reorganize structural model modlue commit
1. We need tools that can merge different components with potentially different (or overlapping) observed time series. This is handled by the Allow combination of component with different numbers of observed states PR. I am confident this code can be improved.
2. Each component needs to have new logic implemented to handle the case where there are multiple observed series. Users can optionally pass a list of names to each component as observed_state_names. Every time you add two components together, all the relevant matrices are padded and expanded, and the total observed states are created as a union between the components.

For now, we assume all states in a component follow the same parameterization. It's now also valid to add together the same component twice with different states to work around this (e.g. AutoRegressive(order=1, observed_state_names=['data_1']) + Autoregressive(order=5, observed_state_names=['data_2'])) would be a valid model with 2 observed states, but each has it's own autoregressive dynamics.

When you pass a batch of observed_state_names, e.g. LevelTrend(order=2, observed_state_names=['data_1', 'data_2']), the parameters will all be given a batch dimension, but will otherwise be the same as the base case.

More docs coming, but I tried obsessively document what in there so far.

The logic for extending the components is pretty straight-forward -- mostly copying + block_diag or concat, but there are some corner cases that need attention.

This PR should be seen as a companion to #450. Instead of vectorizing across the computation of a model, we're concatenating models. There will be cases where this is superior -- for example when you want to explicitly model latent interactions between components. But in other cases, this approach will be worse. I am interested in having both.

[AlexAndorra]

AutoRegressive(order=1, observed_state_names=['data_1']) + Autoregressive(order=5, observed_state_names=['data_2'])) would be a valid model with 2 observed states, but each has it's own autoregressive dynamics.

This is cool! I will review ASAP.

Note that #450 is currently blocked by what I think is a pytensor bug

[AlexAndorra]

This is 🔥 @jessegrabowski 🤯
I just left a suggestion for what I think was a typo in the docstring. I'll merge once this is resolved, and then test all of this for our PyData tutorial -- probably this weekend.

Just a quick question: IIUC, now users can also have batched RegressionComponents, correct?

[AlexAndorra]

This is 🔥 @jessegrabowski 🤯
I just left a suggestion for what I think was a typo in the docstring.

Still missing this feature are:

• Cycle (currently worked on by @AlexAndorra)
• Seasonal
• Regression (currently worked on by @Dekermanjian)

We also need to:

• Make sure that there are tests that combined LevelTrend + AR + error for two observed variables with no interaction model matches two separate models for each, given the same parameters.
• Make sure that pytensor ops are used everywhere for building the SS matrices (no numpy/scipy)

[AlexAndorra]

I think I'm done for a first review from you on the Cycle component @jessegrabowski 🍾

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[jessegrabowski]

@AlexAndorra I left comments for you

Since it's my own PR I can't request changes. It's better in future if you fork the PR branch and open a new PR into this PR, then we can do the usual review workflow on your PR and merge it into this PR when we're ready

[jessegrabowski]

@AlexAndorra @Dekermanjian I want the names of parameters in the components to be really consistent and unsurprising. So please vote on:

1. For the sigma parameters: name_sigma vs sigma_name
2. For the initial state parameters :name_initial vs initial_name vs name
3. For assorted greek things, like rho in Cycle: name_greek vs greek_name vs descriptive_name vs name_descriptive
4. For shock state names: name vs name_shock vs name_innovation

Concrete examples for (3):
a. business_cycle_rho
b. rho_business_cycle
c. dampening_business_cycle
d. business_cycle_dampening

For (4), I'm talking about the internal state names that will end up as labels for the R and Q matrices, nothing else.

Also for default names, since all the are going to depend on the names in the multivariate case, should we:

1. Make all the default names simpler. For example LevelTrend -> level_trend, or Cycle[s={cycle}, dampen={dampen}, innovations={innovations}] -> cycle
2. Keep the complex names for univariate case, but use a simple default name when its multivariate
3. Do away with default names, and force the user to always pass a name
4. As 3, but only in the multivariate case

[jessegrabowski]

Looks like we have consensus on sigma_name and initial_name. Let's agree that in all cases,the name is always at the end of everything and we have a predictable syntax of f'{something}_{name}'.

To that end, let's also go with {name}_shock for the shock dim name. I actually agree with @Dekermanjian that {name}_innovation is a better name, but I can really see myself getting sick of typing "innovation" over and over. Shock is shorter, and it's what we already use.

Regarding greek/descriptive, I'm really on the fence. I understand why @AlexAndorra is super anti-greek, and I've done implementations that specifically remove greek names (changing beta and gamma in BatchNorm to loc and scale here, for example).

On the other hand, we do pay a cost when inventing our own names, unless they are already widely known/used in the specific literature. In the loc and scale example above, ML people probably don't know what that means, so even though it's clearer to us, a practitioner coming in might not agree.

In actual fact, we already mix descriptive and greek names. In CycleComponent, we use cycle_dampening_factor and cycle_length, but then in RegressionComponent we use beta (not coef like in AutoRegressive and TimeSeasonlity :) ). And of course we use sigma everywhere, but I guess that's a "good greek", not one of those bad greeks that take our jobs.

My suggestion is to punt on this and open a new issue to address package-wide naming conventions. For now, let's just make sure everything is of the form f'{something}_{name}'.

[AlexAndorra]

I can really see myself getting sick of typing "innovation" over and over

Ha ha, literally the reasoning behind my choice.

On the other hand, we do pay a cost when inventing our own names, unless they are already widely known/used in the specific literature

Totally agree, it's just that the way I see it, the cost is internal: even if the descriptive names are only our own (which is not the case here, IIUC):

1. We can define them in the docstrings (in the Greek case, one has to go look on the internet and make sure it's a one-to-one correspondence
2. Descriptive is well... more descriptive -- and as Guido says, "explicit is better than implicit"

In CycleComponent, we use cycle_dampening_factor and cycle_length, but then in RegressionComponent we use beta (not coef like in AutoRegressive and TimeSeasonlity :) )

And I like CycleComponent much better for that :)

And of course we use sigma everywhere, but I guess that's a "good greek", not one of those bad greeks that take our jobs.

Masterpiece, nothing to add, you killed me 🤣
In all seriousness though, sigma is so universally associated to standard deviation that this is perfectly fine (and shorter to type)

[jessegrabowski]

This is really close!

I think basically we need to go through and make sure all the names in all the components follow the agreed format. Let's leave {name}_shock for now, but I will open another PR to change that one since it breaks the pattern.

We just need a test for hidden state decomposition with multiple observed, and I think we're pretty much there.

[jessegrabowski]

I added the test for multivariate decomposition. Basically it works, but it leaves a lot to be desired. I will open a separate issue to improve it.

I think this is more or less done. @AlexAndorra and @Dekermanjian , if you could provide a final review then we can merge it.

Also @Dekermanjian what do you want to do with your notebook that got merged in? We can either clean it up a bit to be a simple example, or we can remove it for now and you can do something more elaborate with missingness in the future.

[Dekermanjian]

I think this is more or less done. @AlexAndorra and @Dekermanjian , if you could provide a final review then we can merge it.

That is great! I will go through and review it today as soon as I am off from work.

Also @Dekermanjian what do you want to do with your notebook that got merged in? We can either clean it up a bit to be a simple example, or we can remove it for now and you can do something more elaborate with missingness in the future.

I am not sure. I think that if the simple example doesn’t add anything on top of what Alex is working on then it would be best to remove it and do something more elaborate as you say with missingness. I am interested in figuring out how that will all work in a Bayesian State Space framework.

[Dekermanjian]

I think it is looking great! The main issues I found in my review were related to the doctstring and certain components that still haven't adopted the schema of something_{self.name}. These should be pretty quick to fix. We are almost there!

[AlexAndorra]

I believe... this is it guys 🍾

[Dekermanjian]

This is looking fantastic! 👏