pymc-extras 0.3.1__py3-none-any.whl → 0.4.0__py3-none-any.whl

pymc_extras/statespace/models/structural/components/level_trend.py

@@ -0,0 +1,257 @@
+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 POSITION_DERIVATIVE_NAMES
+
+
+class LevelTrendComponent(Component):
+    r"""
+    Level and trend component of a structural time series model
+
+    Parameters
+    ----------
+    order : int
+
+        Number of time derivatives of the trend to include in the model. For example, when order=3, the trend will
+        be of the form ``y = a + b * t + c * t ** 2``, where the coefficients ``a, b, c`` come from the initial
+        state values.
+
+    innovations_order : int or sequence of int, optional
+
+        The number of stochastic innovations to include in the model. By default, ``innovations_order = order``
+
+    name : str, default "level_trend"
+        A name for this level-trend 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
+    -----
+    This class implements the level and trend components of the general structural time series model. In the most
+    general form, the level and trend is described by a system of two time-varying equations.
+
+    .. math::
+        \begin{align}
+            \mu_{t+1} &= \mu_t + \nu_t + \zeta_t \\
+            \nu_{t+1} &= \nu_t + \xi_t
+            \zeta_t &\sim N(0, \sigma_\zeta) \\
+            \xi_t &\sim N(0, \sigma_\xi)
+        \end{align}
+
+    Where :math:`\mu_{t+1}` is the mean of the timeseries at time t, and :math:`\nu_t` is the drift or the slope of
+    the process. When both innovations :math:`\zeta_t` and :math:`\xi_t` are included in the model, it is known as a
+    *local linear trend* model. This system of two equations, corresponding to ``order=2``, can be expanded or
+    contracted by adding or removing equations. ``order=3`` would add an acceleration term to the sytsem:
+
+    .. math::
+        \begin{align}
+            \mu_{t+1} &= \mu_t + \nu_t + \zeta_t \\
+            \nu_{t+1} &= \nu_t + \eta_t + \xi_t \\
+            \eta_{t+1} &= \eta_{t-1} + \omega_t \\
+            \zeta_t &\sim N(0, \sigma_\zeta) \\
+            \xi_t &\sim N(0, \sigma_\xi) \\
+            \omega_t &\sim N(0, \sigma_\omega)
+        \end{align}
+
+    After setting all innovation terms to zero and defining initial states :math:`\mu_0, \nu_0, \eta_0`, these equations
+    can be collapsed to:
+
+    .. math::
+        \mu_t = \mu_0 + \nu_0 \cdot t + \eta_0 \cdot t^2
+
+    Which clarifies how the order and initial states influence the model. In particular, the initial states are the
+    coefficients on the intercept, slope, acceleration, and so on.
+
+    In this light, allowing for innovations can be understood as allowing these coefficients to vary over time. Each
+    component can be individually selected for time variation by passing a list to the ``innovations_order`` argument.
+    For example, a constant intercept with time varying trend and acceleration is specified as ``order=3,
+    innovations_order=[0, 1, 1]``.
+
+    By choosing the ``order`` and ``innovations_order``, a large variety of models can be obtained. Notable
+    models include:
+
+    * Constant intercept, ``order=1, innovations_order=0``
+
+    .. math::
+        \mu_t = \mu
+
+    * Constant linear slope, ``order=2, innovations_order=0``
+
+    .. math::
+        \mu_t = \mu_{t-1} + \nu
+
+    * Gaussian Random Walk, ``order=1, innovations_order=1``
+
+    .. math::
+        \mu_t = \mu_{t-1} + \zeta_t
+
+    * Gaussian Random Walk with Drift, ``order=2, innovations_order=1``
+
+    .. math::
+        \mu_t = \mu_{t-1} + \nu + \zeta_t
+
+    * Smooth Trend, ``order=2, innovations_order=[0, 1]``
+
+    .. math::
+        \begin{align}
+            \mu_t &= \mu_{t-1} + \nu_{t-1} \\
+            \nu_t &= \nu_{t-1} + \xi_t
+        \end{align}
+
+    * Local Level, ``order=2, innovations_order=2``
+
+    [1] notes that the smooth trend model produces more gradually changing slopes than the full local linear trend
+    model, and is equivalent to an "integrated trend model".
+
+    References
+    ----------
+    .. [1] Durbin, James, and Siem Jan Koopman. 2012.
+        Time Series Analysis by State Space Methods: Second Edition.
+        Oxford University Press.
+
+    """
+
+    def __init__(
+        self,
+        order: int | list[int] = 2,
+        innovations_order: int | list[int] | None = None,
+        name: str = "level_trend",
+        observed_state_names: list[str] | None = None,
+    ):
+        if innovations_order is None:
+            innovations_order = order
+
+        if observed_state_names is None:
+            observed_state_names = ["data"]
+        k_endog = len(observed_state_names)
+
+        self._order_mask = order_to_mask(order)
+        max_state = np.flatnonzero(self._order_mask)[-1].item() + 1
+
+        # If the user passes excess zeros, raise an error. The alternative is to prune them, but this would cause
+        # the shape of the state to be different to what the user expects.
+        if len(self._order_mask) > max_state:
+            raise ValueError(
+                f"order={order} is invalid. The highest derivative should not be set to zero. If you want a "
+                f"lower order model, explicitly omit the zeros."
+            )
+        k_states = max_state
+
+        if isinstance(innovations_order, int):
+            n = innovations_order
+            innovations_order = order_to_mask(k_states)
+            if n > 0:
+                innovations_order[n:] = False
+            else:
+                innovations_order[:] = False
+        else:
+            innovations_order = order_to_mask(innovations_order)
+
+        self.innovations_order = innovations_order[:max_state]
+        k_posdef = int(sum(innovations_order))
+
+        super().__init__(
+            name,
+            k_endog=k_endog,
+            k_states=k_states * k_endog,
+            k_posdef=k_posdef * k_endog,
+            observed_state_names=observed_state_names,
+            measurement_error=False,
+            combine_hidden_states=False,
+            obs_state_idxs=np.tile(np.array([1.0] + [0.0] * (k_states - 1)), k_endog),
+        )
+
+    def populate_component_properties(self):
+        k_endog = self.k_endog
+        k_states = self.k_states // k_endog
+        k_posdef = self.k_posdef // k_endog
+
+        name_slice = POSITION_DERIVATIVE_NAMES[:k_states]
+        self.param_names = [f"initial_{self.name}"]
+        base_names = [name for name, mask in zip(name_slice, self._order_mask) if mask]
+        self.state_names = [
+            f"{name}[{obs_name}]" for obs_name in self.observed_state_names for name in base_names
+        ]
+        self.param_dims = {f"initial_{self.name}": (f"state_{self.name}",)}
+        self.coords = {f"state_{self.name}": base_names}
+
+        if k_endog > 1:
+            self.param_dims[f"state_{self.name}"] = (
+                f"endog_{self.name}",
+                f"state_{self.name}",
+            )
+            self.param_dims = {f"initial_{self.name}": (f"endog_{self.name}", f"state_{self.name}")}
+            self.coords[f"endog_{self.name}"] = self.observed_state_names
+
+        shape = (k_endog, k_states) if k_endog > 1 else (k_states,)
+        self.param_info = {f"initial_{self.name}": {"shape": shape, "constraints": None}}
+
+        if self.k_posdef > 0:
+            self.param_names += [f"sigma_{self.name}"]
+
+            base_shock_names = [
+                name for name, mask in zip(name_slice, self.innovations_order) if mask
+            ]
+
+            self.shock_names = [
+                f"{name}[{obs_name}]"
+                for obs_name in self.observed_state_names
+                for name in base_shock_names
+            ]
+
+            self.param_dims[f"sigma_{self.name}"] = (
+                (f"shock_{self.name}",)
+                if k_endog == 1
+                else (f"endog_{self.name}", f"shock_{self.name}")
+            )
+            self.coords[f"shock_{self.name}"] = base_shock_names
+            self.param_info[f"sigma_{self.name}"] = {
+                "shape": (k_posdef,) if k_endog == 1 else (k_endog, k_posdef),
+                "constraints": "Positive",
+            }
+
+        for name in self.param_names:
+            self.param_info[name]["dims"] = self.param_dims[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_endog
+
+        initial_trend = self.make_and_register_variable(
+            f"initial_{self.name}",
+            shape=(k_states,) if k_endog == 1 else (k_endog, k_states),
+        )
+        self.ssm["initial_state", :] = initial_trend.ravel()
+
+        triu_idx = pt.triu_indices(k_states)
+        T = pt.zeros((k_states, k_states))[triu_idx[0], triu_idx[1]].set(1)
+
+        self.ssm["transition", :, :] = pt.specify_shape(
+            pt.linalg.block_diag(*[T for _ in range(k_endog)]), (self.k_states, self.k_states)
+        )
+
+        R = np.eye(k_states)
+        R = R[:, self.innovations_order]
+
+        self.ssm["selection", :, :] = pt.specify_shape(
+            pt.linalg.block_diag(*[R for _ in range(k_endog)]), (self.k_states, self.k_posdef)
+        )
+
+        Z = np.array([1.0] + [0.0] * (k_states - 1)).reshape((1, -1))
+
+        self.ssm["design", :, :] = pt.specify_shape(
+            pt.linalg.block_diag(*[Z for _ in range(k_endog)]), (self.k_endog, self.k_states)
+        )
+
+        if k_posdef > 0:
+            sigma_trend = self.make_and_register_variable(
+                f"sigma_{self.name}",
+                shape=(k_posdef,) if k_endog == 1 else (k_endog, k_posdef),
+            )
+            diag_idx = np.diag_indices(k_posdef * k_endog)
+            idx = np.s_["state_cov", diag_idx[0], diag_idx[1]]
+            self.ssm[idx] = (sigma_trend**2).ravel()

pymc_extras/statespace/models/structural/components/measurement_error.py

@@ -0,0 +1,137 @@
+import numpy as np
+
+from pymc_extras.statespace.models.structural.core import Component
+
+
+class MeasurementError(Component):
+    r"""
+    Measurement error component for structural time series models.
+
+    This component adds observation noise to the model by introducing a variance parameter
+    that affects the observation covariance matrix H. Unlike other components, it has no
+    hidden states and should only be used in combination with other components.
+
+    Parameters
+    ----------
+    name : str, optional
+        Name of the measurement error component. Default is "MeasurementError".
+    observed_state_names : list[str] | None, optional
+        Names of the observed variables. If None, defaults to ["data"].
+
+    Notes
+    -----
+    The measurement error component models observation noise as:
+
+    .. math::
+
+        y_t = \text{signal}_t + \varepsilon_t, \quad \varepsilon_t \sim N(0, \sigma^2)
+
+    Where :math:`\text{signal}_t` is the true signal from other components and
+    :math:`\sigma^2` is the measurement error variance.
+
+    This component:
+        - Has no hidden states (k_states = 0)
+        - Has no innovations (k_posdef = 0)
+        - Adds a single parameter: sigma_{name}
+        - Modifies the observation covariance matrix H
+
+    Examples
+    --------
+    **Basic usage with trend component:**
+
+    .. code:: python
+
+        from pymc_extras.statespace import structural as st
+        import pymc as pm
+        import pytensor.tensor as pt
+
+        trend = st.LevelTrendComponent(order=2, innovations_order=1)
+        error = st.MeasurementError()
+
+        ss_mod = (trend + error).build()
+
+        # Use with PyMC
+        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'])
+            initial_trend = pm.Normal('initial_trend', sigma=10, dims=ss_mod.param_dims['initial_trend'])
+            sigma_obs = pm.Exponential('sigma_obs', 1, dims=ss_mod.param_dims['sigma_obs'])
+
+            ss_mod.build_statespace_graph(data)
+            idata = pm.sample()
+
+    **Multivariate measurement error:**
+
+    .. code:: python
+
+        # For multiple observed variables
+        # This creates separate measurement error variances for each variable
+        # sigma_obs_error will have shape (3,) for the three variables
+        error = st.MeasurementError(
+            name="obs_error",
+            observed_state_names=["gdp", "unemployment", "inflation"]
+        )
+
+    **Complete model example:**
+
+    .. code:: python
+
+        trend = st.LevelTrendComponent(order=2, innovations_order=1)
+        seasonal = st.TimeSeasonality(season_length=12, innovations=True)
+        error = st.MeasurementError()
+
+        model = (trend + seasonal + error).build()
+
+        # The model now includes:
+        # - Trend parameters: level_trend, sigma_trend
+        # - Seasonal parameters: seasonal_coefs, sigma_seasonal
+        # - Measurement error parameter: sigma_obs
+
+    See Also
+    --------
+    Component : Base class for all structural components.
+    StructuralTimeSeries : Complete model class.
+    """
+
+    def __init__(
+        self, name: str = "MeasurementError", observed_state_names: list[str] | None = None
+    ):
+        if observed_state_names is None:
+            observed_state_names = ["data"]
+
+        k_endog = len(observed_state_names)
+        k_states = 0
+        k_posdef = 0
+
+        super().__init__(
+            name,
+            k_endog,
+            k_states,
+            k_posdef,
+            measurement_error=True,
+            combine_hidden_states=False,
+            observed_state_names=observed_state_names,
+        )
+
+    def populate_component_properties(self):
+        self.param_names = [f"sigma_{self.name}"]
+        self.param_dims = {}
+        self.coords = {}
+
+        if self.k_endog > 1:
+            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"sigma_{self.name}": {
+                "shape": (self.k_endog,) if self.k_endog > 1 else (),
+                "constraints": "Positive",
+                "dims": (f"endog_{self.name}",) if self.k_endog > 1 else None,
+            }
+        }
+
+    def make_symbolic_graph(self) -> None:
+        sigma_shape = () if self.k_endog == 1 else (self.k_endog,)
+        error_sigma = self.make_and_register_variable(f"sigma_{self.name}", shape=sigma_shape)
+        diag_idx = np.diag_indices(self.k_endog)
+        idx = np.s_["obs_cov", diag_idx[0], diag_idx[1]]
+        self.ssm[idx] = error_sigma**2

