pymc-extras 0.5.0__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,6 +12,7 @@ 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,
@@ -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

pymc_extras/statespace/models/VARMAX.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
-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,
@@ -99,9 +99,7 @@ class BayesianVARMAX(PyMCStateSpace):
         self,
         order: tuple[int, int],
         endog_names: list[str] | None = None,
-        k_endog: int | None = None,
         exog_state_names: list[str] | dict[str, list[str]] | None = None,
-        k_exog: int | dict[str, int] | None = None,
         stationary_initialization: bool = False,
         filter_type: str = "standard",
         measurement_error: bool = False,
@@ -118,10 +116,7 @@ class BayesianVARMAX(PyMCStateSpace):
             specified order are included. For restricted models, set zeros directly on the priors.
 
         endog_names: list of str, optional
-            Names of the endogenous variables being modeled. Used to generate names for the state and shock coords. If
-            None, the state names will simply be numbered.
-
-            Exactly one of either ``endog_names`` or ``k_endog`` must be specified.
+            Names of the endogenous variables being modeled. Used to generate names for the state and shock coords.
 
         exog_state_names : list[str] or dict[str, list[str]], optional
             Names of the exogenous state variables. If a list, all endogenous variables will share the same exogenous
@@ -129,12 +124,6 @@ class BayesianVARMAX(PyMCStateSpace):
             exogenous variable names for that endogenous variable. Endogenous variables not included in the dict will
             be assumed to have no exogenous variables. If None, no exogenous variables will be included.
 
-        k_exog : int or dict[str, int], optional
-            Number of exogenous variables. If an int, all endogenous variables will share the same number of exogenous
-            variables. If a dict, keys should be the names of the endogenous variables, and values should be the number of
-            exogenous variables for that endogenous variable. Endogenous variables not included in the dict will be
-            assumed to have no exogenous variables. If None, no exogenous variables will be included.
-
         stationary_initialization: bool, default False
             If true, the initial state and initial state covariance will not be assigned priors. Instead, their steady
             state values will be used. If False, the user is responsible for setting priors on the initial state and
@@ -148,15 +137,6 @@ class BayesianVARMAX(PyMCStateSpace):
             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.
 
-        state_structure: str, default "fast"
-            How to represent the state-space system. When "interpretable", each element of the state vector will have a
-            precise meaning as either lagged data, innovations, or lagged innovations. This comes at the cost of a larger
-            state vector, which may hurt performance.
-
-            When "fast", states are combined to minimize the dimension of the state vector, but lags and innovations are
-            mixed together as a result. Only the first state (the modeled timeseries) will have an obvious interpretation
-            in this case.
-
         measurement_error: bool, default True
             If true, a measurement error term is added to the model.
 
@@ -171,57 +151,23 @@ class BayesianVARMAX(PyMCStateSpace):
             to all sampling methods.
 
         """
-        if (endog_names is None) and (k_endog is None):
-            raise ValueError("Must specify either endog_names or k_endog")
-        if (endog_names is not None) and (k_endog is None):
-            k_endog = len(endog_names)
-        if (endog_names is None) and (k_endog is not None):
-            endog_names = [f"observed_{i}" for i in range(k_endog)]
-        if (endog_names is not None) and (k_endog is not None):
-            if len(endog_names) != k_endog:
-                raise ValueError("Length of provided endog_names does not match provided k_endog")
-
-        if k_exog is not None and not isinstance(k_exog, int | dict):
-            raise ValueError("If not None, k_endog must be either an int or a dict")
+
+        validate_names(endog_names, var_name="endog_names", optional=False)
+        k_endog = len(endog_names)
+
+        needs_exog_data = False
+
         if exog_state_names is not None and not isinstance(exog_state_names, list | dict):
             raise ValueError("If not None, exog_state_names must be either a list or a dict")
 
-        if k_exog is not None and exog_state_names is not None:
-            if isinstance(k_exog, int) and isinstance(exog_state_names, list):
-                if len(exog_state_names) != k_exog:
-                    raise ValueError("Length of exog_state_names does not match provided k_exog")
-            elif isinstance(k_exog, int) and isinstance(exog_state_names, dict):
-                raise ValueError(
-                    "If k_exog is an int, exog_state_names must be a list of the same length (or None)"
-                )
-            elif isinstance(k_exog, dict) and isinstance(exog_state_names, list):
-                raise ValueError(
-                    "If k_exog is a dict, exog_state_names must be a dict as well (or None)"
-                )
-            elif isinstance(k_exog, dict) and isinstance(exog_state_names, dict):
-                if set(k_exog.keys()) != set(exog_state_names.keys()):
-                    raise ValueError("Keys of k_exog and exog_state_names dicts must match")
-                if not all(
-                    len(names) == k for names, k in zip(exog_state_names.values(), k_exog.values())
-                ):
-                    raise ValueError(
-                        "If both k_endog and exog_state_names are provided, lengths of exog_state_names "
-                        "lists must match corresponding values in k_exog"
-                    )
-
-        if k_exog is not None and exog_state_names is None:
-            if isinstance(k_exog, int):
-                exog_state_names = [f"exogenous_{i}" for i in range(k_exog)]
-            elif isinstance(k_exog, dict):
-                exog_state_names = {
-                    name: [f"{name}_exogenous_{i}" for i in range(k)] for name, k in k_exog.items()
-                }
-
-        if k_exog is None and exog_state_names is not None:
+        if exog_state_names is not None:
             if isinstance(exog_state_names, list):
                 k_exog = len(exog_state_names)
             elif isinstance(exog_state_names, dict):
                 k_exog = {name: len(names) for name, names in exog_state_names.items()}
+            needs_exog_data = True
+        else:
+            k_exog = None
 
         # If exog_state_names is a dict but 1) all endog variables are among the keys, and 2) all values are the same
         # then we can drop back to the list case.
@@ -254,6 +200,8 @@ class BayesianVARMAX(PyMCStateSpace):
             mode=mode,
         )
 
+        self._needs_exog_data = needs_exog_data
+
         # Save counts of the number of parameters in each category
         self.param_counts = {
             "x0": k_states * (1 - self.stationary_initialization),
@@ -337,7 +285,7 @@ class BayesianVARMAX(PyMCStateSpace):
 
     @property
     def data_info(self) -> dict[str, dict[str, Any]]:
-        info = None
+        info = {}
 
         if isinstance(self.exog_state_names, list):
             info = {

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

@@ -141,19 +141,19 @@ class AutoregressiveComponent(Component):
 
         self.param_info = {
             f"params_{self.name}": {
-                "shape": (k_states,) if self.k_endog == 1 else (self.k_endog, k_states),
+                "shape": (k_endog_effective, k_states) if k_endog_effective > 1 else (k_states,),
                 "constraints": None,
                 "dims": (AR_PARAM_DIM,)
-                if self.k_endog == 1
+                if k_endog_effective == 1
                 else (
                     f"endog_{self.name}",
                     f"lag_{self.name}",
                 ),
             },
             f"sigma_{self.name}": {
-                "shape": () if self.k_endog == 1 else (self.k_endog,),
+                "shape": (k_endog_effective,) if k_endog_effective > 1 else (),
                 "constraints": "Positive",
-                "dims": None if self.k_endog == 1 else (f"endog_{self.name}",),
+                "dims": (f"endog_{self.name}",) if k_endog_effective > 1 else None,
             },
         }