Multivariate Structural Statespace Components (#529)

* Reorganize structural model module

* Allow combination of components with different numbers of observed states

* Allow multiple observed in LevelTrend component

* Allow multiple observed states in measurement error component

* Allow multiple observed in AutoRegressive component

* Fix typo in docstrings

* Allow multiple observed in Cycle component

* Fix Cycle docstring examples

* Use pytensor block_diag for Cycle

* 1. updated level_trend component coord/param labels
2. Adjusted the regression component to allow multivariate regression component specification
3. Added a notebook for quick evaluation of the adjustments and additions made

* 1. removed incorrectly comitted file test_structural.py
2. replaced scipy block diag with pytensor block diag
3. Added forecast to test model in multivariate ssm notebook

* removed incorrectly committed file structural.py

* Always count names to determine k_endog

* LevelTrend state/shock names depend on component name

* Update tests to new names

* More test updates

* Delay dropping data names from states/coords until `.build`

* Remove docstring typo

* Update autoregressive component and tests

* 1. updated regression component to revert to univariate shape schema if k_endog=1
2. adjusted coord name in test_statespace test_forecast_with_exog test to avoid key error
3. Added forecast plot and missing value state space section to notebook

* Add component name to shock state names

* Allow multiple observed in TimeSeasonality component

* updated coord dimension names in regression component to always use the user given name as a prefix

* updated level_trend component so that component name is added as a prefix to coord dimension names and updated tests in test_statespace accordingly

* Allow multiple observed in FrequencySeasonality component

* Propagate static shape information in join_tensors_by_dim_labels where possible

* Regression component bugfix and tests

* update default name in test

* Improve cycle code with Jesse's feedback

* Explicitly test matrices in test_cycle

* Add test addition of two Cycles with different observed names

* Make code for state cov when no innov clearer

* 1. updated coords for level trend component and regression component
2. extended test_forecast to include multivariate forecasts with no exogenous variables

* 1. added seasonality params coords fix from Alex Andorra
2. added datetime multivariate forecast tests no exogenous variables
3. added test to check that parameter shapes and coordinate dims agree with one another

* 1. updated param dims in regression component
2. added all components to test_param_dims_coord

* Add exog names to shock_names in LevelTrend

* Update tests to respect new naming convention

* added multivariate forecast with exogenous variables test

* level_trend shock coords should be base_shock_names, not self.shock_names

* Save static shape of last data dim

* More static shapes

* Broken test of decomposition with multiple observed

* Don't use `pad`

* fix decompose test

* updated leveltrend, seasonal, cycle components to adhere to naming schema and updated tests in accordance to naming changes

* Delete notebook

* Use nwe name order in autoregressive component

* Improve docstrings

* Refactor test_regression to cover innotations = True

* Fix comment in test_cycle

* Add docstrings to core and measurement_error

* Improve cycle and seasonal docstrings

* Fix shape of AR params

* Some other AR dims fixes

* Cast all component parameters to lists in `Component.__init__`

* Add type checking and errors for property combination

* re-run structural notebook

---------

Co-authored-by: Alexandre Andorra <[email protected]>
Co-authored-by: Alex Andorra <[email protected]>
Co-authored-by: Jonathan Dekermanjian <[email protected]>

Author: 4 people

notebooks/Structural Timeseries Modeling.ipynb

@@ -0,0 +1,21 @@
+from pymc_extras.statespace.models.structural.components.autoregressive import (
+    AutoregressiveComponent,
+)
+from pymc_extras.statespace.models.structural.components.cycle import CycleComponent
+from pymc_extras.statespace.models.structural.components.level_trend import LevelTrendComponent
+from pymc_extras.statespace.models.structural.components.measurement_error import MeasurementError
+from pymc_extras.statespace.models.structural.components.regression import RegressionComponent
+from pymc_extras.statespace.models.structural.components.seasonality import (
+    FrequencySeasonality,
+    TimeSeasonality,
+)
+
+__all__ = [
+    "LevelTrendComponent",
+    "MeasurementError",
+    "AutoregressiveComponent",
+    "TimeSeasonality",
+    "FrequencySeasonality",
+    "RegressionComponent",
+    "CycleComponent",
+]

pymc_extras/statespace/models/structural.py

@@ -0,0 +1,188 @@
+import numpy as np
+import pytensor.tensor as pt
+
+from pymc_extras.statespace.models.structural.core import Component
+from pymc_extras.statespace.models.structural.utils import order_to_mask
+from pymc_extras.statespace.utils.constants import AR_PARAM_DIM
+
+
+class AutoregressiveComponent(Component):
+    r"""
+    Autoregressive timeseries component
+
+    Parameters
+    ----------
+    order: int or sequence of int
+
+        If int, the number of lags to include in the model.
+        If a sequence, an array-like of zeros and ones indicating which lags to include in the model.
+
+    name: str, default "auto_regressive"
+        A name for this autoregressive component. Used to label dimensions and coordinates.
+
+    observed_state_names: list[str] | None, default None
+        List of strings for observed state labels. If None, defaults to ["data"].
+
+    Notes
+    -----
+    An autoregressive component can be thought of as a way o introducing serially correlated errors into the model.
+    The process is modeled:
+
+    .. math::
+        x_t = \sum_{i=1}^p \rho_i x_{t-i}
+
+    Where ``p``, the number of autoregressive terms to model, is the order of the process. By default, all lags up to
+    ``p`` are included in the model. To disable lags, pass a list of zeros and ones to the ``order`` argumnet. For
+    example, ``order=[1, 1, 0, 1]`` would become:
+
+    .. math::
+        x_t = \rho_1 x_{t-1} + \rho_2 x_{t-1} + \rho_4 x_{t-1}
+
+    The coefficient :math:`\rho_3` has been constrained to zero.
+
+    .. warning:: This class is meant to be used as a component in a structural time series model. For modeling of
+              stationary processes with ARIMA, use ``statespace.BayesianSARIMA``.
+
+    Examples
+    --------
+    Model a timeseries as an AR(2) process with non-zero mean:
+
+    .. code:: python
+
+        from pymc_extras.statespace import structural as st
+        import pymc as pm
+        import pytensor.tensor as pt
+
+        trend = st.LevelTrendComponent(order=1, innovations_order=0)
+        ar = st.AutoregressiveComponent(2)
+        ss_mod = (trend + ar).build()
+
+        with pm.Model(coords=ss_mod.coords) as model:
+            P0 = pm.Deterministic('P0', pt.eye(ss_mod.k_states) * 10, dims=ss_mod.param_dims['P0'])
+            intitial_trend = pm.Normal('initial_trend', sigma=10, dims=ss_mod.param_dims['initial_trend'])
+            ar_params = pm.Normal('ar_params', dims=ss_mod.param_dims['ar_params'])
+            sigma_ar = pm.Exponential('sigma_ar', 1, dims=ss_mod.param_dims['sigma_ar'])
+
+            ss_mod.build_statespace_graph(data)
+            idata = pm.sample(nuts_sampler='numpyro')
+
+    """
+
+    def __init__(
+        self,
+        order: int = 1,
+        name: str = "auto_regressive",
+        observed_state_names: list[str] | None = None,
+    ):
+        if observed_state_names is None:
+            observed_state_names = ["data"]
+
+        k_posdef = k_endog = len(observed_state_names)
+
+        order = order_to_mask(order)
+        ar_lags = np.flatnonzero(order).ravel().astype(int) + 1
+        k_states = len(order)
+
+        self.order = order
+        self.ar_lags = ar_lags
+
+        super().__init__(
+            name=name,
+            k_endog=k_endog,
+            k_states=k_states * k_endog,
+            k_posdef=k_posdef,
+            measurement_error=True,
+            combine_hidden_states=True,
+            observed_state_names=observed_state_names,
+            obs_state_idxs=np.tile(np.r_[[1.0], np.zeros(k_states - 1)], k_endog),
+        )
+
+    def populate_component_properties(self):
+        k_states = self.k_states // self.k_endog  # this is also the number of AR lags
+
+        self.state_names = [
+            f"L{i + 1}[{state_name}]"
+            for state_name in self.observed_state_names
+            for i in range(k_states)
+        ]
+
+        self.shock_names = [f"{self.name}[{obs_name}]" for obs_name in self.observed_state_names]
+        self.param_names = [f"params_{self.name}", f"sigma_{self.name}"]
+        self.param_dims = {f"params_{self.name}": (f"lag_{self.name}",)}
+        self.coords = {f"lag_{self.name}": self.ar_lags.tolist()}
+
+        if self.k_endog > 1:
+            self.param_dims[f"params_{self.name}"] = (
+                f"endog_{self.name}",
+                AR_PARAM_DIM,
+            )
+            self.param_dims[f"sigma_{self.name}"] = (f"endog_{self.name}",)
+
+            self.coords[f"endog_{self.name}"] = self.observed_state_names
+
+        self.param_info = {
+            f"params_{self.name}": {
+                "shape": (k_states,) if self.k_endog == 1 else (self.k_endog, k_states),
+                "constraints": None,
+                "dims": (AR_PARAM_DIM,)
+                if self.k_endog == 1
+                else (
+                    f"endog_{self.name}",
+                    f"lag_{self.name}",
+                ),
+            },
+            f"sigma_{self.name}": {
+                "shape": () if self.k_endog == 1 else (self.k_endog,),
+                "constraints": "Positive",
+                "dims": None if self.k_endog == 1 else (f"endog_{self.name}",),
+            },
+        }
+
+    def make_symbolic_graph(self) -> None:
+        k_endog = self.k_endog
+        k_states = self.k_states // k_endog
+        k_posdef = self.k_posdef
+
+        k_nonzero = int(sum(self.order))
+        ar_params = self.make_and_register_variable(
+            f"params_{self.name}", shape=(k_nonzero,) if k_endog == 1 else (k_endog, k_nonzero)
+        )
+        sigma_ar = self.make_and_register_variable(
+            f"sigma_{self.name}", shape=() if k_endog == 1 else (k_endog,)
+        )
+
+        if k_endog == 1:
+            T = pt.eye(k_states, k=-1)
+            ar_idx = (np.zeros(k_nonzero, dtype="int"), np.nonzero(self.order)[0])
+            T = T[ar_idx].set(ar_params)
+
+        else:
+            transition_matrices = []
+
+            for i in range(k_endog):
+                T = pt.eye(k_states, k=-1)
+                ar_idx = (np.zeros(k_nonzero, dtype="int"), np.nonzero(self.order)[0])
+                T = T[ar_idx].set(ar_params[i])
+                transition_matrices.append(T)
+            T = pt.specify_shape(
+                pt.linalg.block_diag(*transition_matrices), (self.k_states, self.k_states)
+            )
+
+        self.ssm["transition", :, :] = T
+
+        R = np.eye(k_states)
+        R_mask = np.full((k_states), False)
+        R_mask[0] = True
+        R = R[:, R_mask]
+
+        self.ssm["selection", :, :] = pt.specify_shape(
+            pt.linalg.block_diag(*[R for _ in range(k_endog)]), (self.k_states, self.k_posdef)
+        )
+
+        Z = pt.zeros((1, k_states))[0, 0].set(1.0)
+        self.ssm["design", :, :] = pt.specify_shape(
+            pt.linalg.block_diag(*[Z for _ in range(k_endog)]), (self.k_endog, self.k_states)
+        )
+
+        cov_idx = ("state_cov", *np.diag_indices(k_posdef))
+        self.ssm[cov_idx] = sigma_ar**2