pymc_extras/statespace/models/structural/components/regression.py

@@ -0,0 +1,228 @@
+import numpy as np
+
+from pytensor import tensor as pt
+
+from pymc_extras.statespace.models.structural.core import Component
+from pymc_extras.statespace.utils.constants import TIME_DIM
+
+
+class RegressionComponent(Component):
+    r"""
+    Regression component for exogenous variables in a structural time series model
+
+    Parameters
+    ----------
+    k_exog : int | None, default None
+        Number of exogenous variables to include in the regression. Must be specified if
+        state_names is not provided.
+
+    name : str | None, default "regression"
+        A name for this regression component. Used to label dimensions and coordinates.
+
+    state_names : list[str] | None, default None
+        List of strings for regression coefficient labels. If provided, must be of length
+        k_exog. If None and k_exog is provided, coefficients will be named
+        "{name}_1, {name}_2, ...".
+
+    observed_state_names : list[str] | None, default None
+        List of strings for observed state labels. If None, defaults to ["data"].
+
+    innovations : bool, default False
+        Whether to include stochastic innovations in the regression coefficients,
+        allowing them to vary over time. If True, coefficients follow a random walk.
+
+    Notes
+    -----
+    This component implements regression with exogenous variables in a structural time series
+    model. The regression component can be expressed as:
+
+    .. math::
+        y_t = \beta_t^T x_t + \epsilon_t
+
+    Where :math:`y_t` is the dependent variable, :math:`x_t` is the vector of exogenous
+    variables, :math:`\beta_t` is the vector of regression coefficients, and :math:`\epsilon_t`
+    is the error term.
+
+    When ``innovations=False`` (default), the coefficients are constant over time:
+    :math:`\beta_t = \beta_0` for all t.
+
+    When ``innovations=True``, the coefficients follow a random walk:
+    :math:`\beta_{t+1} = \beta_t + \eta_t`, where :math:`\eta_t \sim N(0, \Sigma_\beta)`.
+
+    The component supports both univariate and multivariate regression. In the multivariate
+    case, separate coefficients are estimated for each endogenous variable (i.e time series).
+
+    Examples
+    --------
+    Simple regression with constant coefficients:
+
+    .. 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=1)
+        regression = st.RegressionComponent(k_exog=2, state_names=['intercept', 'slope'])
+        ss_mod = (trend + regression).build()
+
+        with pm.Model(coords=ss_mod.coords) as model:
+            # Prior for regression coefficients
+            betas = pm.Normal('betas', dims=ss_mod.param_dims['beta_regression'])
+
+            # Prior for trend innovations
+            sigma_trend = pm.Exponential('sigma_trend', 1)
+
+            ss_mod.build_statespace_graph(data)
+            idata = pm.sample()
+
+    Multivariate regression with time-varying coefficients:
+    - There are 2 exogenous variables (price and income effects)
+    - There are 2 endogenous variables (sales and revenue)
+    - The regression coefficients are allowed to vary over time (`innovations=True`)
+
+    .. code:: python
+
+        regression = st.RegressionComponent(
+            k_exog=2,
+            state_names=['price_effect', 'income_effect'],
+            observed_state_names=['sales', 'revenue'],
+            innovations=True
+        )
+
+        with pm.Model(coords=ss_mod.coords) as model:
+            betas = pm.Normal('betas', dims=ss_mod.param_dims['beta_regression'])
+
+            # Innovation variance for time-varying coefficients
+            sigma_beta = pm.Exponential('sigma_beta', 1, dims=ss_mod.param_dims['sigma_beta_regression'])
+
+            ss_mod.build_statespace_graph(data)
+            idata = pm.sample()
+    """
+
+    def __init__(
+        self,
+        k_exog: int | None = None,
+        name: str | None = "regression",
+        state_names: list[str] | None = None,
+        observed_state_names: list[str] | None = None,
+        innovations=False,
+    ):
+        if observed_state_names is None:
+            observed_state_names = ["data"]
+
+        self.innovations = innovations
+        k_exog = self._handle_input_data(k_exog, state_names, name)
+
+        k_states = k_exog
+        k_endog = len(observed_state_names)
+        k_posdef = k_exog
+
+        super().__init__(
+            name=name,
+            k_endog=k_endog,
+            k_states=k_states * k_endog,
+            k_posdef=k_posdef * k_endog,
+            state_names=self.state_names,
+            observed_state_names=observed_state_names,
+            measurement_error=False,
+            combine_hidden_states=False,
+            exog_names=[f"data_{name}"],
+            obs_state_idxs=np.ones(k_states),
+        )
+
+    @staticmethod
+    def _get_state_names(k_exog: int | None, state_names: list[str] | None, name: str):
+        if k_exog is None and state_names is None:
+            raise ValueError("Must specify at least one of k_exog or state_names")
+        if state_names is not None and k_exog is not None:
+            if len(state_names) != k_exog:
+                raise ValueError(f"Expected {k_exog} state names, found {len(state_names)}")
+        elif k_exog is None:
+            k_exog = len(state_names)
+        else:
+            state_names = [f"{name}_{i + 1}" for i in range(k_exog)]
+
+        return k_exog, state_names
+
+    def _handle_input_data(self, k_exog: int, state_names: list[str] | None, name) -> int:
+        k_exog, state_names = self._get_state_names(k_exog, state_names, name)
+        self.state_names = state_names
+
+        return k_exog
+
+    def make_symbolic_graph(self) -> None:
+        k_endog = self.k_endog
+        k_states = self.k_states // k_endog
+
+        betas = self.make_and_register_variable(
+            f"beta_{self.name}", shape=(k_endog, k_states) if k_endog > 1 else (k_states,)
+        )
+        regression_data = self.make_and_register_data(f"data_{self.name}", shape=(None, k_states))
+
+        self.ssm["initial_state", :] = betas.ravel()
+        self.ssm["transition", :, :] = pt.eye(self.k_states)
+        self.ssm["selection", :, :] = pt.eye(self.k_states)
+
+        Z = pt.linalg.block_diag(*[pt.expand_dims(regression_data, 1) for _ in range(k_endog)])
+        self.ssm["design"] = pt.specify_shape(
+            Z, (None, k_endog, regression_data.type.shape[1] * k_endog)
+        )
+
+        if self.innovations:
+            sigma_beta = self.make_and_register_variable(
+                f"sigma_beta_{self.name}", (k_states,) if k_endog == 1 else (k_endog, k_states)
+            )
+            row_idx, col_idx = np.diag_indices(self.k_states)
+            self.ssm["state_cov", row_idx, col_idx] = sigma_beta.ravel() ** 2
+
+    def populate_component_properties(self) -> None:
+        k_endog = self.k_endog
+        k_states = self.k_states // k_endog
+
+        self.shock_names = self.state_names
+
+        self.param_names = [f"beta_{self.name}"]
+        self.data_names = [f"data_{self.name}"]
+        self.param_dims = {
+            f"beta_{self.name}": (f"endog_{self.name}", f"state_{self.name}")
+            if k_endog > 1
+            else (f"state_{self.name}",)
+        }
+
+        base_names = self.state_names
+        self.state_names = [
+            f"{name}[{obs_name}]" for obs_name in self.observed_state_names for name in base_names
+        ]
+
+        self.param_info = {
+            f"beta_{self.name}": {
+                "shape": (k_endog, k_states) if k_endog > 1 else (k_states,),
+                "constraints": None,
+                "dims": (f"endog_{self.name}", f"state_{self.name}")
+                if k_endog > 1
+                else (f"state_{self.name}",),
+            },
+        }
+
+        self.data_info = {
+            f"data_{self.name}": {
+                "shape": (None, k_states),
+                "dims": (TIME_DIM, f"state_{self.name}"),
+            },
+        }
+        self.coords = {
+            f"state_{self.name}": base_names,
+            f"endog_{self.name}": self.observed_state_names,
+        }
+
+        if self.innovations:
+            self.param_names += [f"sigma_beta_{self.name}"]
+            self.param_dims[f"sigma_beta_{self.name}"] = (f"state_{self.name}",)
+            self.param_info[f"sigma_beta_{self.name}"] = {
+                "shape": (k_states,),
+                "constraints": "Positive",
+                "dims": (f"state_{self.name}",)
+                if k_endog == 1
+                else (f"endog_{self.name}", f"state_{self.name}"),
+            }