pymc-extras 0.4.1__py3-none-any.whl → 0.6.0__py3-none-any.whl

pymc_extras/statespace/models/ETS.py

@@ -9,7 +9,7 @@ from pytensor.compile.mode import Mode
 from pytensor.tensor.slinalg import solve_discrete_lyapunov
 
 from pymc_extras.statespace.core.statespace import PyMCStateSpace, floatX
-from pymc_extras.statespace.models.utilities import make_default_coords
+from pymc_extras.statespace.models.utilities import make_default_coords, validate_names
 from pymc_extras.statespace.utils.constants import (
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
@@ -20,213 +20,209 @@ from pymc_extras.statespace.utils.constants import (
 
 
 class BayesianETS(PyMCStateSpace):
-    def __init__(
-        self,
-        order: tuple[str, str, str] | None = None,
-        endog_names: str | list[str] | None = None,
-        k_endog: int = 1,
-        trend: bool = True,
-        damped_trend: bool = False,
-        seasonal: bool = False,
-        seasonal_periods: int | None = None,
-        measurement_error: bool = False,
-        use_transformed_parameterization: bool = False,
-        dense_innovation_covariance: bool = False,
-        stationary_initialization: bool = False,
-        initialization_dampening: float = 0.8,
-        filter_type: str = "standard",
-        verbose: bool = True,
-        mode: str | Mode | None = None,
-    ):
-        r"""
-        Exponential Smoothing State Space Model
+    r"""
+    Exponential Smoothing State Space Model
+
+    This class can represent a subset of exponential smoothing state space models, specifically those with additive
+    errors. Following .. [1], The general form of the model is:
+
+    .. math::
+
+        \begin{align}
+        y_t &= l_{t-1} + b_{t-1} + s_{t-m} + \epsilon_t \\
+        \epsilon_t &\sim N(0, \sigma)
+        \end{align}
+
+    where :math:`l_t` is the level component, :math:`b_t` is the trend component, and :math:`s_t` is the seasonal
+    component. These components can be included or excluded, leading to different model specifications. The following
+    models are possible:
+
+    * `ETS(A,N,N)`: Simple exponential smoothing
+
+        .. math::
+
+            \begin{align}
+            y_t &= l_{t-1} + \epsilon_t \\
+            l_t &= l_{t-1} + \alpha \epsilon_t
+            \end{align}
+
+    Where :math:`\alpha \in [0, 1]` is a mixing parameter between past observations and current innovations.
+    These equations arise by starting from the "component form":
+
+        .. math::
+
+            \begin{align}
+            \hat{y}_{t+1 | t} &= l_t \\
+            l_t &= \alpha y_t + (1 - \alpha) l_{t-1} \\
+            &= l_{t-1} + \alpha (y_t - l_{t-1})
+            &= l_{t-1} + \alpha \epsilon_t
+            \end{align}
+
+    Where $\epsilon_t$ are the forecast errors, assumed to be IID mean zero and normally distributed. The role of
+    :math:`\alpha` is clearest in the second line. The level of the time series at each time is a mixture of
+    :math:`\alpha` percent of the incoming data, and :math:`1 - \alpha` percent of the previous level. Recursive
+    substitution reveals that the level is a weighted composite of all previous observations; thus the name
+    "Exponential Smoothing".
+
+    Additional supposed specifications include:
+
+    * `ETS(A,A,N)`: Holt's linear trend method
+
+        .. math::
+
+            \begin{align}
+            y_t &= l_{t-1} + b_{t-1} + \epsilon_t \\
+            l_t &= l_{t-1} + b_{t-1} + \alpha \epsilon_t \\
+            b_t &= b_{t-1} + \alpha \beta^\star \epsilon_t
+            \end{align}
+
+        [1]_ also consider an alternative parameterization with :math:`\beta = \alpha \beta^\star`.
+
+    * `ETS(A,N,A)`: Additive seasonal method
+
+        .. math::
 
-        This class can represent a subset of exponential smoothing state space models, specifically those with additive
-        errors. Following .. [1], The general form of the model is:
+            \begin{align}
+            y_t &= l_{t-1} + s_{t-m} + \epsilon_t \\
+            l_t &= l_{t-1} + \alpha \epsilon_t \\
+            s_t &= s_{t-m} + (1 - \alpha)\gamma^\star \epsilon_t
+            \end{align}
+
+        [1]_ also consider an alternative parameterization with :math:`\gamma = (1 - \alpha) \gamma^\star`.
+
+    * `ETS(A,A,A)`: Additive Holt-Winters method
 
         .. math::
 
             \begin{align}
             y_t &= l_{t-1} + b_{t-1} + s_{t-m} + \epsilon_t \\
-            \epsilon_t &\sim N(0, \sigma)
+            l_t &= l_{t-1} + \alpha \epsilon_t \\
+            b_t &= b_{t-1} + \alpha \beta^\star \epsilon_t \\
+            s_t &= s_{t-m} + (1 - \alpha) \gamma^\star \epsilon_t
             \end{align}
 
-        where :math:`l_t` is the level component, :math:`b_t` is the trend component, and :math:`s_t` is the seasonal
-        component. These components can be included or excluded, leading to different model specifications. The following
-        models are possible:
+        [1]_ also consider an alternative parameterization with :math:`\beta = \alpha \beta^star` and
+        :math:`\gamma = (1 - \alpha) \gamma^\star`.
+
+    * `ETS(A, Ad, N)`: Dampened trend method
 
-        * `ETS(A,N,N)`: Simple exponential smoothing
+        .. math::
+
+            \begin{align}
+            y_t &= l_{t-1} + b_{t-1} + \epsilon_t \\
+            l_t &= l_{t-1} + \alpha \epsilon_t \\
+            b_t &= \phi b_{t-1} + \alpha \beta^\star \epsilon_t
+            \end{align}
+
+        [1]_ also consider an alternative parameterization with :math:`\beta = \alpha \beta^\star`.
 
-            .. math::
+    * `ETS(A, Ad, A)`: Dampened trend with seasonal method
 
-                \begin{align}
-                y_t &= l_{t-1} + \epsilon_t \\
-                l_t &= l_{t-1} + \alpha \epsilon_t
-                \end{align}
+        .. math::
 
-        Where :math:`\alpha \in [0, 1]` is a mixing parameter between past observations and current innovations.
-        These equations arise by starting from the "component form":
+            \begin{align}
+            y_t &= l_{t-1} + b_{t-1} + s_{t-m} + \epsilon_t \\
+            l_t &= l_{t-1} + \alpha \epsilon_t \\
+            b_t &= \phi b_{t-1} + \alpha \beta^\star \epsilon_t \\
+            s_t &= s_{t-m} + (1 - \alpha) \gamma^\star \epsilon_t
+            \end{align}
 
-            .. math::
+        [1]_ also consider an alternative parameterization with :math:`\beta = \alpha \beta^star` and
+        :math:`\gamma = (1 - \alpha) \gamma^\star`.
+
+
+    Parameters
+    ----------
+    order: tuple of string, Optional
+        The exponential smoothing "order". This is a tuple of three strings, each of which should be one of 'A', 'Ad',
+        or 'N'.
+        If provided, the model will be initialized from the given order, and the `trend`, `damped_trend`, and `seasonal`
+        arguments will be ignored.
+    endog_names: str or list of str
+        Names associated with observed states. If a list, the length should be equal to the number of time series
+        to be estimated.
+    trend: bool
+        Whether to include a trend component. Setting ``trend=True`` is equivalent to ``order[1] == 'A'``.
+    damped_trend: bool
+        Whether to include a damping parameter on the trend component. Ignored if `trend` is `False`. Setting
+        ``trend=True`` and ``damped_trend=True`` is equivalent to order[1] == 'Ad'.
+    seasonal: bool
+        Whether to include a seasonal component. Setting ``seasonal=True`` is equivalent to ``order[2] = 'A'``.
+    seasonal_periods: int
+        The number of periods in a complete seasonal cycle. Ignored if `seasonal` is `False`
+        (or if ``order[2] == "N"``)
+    measurement_error: bool
+        Whether to include a measurement error term in the model. Default is `False`.
+    use_transformed_parameterization: bool, default False
+        If true, use the :math:`\alpha, \beta, \gamma` parameterization, otherwise use the :math:`\alpha, \beta^\star,
+        \gamma^\star` parameterization. This will change the admissible region for the priors.
+
+        - Under the **non-transformed** parameterization, all of :math:`\alpha, \beta^\star, \gamma^\star` should be
+          between 0 and 1.
+        - Under the **transformed**  parameterization, :math:`\alpha \in (0, 1)`, :math:`\beta \in (0, \alpha)`, and
+          :math:`\gamma \in (0, 1 - \alpha)`
+
+        The :meth:`param_info` method will change to reflect the suggested intervals based on the value of this
+        argument.
+    dense_innovation_covariance: bool, default False
+        Whether to estimate a dense covariance for statespace innovations. In an ETS models, each observed variable
+        has a single source of stochastic variation. If True, these innovations are allowed to be correlated.
+        Ignored if ``k_endog == 1``
+    stationary_initialization: bool, default False
+        If True, the Kalman Filter's initial covariance matrix will be set to an approximate steady-state value.
+        The approximation is formed by adding a small dampening factor to each state. Specifically, the level state
+        for a ('A', 'N', 'N') model is written:
 
-                \begin{align}
-                \hat{y}_{t+1 | t} &= l_t \\
-                l_t &= \alpha y_t + (1 - \alpha) l_{t-1} \\
-                &= l_{t-1} + \alpha (y_t - l_{t-1})
-                &= l_{t-1} + \alpha \epsilon_t
-                \end{align}
+        .. math::
+            \ell_t = \ell_{t-1} + \alpha * e_t
 
-        Where $\epsilon_t$ are the forecast errors, assumed to be IID mean zero and normally distributed. The role of
-        :math:`\alpha` is clearest in the second line. The level of the time series at each time is a mixture of
-        :math:`\alpha` percent of the incoming data, and :math:`1 - \alpha` percent of the previous level. Recursive
-        substitution reveals that the level is a weighted composite of all previous observations; thus the name
-        "Exponential Smoothing".
-
-        Additional supposed specifications include:
+        That this system is not stationary can be understood in ARIMA terms: the level is a random walk; that is,
+        :math:`rho = 1`. This can be remedied by pretending that we instead have a dampened system:
 
-        * `ETS(A,A,N)`: Holt's linear trend method
-
-            .. math::
-
-                \begin{align}
-                y_t &= l_{t-1} + b_{t-1} + \epsilon_t \\
-                l_t &= l_{t-1} + b_{t-1} + \alpha \epsilon_t \\
-                b_t &= b_{t-1} + \alpha \beta^\star \epsilon_t
-                \end{align}
-
-            [1]_ also consider an alternative parameterization with :math:`\beta = \alpha \beta^\star`.
-
-        * `ETS(A,N,A)`: Additive seasonal method
-
-            .. math::
-
-                \begin{align}
-                y_t &= l_{t-1} + s_{t-m} + \epsilon_t \\
-                l_t &= l_{t-1} + \alpha \epsilon_t \\
-                s_t &= s_{t-m} + (1 - \alpha)\gamma^\star \epsilon_t
-                \end{align}
-
-            [1]_ also consider an alternative parameterization with :math:`\gamma = (1 - \alpha) \gamma^\star`.
-
-        * `ETS(A,A,A)`: Additive Holt-Winters method
-
-            .. math::
-
-                \begin{align}
-                y_t &= l_{t-1} + b_{t-1} + s_{t-m} + \epsilon_t \\
-                l_t &= l_{t-1} + \alpha \epsilon_t \\
-                b_t &= b_{t-1} + \alpha \beta^\star \epsilon_t \\
-                s_t &= s_{t-m} + (1 - \alpha) \gamma^\star \epsilon_t
-                \end{align}
-
-            [1]_ also consider an alternative parameterization with :math:`\beta = \alpha \beta^star` and
-            :math:`\gamma = (1 - \alpha) \gamma^\star`.
-
-        * `ETS(A, Ad, N)`: Dampened trend method
-
-            .. math::
-
-                \begin{align}
-                y_t &= l_{t-1} + b_{t-1} + \epsilon_t \\
-                l_t &= l_{t-1} + \alpha \epsilon_t \\
-                b_t &= \phi b_{t-1} + \alpha \beta^\star \epsilon_t
-                \end{align}
-
-            [1]_ also consider an alternative parameterization with :math:`\beta = \alpha \beta^\star`.
-
-        * `ETS(A, Ad, A)`: Dampened trend with seasonal method
-
-            .. math::
-
-                \begin{align}
-                y_t &= l_{t-1} + b_{t-1} + s_{t-m} + \epsilon_t \\
-                l_t &= l_{t-1} + \alpha \epsilon_t \\
-                b_t &= \phi b_{t-1} + \alpha \beta^\star \epsilon_t \\
-                s_t &= s_{t-m} + (1 - \alpha) \gamma^\star \epsilon_t
-                \end{align}
-
-            [1]_ also consider an alternative parameterization with :math:`\beta = \alpha \beta^star` and
-            :math:`\gamma = (1 - \alpha) \gamma^\star`.
-
-
-        Parameters
-        ----------
-        order: tuple of string, Optional
-            The exponential smoothing "order". This is a tuple of three strings, each of which should be one of 'A', 'Ad',
-            or 'N'.
-            If provided, the model will be initialized from the given order, and the `trend`, `damped_trend`, and `seasonal`
-            arguments will be ignored.
-        endog_names: str or list of str, Optional
-            Names associated with observed states. If a list, the length should be equal to the number of time series
-            to be estimated.
-        k_endog: int, Optional
-            Number of time series to estimate. If endog_names are provided, this is ignored and len(endog_names) is
-            used instead.
-        trend: bool
-            Whether to include a trend component. Setting ``trend=True`` is equivalent to ``order[1] == 'A'``.
-        damped_trend: bool
-            Whether to include a damping parameter on the trend component. Ignored if `trend` is `False`. Setting
-            ``trend=True`` and ``damped_trend=True`` is equivalent to order[1] == 'Ad'.
-        seasonal: bool
-            Whether to include a seasonal component. Setting ``seasonal=True`` is equivalent to ``order[2] = 'A'``.
-        seasonal_periods: int
-            The number of periods in a complete seasonal cycle. Ignored if `seasonal` is `False`
-            (or if ``order[2] == "N"``)
-        measurement_error: bool
-            Whether to include a measurement error term in the model. Default is `False`.
-        use_transformed_parameterization: bool, default False
-            If true, use the :math:`\alpha, \beta, \gamma` parameterization, otherwise use the :math:`\alpha, \beta^\star,
-            \gamma^\star` parameterization. This will change the admissible region for the priors.
-
-            - Under the **non-transformed** parameterization, all of :math:`\alpha, \beta^\star, \gamma^\star` should be
-              between 0 and 1.
-            - Under the **transformed**  parameterization, :math:`\alpha \in (0, 1)`, :math:`\beta \in (0, \alpha)`, and
-              :math:`\gamma \in (0, 1 - \alpha)`
-
-            The :meth:`param_info` method will change to reflect the suggested intervals based on the value of this
-            argument.
-        dense_innovation_covariance: bool, default False
-            Whether to estimate a dense covariance for statespace innovations. In an ETS models, each observed variable
-            has a single source of stochastic variation. If True, these innovations are allowed to be correlated.
-            Ignored if ``k_endog == 1``
-        stationary_initialization: bool, default False
-            If True, the Kalman Filter's initial covariance matrix will be set to an approximate steady-state value.
-            The approximation is formed by adding a small dampening factor to each state. Specifically, the level state
-            for a ('A', 'N', 'N') model is written:
-
-            .. math::
-                \ell_t = \ell_{t-1} + \alpha * e_t
-
-            That this system is not stationary can be understood in ARIMA terms: the level is a random walk; that is,
-            :math:`rho = 1`. This can be remedied by pretending that we instead have a dampened system:
-
-            .. math::
-                \ell_t = \rho \ell_{t-1} + \alpha * e_t
-
-            With :math:`\rho \approx 1`, the system is stationary, and we can solve for the steady-state covariance
-            matrix. This is then used as the initial covariance matrix for the Kalman Filter. This is a heuristic
-            method that helps avoid setting a prior on the initial covariance matrix.
-        initialization_dampening: float, default 0.8
-            Dampening factor to add to non-stationary model components. This is only used for initialization, it does
-            *not* add dampening to the model. Ignored if `stationary_initialization` is `False`.
-        filter_type: str, default "standard"
-            The type of Kalman Filter to use. Options are "standard", "single", "univariate", "steady_state",
-            and "cholesky". See the docs for kalman filters for more details.
-        verbose: bool, default True
-            If true, a message will be logged to the terminal explaining the variable names, dimensions, and supports.
-        mode: str or Mode, optional
-            Pytensor compile mode, used in auxiliary sampling methods such as ``sample_conditional_posterior`` and
-            ``forecast``. The mode does **not** effect calls to ``pm.sample``.
-
-            Regardless of whether a mode is specified, it can always be overwritten via the ``compile_kwargs`` argument
-            to all sampling methods.
-
-
-        References
-        ----------
-        .. [1] Hyndman, Rob J., and George Athanasopoulos. Forecasting: principles and practice. OTexts, 2018.
-        """
+        .. math::
+            \ell_t = \rho \ell_{t-1} + \alpha * e_t
+
+        With :math:`\rho \approx 1`, the system is stationary, and we can solve for the steady-state covariance
+        matrix. This is then used as the initial covariance matrix for the Kalman Filter. This is a heuristic
+        method that helps avoid setting a prior on the initial covariance matrix.
+    initialization_dampening: float, default 0.8
+        Dampening factor to add to non-stationary model components. This is only used for initialization, it does
+        *not* add dampening to the model. Ignored if `stationary_initialization` is `False`.
+    filter_type: str, default "standard"
+        The type of Kalman Filter to use. Options are "standard", "single", "univariate", "steady_state",
+        and "cholesky". See the docs for kalman filters for more details.
+    verbose: bool, default True
+        If true, a message will be logged to the terminal explaining the variable names, dimensions, and supports.
+    mode: str or Mode, optional
+        Pytensor compile mode, used in auxiliary sampling methods such as ``sample_conditional_posterior`` and
+        ``forecast``. The mode does **not** effect calls to ``pm.sample``.
+
+        Regardless of whether a mode is specified, it can always be overwritten via the ``compile_kwargs`` argument
+        to all sampling methods.
+
+
+    References
+    ----------
+    .. [1] Hyndman, Rob J., and George Athanasopoulos. Forecasting: principles and practice. OTexts, 2018.
+    """
 
+    def __init__(
+        self,
+        order: tuple[str, str, str] | None = None,
+        endog_names: str | list[str] | None = None,
+        trend: bool = True,
+        damped_trend: bool = False,
+        seasonal: bool = False,
+        seasonal_periods: int | None = None,
+        measurement_error: bool = False,
+        use_transformed_parameterization: bool = False,
+        dense_innovation_covariance: bool = False,
+        stationary_initialization: bool = False,
+        initialization_dampening: float = 0.8,
+        filter_type: str = "standard",
+        verbose: bool = True,
+        mode: str | Mode | None = None,
+    ):
         if order is not None:
             if len(order) != 3 or any(not isinstance(o, str) for o in order):
                 raise ValueError("Order must be a tuple of three strings.")
@@ -265,13 +261,9 @@ class BayesianETS(PyMCStateSpace):
         if self.seasonal and self.seasonal_periods is None:
             raise ValueError("If seasonal is True, seasonal_periods must be provided.")
 
-        if endog_names is not None:
-            endog_names = list(endog_names)
-            k_endog = len(endog_names)
-        else:
-            endog_names = [f"data_{i}" for i in range(k_endog)] if k_endog > 1 else ["data"]
-
-        self.endog_names = endog_names
+        validate_names(endog_names, var_name="endog_names", optional=False)
+        k_endog = len(endog_names)
+        self.endog_names = list(endog_names)
 
         if dense_innovation_covariance and k_endog == 1:
             dense_innovation_covariance = False

pymc_extras/statespace/models/SARIMAX.py

@@ -12,12 +12,13 @@ from pymc_extras.statespace.models.utilities import (
     make_default_coords,
     make_harvey_state_names,
     make_SARIMA_transition_matrix,
+    validate_names,
 )
 from pymc_extras.statespace.utils.constants import (
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
     AR_PARAM_DIM,
-    EXOGENOUS_DIM,
+    EXOG_STATE_DIM,
     MA_PARAM_DIM,
     OBS_STATE_DIM,
     SARIMAX_STATE_STRUCTURES,
@@ -132,7 +133,6 @@ class BayesianSARIMAX(PyMCStateSpace):
         order: tuple[int, int, int],
         seasonal_order: tuple[int, int, int, int] | None = None,
         exog_state_names: list[str] | None = None,
-        k_exog: int | None = None,
         stationary_initialization: bool = True,
         filter_type: str = "standard",
         state_structure: str = "fast",
@@ -166,10 +166,6 @@ class BayesianSARIMAX(PyMCStateSpace):
         exog_state_names : list[str], optional
             Names of the exogenous state variables.
 
-        k_exog : int, optional
-            Number of exogenous variables. If provided, must match the length of
-            `exog_state_names`.
-
         stationary_initialization : bool, default True
             If true, the initial state and initial state covariance will not be assigned priors. Instead, their steady
             state values will be used.
@@ -212,18 +208,10 @@ class BayesianSARIMAX(PyMCStateSpace):
         if seasonal_order is None:
             seasonal_order = (0, 0, 0, 0)
 
-        if exog_state_names is None and k_exog is not None:
-            exog_state_names = [f"exogenous_{i}" for i in range(k_exog)]
-        elif exog_state_names is not None and k_exog is None:
-            k_exog = len(exog_state_names)
-        elif exog_state_names is not None and k_exog is not None:
-            if len(exog_state_names) != k_exog:
-                raise ValueError(
-                    f"Based on provided inputs, expected exog_state_names to have {k_exog} elements, but "
-                    f"found {len(exog_state_names)}"
-                )
-        else:
-            k_exog = 0
+        validate_names(
+            exog_state_names, var_name="exog_state_names", optional=True
+        )  # Not sure if this adds anything
+        k_exog = len(exog_state_names) if exog_state_names is not None else 0
 
         self.exog_state_names = exog_state_names
         self.k_exog = k_exog
@@ -315,7 +303,7 @@ class BayesianSARIMAX(PyMCStateSpace):
     def data_info(self) -> dict[str, dict[str, Any]]:
         info = {
             "exogenous_data": {
-                "dims": (TIME_DIM, EXOGENOUS_DIM),
+                "dims": (TIME_DIM, EXOG_STATE_DIM),
                 "shape": (None, self.k_exog),
             }
         }
@@ -403,7 +391,7 @@ class BayesianSARIMAX(PyMCStateSpace):
             "ma_params": (MA_PARAM_DIM,),
             "seasonal_ar_params": (SEASONAL_AR_PARAM_DIM,),
             "seasonal_ma_params": (SEASONAL_MA_PARAM_DIM,),
-            "beta_exog": (EXOGENOUS_DIM,),
+            "beta_exog": (EXOG_STATE_DIM,),
         }
         if self.k_endog == 1:
             coord_map["sigma_state"] = None
@@ -438,7 +426,7 @@ class BayesianSARIMAX(PyMCStateSpace):
         if self.Q > 0:
             coords.update({SEASONAL_MA_PARAM_DIM: list(range(1, self.Q + 1))})
         if self.k_exog > 0:
-            coords.update({EXOGENOUS_DIM: self.exog_state_names})
+            coords.update({EXOG_STATE_DIM: self.exog_state_names})
         return coords
 
     def _stationary_initialization(self):