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

pymc_extras/inference/laplace_approx/find_map.py

@@ -326,7 +326,7 @@ def find_MAP(
     )
 
     raveled_optimized = RaveledVars(optimizer_result.x, initial_params.point_map_info)
-    unobserved_vars = get_default_varnames(model.unobserved_value_vars, include_transformed)
+    unobserved_vars = get_default_varnames(model.unobserved_value_vars, include_transformed=True)
     unobserved_vars_values = model.compile_fn(unobserved_vars, mode="FAST_COMPILE")(
         DictToArrayBijection.rmap(raveled_optimized)
     )
@@ -335,13 +335,20 @@ def find_MAP(
         var.name: value for var, value in zip(unobserved_vars, unobserved_vars_values)
     }
 
-    idata = map_results_to_inference_data(optimized_point, frozen_model)
-    idata = add_fit_to_inference_data(idata, raveled_optimized, H_inv)
+    idata = map_results_to_inference_data(
+        map_point=optimized_point, model=frozen_model, include_transformed=include_transformed
+    )
+
+    idata = add_fit_to_inference_data(
+        idata=idata, mu=raveled_optimized, H_inv=H_inv, model=frozen_model
+    )
+
     idata = add_optimizer_result_to_inference_data(
-        idata, optimizer_result, method, raveled_optimized, model
+        idata=idata, result=optimizer_result, method=method, mu=raveled_optimized, model=model
     )
+
     idata = add_data_to_inference_data(
-        idata, progressbar=False, model=model, compile_kwargs=compile_kwargs
+        idata=idata, progressbar=False, model=model, compile_kwargs=compile_kwargs
     )
 
     return idata

pymc_extras/inference/laplace_approx/idata.py

