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]

No need for a special case here, block_diag([x]) = x and we rewrite away the useless block_diag

[jessegrabowski]

What if innovations=False, does R need to be adjusted in that case?

It was like this before your changes so don't worry about it, but it might need a separate issue.

[jessegrabowski]

Same as below, no need for a special case here

[jessegrabowski]

In this test, eval the transition, design, and selection matrices and make sure they are what they are supposed to be (check the level_trend tests for an example)

[jessegrabowski]

Same here, directly inspect the 3 relevant matrices (plus state_cov I suppose)

[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]

@AlexAndorra Also please add a test adding a cycle component to another cycle component with a different observed state name. Check the resulting matrices come out as expected.

[jessegrabowski]

@Dekermanjian the regression component tests are failing because of this line:

        betas = self.make_and_register_variable(f"beta_{self.name}", shape=(k_endog, k_states))

You need to drop the k_endog part of the shape if k_endog == 1, because we expect a vector in that case (don't want to give all parameters in models a dummy observed state if it doesn't matter). So something like:

        betas = self.make_and_register_variable(f"beta_{self.name}", shape=(k_endog, k_states) if k_endog > 1 else (k_states, )

Also you have beta.reshape((-1, 1)).squeeze(); this is just beta.ravel().

Finally, there are some unused computations -- you had assigned k_posdef = self.k_posdef // self.k_endog, but then didn't use it, and i guess the linter removed the unused variable but left the computation. We can just remove the whole thing.

PS: Please add some multivariate regression tests (I think you already are, but just so it's on the record somewhere, here it is), including a test where you add together two regression components with different state names

[Dekermanjian]

@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

I personally prefer these:

1. sigma_name
2. initial_name
3. greek_name
4. innovation_name

for multivariate, I personally prefer going with option 2 or option 3 if it will also apply to the univariate case. I would prefer that univariate and multivariate are as consistent as would be possible.

[jessegrabowski]

@Dekermanjian I went ahead and fixed up the regression component, so please make sure to pull before you keep working.

The last steps before we merge this are to:

• Address issues with Cycle
• Add a test for forecasting with multiple observed
• Add a test for hidden state decomposition with multiple observed

I think we can also take the rough notebook that @Dekermanjian started and turn it into a tutorial about multiple observed, in the same spirit as the existing one. We can make that a separate PR if we want, but if so, we should drop the notebook from this PR.

@AlexAndorra I need your input on the naming questions I posed above, then we can make a final decision and go with it. Once that's settled, the Cycle fixes are in, and the two tests I'm asking for above are in, I'll rebase this PR so that we have one commit per component (plus one for the utilities) and merge.

[AlexAndorra]

Thanks for the incredibly fast progress @jessegrabowski !! Do you need a new review from me?
I'll address your Cycle comment ASAP.

Here are my votes:

1. For the sigma parameters: sigma_name
2. For the initial state parameters initial_name
3. For assorted greek things, like rho in Cycle: descriptive_name (I feel very strongly against Greek names in general, which are non-descriptive at all and increase the entry cost for users, especially since the state space modeling space uses different Greek names all over the place to designate the same quantities)
4. For shock state names: name_shock

For default names:
2. Keep the complex names for univariate case, but use a simple default name when its multivariate