@@ -59,6 +59,7 @@ def make_unpacked_variable_names(names: list[str], model: pm.Model) -> list[str]
 def map_results_to_inference_data(
     map_point: dict[str, float | int | np.ndarray],
     model: pm.Model | None = None,
+    include_transformed: bool = True,
 ):
     """
     Add the MAP point to an InferenceData object in the posterior group.
@@ -68,13 +69,13 @@ def map_results_to_inference_data(
 
     Parameters
     ----------
-    idata: az.InferenceData
-        An InferenceData object to which the MAP point will be added.
     map_point: dict
         A dictionary containing the MAP point estimates for each variable. The keys should be the variable names, and
         the values should be the corresponding MAP estimates.
     model: Model, optional
         A PyMC model. If None, the model is taken from the current model context.
+    include_transformed: bool
+        Whether to return transformed (unconstrained) variables in the constrained_posterior group. Default is True.
 
     Returns
     -------
@@ -118,7 +119,7 @@ def map_results_to_inference_data(
         dims=dims,
     )
 
-    if unconstrained_names:
+    if unconstrained_names and include_transformed:
         unconstrained_posterior = az.from_dict(
             posterior={
                 k: np.expand_dims(v, (0, 1))

pymc_extras/inference/laplace_approx/laplace.py

@@ -302,7 +302,7 @@ def fit_laplace(
     ----------
     model : pm.Model
         The PyMC model to be fit. If None, the current model context is used.
-    method : str
+    optimize_method : str
         The optimization method to use. Valid choices are: Nelder-Mead, Powell, CG, BFGS, L-BFGS-B, TNC, SLSQP,
         trust-constr, dogleg, trust-ncg, trust-exact, trust-krylov, and basinhopping.
 
@@ -441,9 +441,11 @@ def fit_laplace(
             .rename({"temp_chain": "chain", "temp_draw": "draw"})
         )
 
-        idata.unconstrained_posterior = unstack_laplace_draws(
-            new_posterior.laplace_approximation.values, model, chains=chains, draws=draws
-        )
+        if include_transformed:
+            idata.unconstrained_posterior = unstack_laplace_draws(
+                new_posterior.laplace_approximation.values, model, chains=chains, draws=draws
+            )
+
         idata.posterior = new_posterior.drop_vars(
             ["laplace_approximation", "unpacked_variable_names"]
         )

pymc_extras/inference/pathfinder/pathfinder.py

@@ -38,16 +38,15 @@ from pymc.blocking import DictToArrayBijection, RaveledVars
 from pymc.initial_point import make_initial_point_fn
 from pymc.model import modelcontext
 from pymc.model.core import Point
+from pymc.progress_bar import CustomProgress, default_progress_theme
 from pymc.pytensorf import (
     compile,
     find_rng_nodes,
     reseed_rngs,
 )
 from pymc.util import (
-    CustomProgress,
     RandomSeed,
     _get_seeds_per_chain,
-    default_progress_theme,
     get_default_varnames,
 )
 from pytensor.compile.function.types import Function

pymc_extras/statespace/models/structural/__init__.py

@@ -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/components/autoregressive.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

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

@@ -0,0 +1,305 @@
+import numpy as np
+
+from pytensor import tensor as pt
+from pytensor.tensor.slinalg import block_diag
+
+from pymc_extras.statespace.models.structural.core import Component
+from pymc_extras.statespace.models.structural.utils import _frequency_transition_block
+
+
+class CycleComponent(Component):
+    r"""
+    A component for modeling longer-term cyclical effects
+
+    Supports both univariate and multivariate time series. For multivariate time series,
+    each endogenous variable gets its own independent cycle component with separate
+    cosine/sine states and optional variable-specific innovation variances.
+
+    Parameters
+    ----------
+    name: str
+        Name of the component. Used in generated coordinates and state names. If None, a descriptive name will be
+        used.
+
+    cycle_length: int, optional
+        The length of the cycle, in the calendar units of your data. For example, if your data is monthly, and you
+        want to model a 12-month cycle, use ``cycle_length=12``. You cannot specify both ``cycle_length`` and
+        ``estimate_cycle_length``.
+
+    estimate_cycle_length: bool, default False
+        Whether to estimate the cycle length. If True, an additional parameter, ``cycle_length`` will be added to the
+        model. You cannot specify both ``cycle_length`` and ``estimate_cycle_length``.
+
+    dampen: bool, default False
+        Whether to dampen the cycle by multiplying by a dampening factor :math:`\rho` at every timestep. If true,
+        an additional parameter, ``dampening_factor`` will be added to the model.
+
+    innovations: bool, default True
+        Whether to include stochastic innovations in the strength of the seasonal effect. If True, an additional
+        parameter, ``sigma_{name}`` will be added to the model.
+        For multivariate time series, this is a vector (variable-specific innovation variances).
+
+    observed_state_names: list[str], optional
+        Names of the observed state variables. For univariate time series, defaults to ``["data"]``.
+        For multivariate time series, specify a list of names for each endogenous variable.
+
+    Notes
+    -----
+    The cycle component is very similar in implementation to the frequency domain seasonal component, expect that it
+    is restricted to n=1. The cycle component can be expressed:
+
+    .. math::
+        \begin{align}
+            \gamma_t &= \rho \gamma_{t-1} \cos \lambda + \rho \gamma_{t-1}^\star \sin \lambda + \omega_{t} \\
+            \gamma_{t}^\star &= -\rho \gamma_{t-1} \sin \lambda + \rho \gamma_{t-1}^\star \cos \lambda + \omega_{t}^\star \\
+            \lambda &= \frac{2\pi}{s}
+        \end{align}
+
+    Where :math:`s` is the ``cycle_length``. [1] recommend that this component be used for longer term cyclical
+    effects, such as business cycles, and that the seasonal component be used for shorter term effects, such as
+    weekly or monthly seasonality.
+
+    Unlike a FrequencySeasonality component, the length of a CycleComponent can be estimated.
+
+    **Multivariate Support:**
+    For multivariate time series with k endogenous variables, the component creates:
+    - 2k states (cosine and sine components for each variable)
+    - Block diagonal transition and selection matrices
+    - Variable-specific innovation variances (optional)
+    - Proper parameter shapes: (k, 2) for initial states, (k,) for innovation variances
+
+    Examples
+    --------
+    **Univariate Example:**
+    Estimate a business cycle with length between 6 and 12 years:
+
+    .. code:: python
+
+        from pymc_extras.statespace import structural as st
+        import pymc as pm
+        import pytensor.tensor as pt
+        import pandas as pd
+        import numpy as np
+
+        data = np.random.normal(size=(100, 1))
+
+        # Build the structural model
+        grw = st.LevelTrendComponent(order=1, innovations_order=1)
+        cycle = st.CycleComponent(
+            "business_cycle", cycle_length=12, estimate_cycle_length=False, innovations=True, dampen=True
+        )
+        ss_mod = (grw + cycle).build()
+
+        # Estimate with PyMC
+        with pm.Model(coords=ss_mod.coords) as model:
+            P0 = pm.Deterministic('P0', pt.eye(ss_mod.k_states), dims=ss_mod.param_dims['P0'])
+
+            initial_level_trend = pm.Normal('initial_level_trend', dims=ss_mod.param_dims['initial_level_trend'])
+            sigma_level_trend = pm.HalfNormal('sigma_level_trend', dims=ss_mod.param_dims['sigma_level_trend'])
+
+            business_cycle = pm.Normal("business_cycle", dims=ss_mod.param_dims["business_cycle"])
+            dampening = pm.Beta("dampening_factor_business_cycle", 2, 2)
+            sigma_cycle = pm.HalfNormal("sigma_business_cycle", sigma=1)
+
+            ss_mod.build_statespace_graph(data)
+            idata = pm.sample(
+                nuts_sampler="nutpie", nuts_sampler_kwargs={"backend": "JAX", "gradient_backend": "JAX"}
+            )
+
+    **Multivariate Example:**
+    Model cycles for multiple economic indicators with variable-specific innovation variances:
+
+    .. code:: python
+
+        # Multivariate cycle component
+        cycle = st.CycleComponent(
+            name='business_cycle',
+            cycle_length=12,
+            estimate_cycle_length=False,
+            innovations=True,
+            dampen=True,
+            observed_state_names=['gdp', 'unemployment', 'inflation']
+        )
+        ss_mod = cycle.build()
+
+        with pm.Model(coords=ss_mod.coords) as model:
+            P0 = pm.Deterministic("P0", pt.eye(ss_mod.k_states), dims=ss_mod.param_dims["P0"])
+            # Initial states: shape (3, 2) for 3 variables, 2 states each
+            business_cycle = pm.Normal('business_cycle', dims=ss_mod.param_dims["business_cycle"])
+
+            # Dampening factor: scalar (shared across variables)
+            dampening = pm.Beta("dampening_factor_business_cycle", 2, 2)
+
+            # Innovation variances: shape (3,) for variable-specific variances
+            sigma_cycle = pm.HalfNormal(
+                "sigma_business_cycle", dims=ss_mod.param_dims["sigma_business_cycle"]
+            )
+
+            ss_mod.build_statespace_graph(data)
+            idata = pm.sample(
+                nuts_sampler="nutpie", nuts_sampler_kwargs={"backend": "JAX", "gradient_backend": "JAX"}
+            )
+
+    References
+    ----------
+    .. [1] Durbin, James, and Siem Jan Koopman. 2012.
+        Time Series Analysis by State Space Methods: Second Edition.
+        Oxford University Press.
+    """
+
+    def __init__(
+        self,
+        name: str | None = None,
+        cycle_length: int | None = None,
+        estimate_cycle_length: bool = False,
+        dampen: bool = False,
+        innovations: bool = True,
+        observed_state_names: list[str] | None = None,
+    ):
+        if observed_state_names is None:
+            observed_state_names = ["data"]
+
+        if cycle_length is None and not estimate_cycle_length:
+            raise ValueError("Must specify cycle_length if estimate_cycle_length is False")
+        if cycle_length is not None and estimate_cycle_length:
+            raise ValueError("Cannot specify cycle_length if estimate_cycle_length is True")
+        if name is None:
+            cycle = int(cycle_length) if cycle_length is not None else "Estimate"
+            name = f"Cycle[s={cycle}, dampen={dampen}, innovations={innovations}]"
+
+        self.estimate_cycle_length = estimate_cycle_length
+        self.cycle_length = cycle_length
+        self.innovations = innovations
+        self.dampen = dampen
+        self.n_coefs = 1
+
+        k_endog = len(observed_state_names)
+
+        k_states = 2 * k_endog
+        k_posdef = 2 * k_endog
+
+        obs_state_idx = np.zeros(k_states)
+        obs_state_idx[slice(0, k_states, 2)] = 1
+
+        super().__init__(
+            name=name,
+            k_endog=k_endog,
+            k_states=k_states,
+            k_posdef=k_posdef,
+            measurement_error=False,
+            combine_hidden_states=True,
+            obs_state_idxs=obs_state_idx,
+            observed_state_names=observed_state_names,
+        )
+
+    def make_symbolic_graph(self) -> None:
+        Z = np.array([1.0, 0.0]).reshape((1, -1))
+        design_matrix = block_diag(*[Z for _ in range(self.k_endog)])
+        self.ssm["design", :, :] = pt.as_tensor_variable(design_matrix)
+
+        # selection matrix R defines structure of innovations (always identity for cycle components)
+        # when innovations=False, state cov Q=0, hence R @ Q @ R.T = 0
+        R = np.eye(2)  # 2x2 identity for each cycle component
+        selection_matrix = block_diag(*[R for _ in range(self.k_endog)])
+        self.ssm["selection", :, :] = pt.as_tensor_variable(selection_matrix)
+
+        init_state = self.make_and_register_variable(
+            f"{self.name}", shape=(self.k_endog, 2) if self.k_endog > 1 else (self.k_states,)
+        )
+        self.ssm["initial_state", :] = init_state.ravel()
+
+        if self.estimate_cycle_length:
+            lamb = self.make_and_register_variable(f"length_{self.name}", shape=())
+        else:
+            lamb = self.cycle_length
+
+        if self.dampen:
+            rho = self.make_and_register_variable(f"dampening_factor_{self.name}", shape=())
+        else:
+            rho = 1
+
+        T = rho * _frequency_transition_block(lamb, j=1)
+        transition = block_diag(*[T for _ in range(self.k_endog)])
+        self.ssm["transition"] = pt.specify_shape(transition, (self.k_states, self.k_states))
+
+        if self.innovations:
+            if self.k_endog == 1:
+                sigma_cycle = self.make_and_register_variable(f"sigma_{self.name}", shape=())
+                self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * sigma_cycle**2
+            else:
+                sigma_cycle = self.make_and_register_variable(
+                    f"sigma_{self.name}", shape=(self.k_endog,)
+                )
+                state_cov = block_diag(
+                    *[pt.eye(2) * sigma_cycle[i] ** 2 for i in range(self.k_endog)]
+                )
+                self.ssm["state_cov"] = pt.specify_shape(state_cov, (self.k_states, self.k_states))
+        else:
+            # explicitly set state cov to 0 when no innovations
+            self.ssm["state_cov", :, :] = pt.zeros((self.k_posdef, self.k_posdef))
+
+    def populate_component_properties(self):
+        self.state_names = [
+            f"{f}_{self.name}[{var_name}]" if self.k_endog > 1 else f"{f}_{self.name}"
+            for var_name in self.observed_state_names
+            for f in ["Cos", "Sin"]
+        ]
+
+        self.param_names = [f"{self.name}"]
+
+        if self.k_endog == 1:
+            self.param_dims = {self.name: (f"state_{self.name}",)}
+            self.coords = {f"state_{self.name}": self.state_names}
+            self.param_info = {
+                f"{self.name}": {
+                    "shape": (2,),
+                    "constraints": None,
+                    "dims": (f"state_{self.name}",),
+                }
+            }
+        else:
+            self.param_dims = {self.name: (f"endog_{self.name}", f"state_{self.name}")}
+            self.coords = {
+                f"state_{self.name}": [f"Cos_{self.name}", f"Sin_{self.name}"],
+                f"endog_{self.name}": self.observed_state_names,
+            }
+            self.param_info = {
+                f"{self.name}": {
+                    "shape": (self.k_endog, 2),
+                    "constraints": None,
+                    "dims": (f"endog_{self.name}", f"state_{self.name}"),
+                }
+            }
+
+        if self.estimate_cycle_length:
+            self.param_names += [f"length_{self.name}"]
+            self.param_info[f"length_{self.name}"] = {
+                "shape": () if self.k_endog == 1 else (self.k_endog,),
+                "constraints": "Positive, non-zero",
+                "dims": None if self.k_endog == 1 else f"endog_{self.name}",
+            }
+
+        if self.dampen:
+            self.param_names += [f"dampening_factor_{self.name}"]
+            self.param_info[f"dampening_factor_{self.name}"] = {
+                "shape": () if self.k_endog == 1 else (self.k_endog,),
+                "constraints": "0 < x ≤ 1",
+                "dims": None if self.k_endog == 1 else (f"endog_{self.name}",),
+            }
+
+        if self.innovations:
+            self.param_names += [f"sigma_{self.name}"]
+            if self.k_endog == 1:
+                self.param_info[f"sigma_{self.name}"] = {
+                    "shape": (),
+                    "constraints": "Positive",
+                    "dims": None,
+                }
+            else:
+                self.param_dims[f"sigma_{self.name}"] = (f"endog_{self.name}",)
+                self.param_info[f"sigma_{self.name}"] = {
+                    "shape": (self.k_endog,),
+                    "constraints": "Positive",
+                    "dims": (f"endog_{self.name}",),
+                }
+            self.shock_names = self.state_names.copy()