pymc-extras 0.2.0__py3-none-any.whl

pymc_extras/statespace/core/statespace.py

@@ -0,0 +1,2268 @@
+import logging
+
+from collections.abc import Callable, Sequence
+from typing import Any
+
+import numpy as np
+import pandas as pd
+import pymc as pm
+import pytensor
+import pytensor.tensor as pt
+
+from arviz import InferenceData
+from pymc.model import modelcontext
+from pymc.model.transform.optimization import freeze_dims_and_data
+from pymc.util import RandomState
+from pytensor import Variable, graph_replace
+from pytensor.compile import get_mode
+
+from pymc_extras.statespace.core.representation import PytensorRepresentation
+from pymc_extras.statespace.filters import (
+    KalmanSmoother,
+    SquareRootFilter,
+    StandardFilter,
+    UnivariateFilter,
+)
+from pymc_extras.statespace.filters.distributions import (
+    LinearGaussianStateSpace,
+    MvNormalSVD,
+    SequenceMvNormal,
+)
+from pymc_extras.statespace.filters.utilities import stabilize
+from pymc_extras.statespace.utils.constants import (
+    ALL_STATE_AUX_DIM,
+    ALL_STATE_DIM,
+    FILTER_OUTPUT_DIMS,
+    FILTER_OUTPUT_TYPES,
+    JITTER_DEFAULT,
+    MATRIX_DIMS,
+    MATRIX_NAMES,
+    OBS_STATE_DIM,
+    SHOCK_DIM,
+    SHORT_NAME_TO_LONG,
+    TIME_DIM,
+    VECTOR_VALUED,
+)
+from pymc_extras.statespace.utils.data_tools import register_data_with_pymc
+
+_log = logging.getLogger("pymc.experimental.statespace")
+
+floatX = pytensor.config.floatX
+FILTER_FACTORY = {
+    "standard": StandardFilter,
+    "univariate": UnivariateFilter,
+    "cholesky": SquareRootFilter,
+}
+
+
+def _validate_filter_arg(filter_arg):
+    if filter_arg.lower() not in FILTER_OUTPUT_TYPES:
+        raise ValueError(
+            f'filter_output should be one of {", ".join(FILTER_OUTPUT_TYPES)}, received {filter_arg}'
+        )
+
+
+def _verify_group(group):
+    if group not in ["prior", "posterior"]:
+        raise ValueError(f'Argument "group" must be one of "prior" or "posterior", found {group}')
+
+
+class PyMCStateSpace:
+    r"""
+    Base class for Linear Gaussian Statespace models in PyMC.
+
+    Holds a ``PytensorRepresentation`` and ``KalmanFilter``, and provides a mapping between a PyMC model and the
+    statespace model.
+
+    Parameters
+    ----------
+    k_endog : int
+        The number of endogenous variables (observed time series).
+
+    k_states : int
+        The number of state variables.
+
+    k_posdef : int
+        The number of shocks in the model
+
+    filter_type : str, optional
+        The type of Kalman filter to use. Valid options are "standard", "univariate", "single", "cholesky", and
+        "steady_state". For more information, see the docs for each filter. Default is "standard".
+
+    verbose : bool, optional
+        If True, displays information about the initialized model. Defaults to True.
+
+    measurement_error : bool, optional
+        If true, the model contains measurement error. Needed by post-estimation sampling methods to decide how to
+        compute the observation errors. If False, these errors are deterministically zero; if True, they are sampled
+        from a multivariate normal.
+
+    Notes
+    -----
+    Based on the statsmodels statespace implementation https://github.com/statsmodels/statsmodels/blob/main/statsmodels/tsa/statespace/representation.py,
+    described in [1].
+
+    All statespace models inherit from this base class, which is responsible for providing an interface between a
+    PyMC model and a PytensorRepresentation of a linear statespace model. This is done via the ``update`` method,
+    which takes as input a vector of PyMC random variables and assigns them to their correct positions inside the
+    underlying ``PytensorRepresentation``. Construction of the parameter vector, called ``theta``, is done
+    automatically, but depend on the names provided in the ``param_names`` property.
+
+    To implement a new statespace model, one needs to:
+
+    1. Overload the ``param_names`` property to return a list of parameter names.
+    2. Overload the ``update`` method to put these parameters into their respective statespace matrices
+
+    In addition, a number of additional properties can be overloaded to provide users with additional resources
+    when writing their PyMC models. For details, see the attributes section of the docs for this class.
+
+    Finally, this class holds post-estimation methods common to all statespace models, which do not need to be
+    overloaded when writing a custom statespace model.
+
+    Examples
+    --------
+    The local level model is a simple statespace model. It is a Gaussian random walk with a drift term that itself also
+    follows a Gaussian random walk, as described by the following two equations:
+
+    .. math::
+        \begin{align}
+            y_{t} &= y_{t-1} + x_t + \nu_t \tag{1} \\
+            x_{t} &= x_{t-1} + \eta_t \tag{2}
+        \end{align}
+
+    Where :math:`y_t` is the observed data, and :math:`x_t` is an unobserved trend term. The model has two unknown
+    parameters, the variances on the two innovations, :math:`sigma_\nu` and :math:`sigma_\eta`. Take the hidden state
+    vector to be :math:`\begin{bmatrix} y_t & x_t \end{bmatrix}^T` and the shock vector
+    :math:`\varepsilon_t = \begin{bmatrix} \nu_t & \eta_t \end{bmatrix}^T`. Then this model can be cast into state-space
+    form with the following matrices:
+
+    .. math::
+        \begin{align}
+            T &= \begin{bmatrix}1 & 1 \\ 0 & 1 \end{bmatrix} &
+            R &= \begin{bmatrix}1 & 0 \\ 0 & 1 \end{bmatrix} &
+            Q &= \begin{bmatrix} \sigma_\nu & 0 \\ 0 & \sigma_\eta \end{bmatrix} &
+            Z &= \begin{bmatrix} 1 & 0 \end{bmatrix}
+        \end{align}
+
+    With the remaining statespace matrices as zero matrices of the appropriate sizes. The model has two states,
+    two shocks, and one observed state. Knowing all this, a very simple local level model can be implemented as
+    follows:
+
+    .. code:: python
+
+        from pymc_extras.statespace.core import PyMCStateSpace
+        import numpy as np
+
+        class LocalLevel(PyMCStateSpace):
+            def __init__():
+                # Initialize the superclass. This creates the PytensorRepresentation and the Kalman Filter
+                super().__init__(k_endog=1, k_states=2, k_posdef=2)
+
+                # Declare the non-zero, non-parameterized matrices
+                self.ssm['transition', :, :] = np.array([[1.0, 1.0], [0.0, 1.0]])
+                self.ssm['selection', :, :] = np.eye(2)
+                self.ssm['design', :, :] = np.array([[1.0, 0.0]])
+
+            @property
+            def param_names(self):
+                return ['x0', 'P0', 'sigma_nu', 'sigma_eta']
+
+            def make_symbolic_graph(self):
+                # Declare symbolic variables that represent parameters of the model
+                # In this case, we have 4: x0 (initial state), P0 (initial state covariance), sigma_nu, and sigma_eta
+
+                x0 = self.make_and_register_variable('x0', shape=(2,))
+                P0 = self.make_and_register_variable('P0', shape=(2,2))
+                sigma_mu = self.make_and_register_variable('sigma_nu')
+                sigma_eta = self.make_and_register_variable('sigma_eta')
+
+                # Next, use these symbolic variables to build the statespace matrices by assigning each parameter
+                # to its correct location in the correct matrix
+
+                self.ssm['initial_state', :] = x0
+                self.ssm['initial_state_cov', :, :] = P0
+                self.ssm['state_cov', 0, 0] = sigma_nu
+                self.ssm['state_cov', 1, 1] = sigma_eta
+
+    After defining priors over the named parameters ``P0``, ``x0``, ``sigma_eta``, and ``sigma_nu``, we can sample
+    from this model:
+
+    .. code:: python
+
+        import pymc as pm
+
+        ll = LocalLevel()
+
+        with pm.Model() as mod:
+            x0 = pm.Normal('x0', shape=(2,))
+            P0_diag = pm.Exponential('P0_diag', 1, shape=(2,))
+            P0 = pm.Deterministic('P0', pt.diag(P0_diag))
+            sigma_nu = pm.Exponential('sigma_nu', 1)
+            sigma_eta = pm.Exponential('sigma_eta', 1)
+
+            ll.build_statespace_graph(data = data)
+            idata = pm.sample()
+
+
+    References
+    ----------
+    .. [1] Fulton, Chad. "Estimating time series models by state space methods in Python: Statsmodels." (2015).
+       http://www.chadfulton.com/files/fulton_statsmodels_2017_v1.pdf
+
+    """
+
+    def __init__(
+        self,
+        k_endog: int,
+        k_states: int,
+        k_posdef: int,
+        filter_type: str = "standard",
+        verbose: bool = True,
+        measurement_error: bool = False,
+    ):
+        self._fit_mode: str | None = None
+        self._fit_coords: dict[str, Sequence[str]] | None = None
+        self._fit_dims: dict[str, Sequence[str]] | None = None
+        self._fit_data: pt.TensorVariable | None = None
+
+        self._needs_exog_data = None
+        self._exog_names = []
+        self._exog_data_info = {}
+        self._name_to_variable = {}
+        self._name_to_data = {}
+
+        self.k_endog = k_endog
+        self.k_states = k_states
+        self.k_posdef = k_posdef
+        self.measurement_error = measurement_error
+
+        # All models contain a state space representation and a Kalman filter
+        self.ssm = PytensorRepresentation(k_endog, k_states, k_posdef)
+
+        # This will be populated with PyMC random matrices after calling _insert_random_variables
+        self.subbed_ssm: list[pt.TensorVariable] | None = None
+
+        if filter_type.lower() not in FILTER_FACTORY.keys():
+            raise NotImplementedError(
+                "The following are valid filter types: " + ", ".join(list(FILTER_FACTORY.keys()))
+            )
+
+        if filter_type == "single" and self.k_endog > 1:
+            raise ValueError('Cannot use filter_type = "single" with multiple observed time series')
+
+        self.kalman_filter = FILTER_FACTORY[filter_type.lower()]()
+        self.kalman_smoother = KalmanSmoother()
+        self.make_symbolic_graph()
+
+        if verbose:
+            # These are split into separate try-except blocks, because it will be quite rare of models to implement
+            # _print_data_requirements, but we still want to print the prior requirements.
+            try:
+                self._print_prior_requirements()
+            except NotImplementedError:
+                pass
+            try:
+                self._print_data_requirements()
+            except NotImplementedError:
+                pass
+
+    def _print_prior_requirements(self) -> None:
+        """
+        Prints a short report to the terminal about the priors needed for the model, including their names,
+        shapes, named dimensions, and any parameter constraints.
+        """
+        out = ""
+        for param, info in self.param_info.items():
+            out += f'\t{param} -- shape: {info["shape"]}, constraints: {info["constraints"]}, dims: {info["dims"]}\n'
+        out = out.rstrip()
+
+        _log.info(
+            "The following parameters should be assigned priors inside a PyMC "
+            f"model block: \n"
+            f"{out}"
+        )
+
+    def _print_data_requirements(self) -> None:
+        """
+        Prints a short report to the terminal about the data needed for the model, including their names, shapes,
+        and named dimensions.
+        """
+        if not self.data_info:
+            return
+
+        out = ""
+        for data, info in self.data_info.items():
+            out += f'\t{data} -- shape: {info["shape"]}, dims: {info["dims"]}\n'
+        out = out.rstrip()
+
+        _log.info(
+            "The following Data variables should be assigned to the model inside a PyMC "
+            f"model block: \n"
+            f"{out}"
+        )
+
+    def _unpack_statespace_with_placeholders(
+        self,
+    ) -> tuple[
+        pt.TensorVariable,
+        pt.TensorVariable,
+        pt.TensorVariable,
+        pt.TensorVariable,
+        pt.TensorVariable,
+        pt.TensorVariable,
+        pt.TensorVariable,
+        pt.TensorVariable,
+        pt.TensorVariable,
+    ]:
+        """
+        Helper function to quickly obtain all statespace matrices in the standard order. Matrices returned by this
+        method will include pytensor placeholders.
+        """
+
+        a0 = self.ssm["initial_state"]
+        P0 = self.ssm["initial_state_cov"]
+        c = self.ssm["state_intercept"]
+        d = self.ssm["obs_intercept"]
+        T = self.ssm["transition"]
+        Z = self.ssm["design"]
+        R = self.ssm["selection"]
+        H = self.ssm["obs_cov"]
+        Q = self.ssm["state_cov"]
+
+        return a0, P0, c, d, T, Z, R, H, Q
+
+    def unpack_statespace(self) -> list[pt.TensorVariable]:
+        """
+        Helper function to quickly obtain all statespace matrices in the standard order.
+        """
+
+        if self.subbed_ssm is None:
+            raise ValueError(
+                "Cannot unpack the complete statespace system until PyMC model variables have been "
+                "inserted. To build the random statespace matrices, call build_statespace_graph() inside"
+                "a PyMC model context. "
+            )
+
+        return self.subbed_ssm
+
+    @property
+    def param_names(self) -> list[str]:
+        """
+        Names of model parameters
+
+        A list of all parameters expected by the model. Each parameter will be sought inside the active PyMC model
+        context when ``build_statespace_graph`` is invoked.
+        """
+        raise NotImplementedError("The param_names property has not been implemented!")
+
+    @property
+    def data_names(self) -> list[str]:
+        """
+        Names of data variables expected by the model.
+
+        This does not include the observed data series, which is automatically handled by PyMC. This property only
+        needs to be implemented for models that expect exogenous data.
+        """
+        return []
+
+    @property
+    def param_info(self) -> dict[str, dict[str, Any]]:
+        """
+        Information about parameters needed to declare priors
+
+        A dictionary of param_name: dictionary key-value pairs. The return value is used by the
+        ``_print_prior_requirements`` method, to print a message telling users how to define the necessary priors for
+        the model. Each dictionary should have the following key-value pairs:
+            * key: "shape", value: a tuple of integers
+            * key: "constraints", value: a string describing the support of the prior (positive,
+              positive semi-definite, etc)
+            * key: "dims", value: tuple of strings
+        """
+        raise NotImplementedError("The params_info property has not been implemented!")
+
+    @property
+    def data_info(self) -> dict[str, dict[str, Any]]:
+        """
+        Information about Data variables that need to be declared in the PyMC model block.
+
+        Returns a dictionary of data_name: dictionary of property-name:property description pairs. The return value is
+        used by the ``_print_data_requirements`` method, to print a message telling users how to define the necessary
+        data for the model. Each dictionary should have the following key-value pairs:
+            * key: "shape", value: a tuple of integers
+            * key: "dims", value: tuple of strings
+        """
+        raise NotImplementedError("The data_info property has not been implemented!")
+
+    @property
+    def state_names(self) -> list[str]:
+        """
+        A k_states length list of strings, associated with the model's hidden states
+
+        """
+
+        raise NotImplementedError("The state_names property has not been implemented!")
+
+    @property
+    def observed_states(self) -> list[str]:
+        """
+        A k_endog length list of strings, associated with the model's observed states
+        """
+        raise NotImplementedError("The observed_states property has not been implemented!")
+
+    @property
+    def shock_names(self) -> list[str]:
+        """
+        A k_posdef length list of strings, associated with the model's shock processes
+
+        """
+        raise NotImplementedError("The shock_names property has not been implemented!")
+
+    @property
+    def default_priors(self) -> dict[str, Callable]:
+        """
+        Dictionary of parameter names and callable functions to construct default priors for the model
+
+        Returns a dictionary with param_name: Callable key-value pairs. Used by the ``add_default_priors()`` method
+        to automatically add priors to the PyMC model.
+        """
+        raise NotImplementedError("The default_priors property has not been implemented!")
+
+    @property
+    def coords(self) -> dict[str, Sequence[str]]:
+        """
+        PyMC model coordinates
+
+        Returns a dictionary of dimension: coordinate key-value pairs, to be provided to ``pm.Model``. Dimensions
+        should come from the default names defined in ``statespace.utils.constants`` for them to be detected by
+        sampling methods.
+        """
+        raise NotImplementedError("The coords property has not been implemented!")
+
+    @property
+    def param_dims(self) -> dict[str, Sequence[str]]:
+        """
+        Dictionary of named dimensions for each model parameter
+
+        Returns a dictionary of param_name: dimension key-value pairs, to be provided to the ``dims`` argument of a
+        PyMC random variable. Dimensions should come from the default names defined in ``statespace.utils.constants``
+        for them to be detected by sampling methods.
+
+        """
+        raise NotImplementedError("The param_dims property has not been implemented!")
+
+    def add_default_priors(self) -> None:
+        """
+        Add default priors to the active PyMC model context
+        """
+        raise NotImplementedError("The add_default_priors property has not been implemented!")
+
+    def make_and_register_variable(
+        self, name, shape: int | tuple[int, ...] | None = None, dtype=floatX
+    ) -> pt.TensorVariable:
+        """
+        Helper function to create a pytensor symbolic variable and register it in the _name_to_variable dictionary
+
+        Parameters
+        ----------
+        name : str
+            The name of the placeholder variable. Must be the name of a model parameter.
+        shape : int or tuple of int
+            Shape of the parameter
+        dtype : str, default pytensor.config.floatX
+            dtype of the parameter
+
+        Notes
+        -----
+        Symbolic pytensor variables are used in the ``make_symbolic_graph`` method as placeholders for PyMC random
+        variables. The change is made in the ``_insert_random_variables`` method via ``pytensor.graph_replace``. To
+        make the change, a dictionary mapping pytensor variables to PyMC random variables needs to be constructed.
+
+        The purpose of this method is to:
+            1.  Create the placeholder symbolic variables
+            2.  Register the placeholder variable in the ``_name_to_variable`` dictionary
+
+        The shape provided here will define the shape of the prior that will need to be provided by the user.
+
+        An error is raised if the provided name has already been registered, or if the name is not present in the
+        ``param_names`` property.
+        """
+        if name not in self.param_names:
+            raise ValueError(
+                f"{name} is not a model parameter. All placeholder variables should correspond to model "
+                f"parameters."
+            )
+
+        if name in self._name_to_variable.keys():
+            raise ValueError(
+                f"{name} is already a registered placeholder variable with shape "
+                f"{self._name_to_variable[name].type.shape}"
+            )
+
+        placeholder = pt.tensor(name, shape=shape, dtype=dtype)
+        self._name_to_variable[name] = placeholder
+        return placeholder
+
+    def make_and_register_data(
+        self, name: str, shape: int | tuple[int], dtype: str = floatX
+    ) -> Variable:
+        r"""
+        Helper function to create a pytensor symbolic variable and register it in the _name_to_data dictionary
+
+        Parameters
+        ----------
+        name : str
+            The name of the placeholder data. Must be the name of an expected data variable.
+        shape : int or tuple of int
+            Shape of the parameter
+        dtype : str, default pytensor.config.floatX
+            dtype of the parameter
+
+        Notes
+        -----
+        See docstring for make_and_register_variable for more details. This function is similar, but handles data
+        inputs instead of model parameters.
+
+        An error is raised if the provided name has already been registered, or if the name is not present in the
+        ``data_names`` property.
+        """
+        if name not in self.data_names:
+            raise ValueError(
+                f"{name} is not a model parameter. All placeholder variables should correspond to model "
+                f"parameters."
+            )
+
+        if name in self._name_to_data.keys():
+            raise ValueError(
+                f"{name} is already a registered placeholder variable with shape "
+                f"{self._name_to_data[name].type.shape}"
+            )
+
+        placeholder = pt.tensor(name, shape=shape, dtype=dtype)
+        self._name_to_data[name] = placeholder
+        return placeholder
+
+    def make_symbolic_graph(self) -> None:
+        """
+        The purpose of the make_symbolic_graph function is to hide tedious parameter allocations from the user.
+        In statespace models, it is extremely rare for an entire matrix to be defined by a single prior distribution.
+        Instead, users expect to place priors over single entries of the matrix. The purpose of this function is to
+        meet that expectation.
+
+        Every statespace model needs to implement this function.
+
+        Examples
+        --------
+        As an example, consider an ARMA(2,2) model, which has five parameters (excluding the initial state distribution):
+        2 AR parameters (:math:`\rho_1` and :math:`\rho_2`), 2 MA parameters (:math:`\theta_1` and :math:`theta_2`),
+        and a single innovation covariance (:math:`\\sigma`). A common way of writing this statespace is:
+
+        ..math::
+
+            \begin{align}
+                T &= \begin{bmatrix} \rho_1 & 1 & 0 \\
+                                     \rho_2 & 0 & 1 \\
+                                     0      & 0 & 0
+                      \\end{bmatrix} \\
+                R & = \begin{bmatrix} 1 \\ \theta_1 \\ \theta_2 \\end{bmatrix} \\
+                Q &= \begin{bmatrix} \\sigma \\end{bmatrix}
+            \\end{align}
+
+        To implement this model, we begin by creating the required matrices, and fill in the "fixed" values -- the ones
+        at position (0, 1) and (0, 2) in the T matrix, and at position (0, 0) in the R matrix. These are then saved
+        to the class's PytensorRepresentation -- called ``ssm``.
+
+        .. code:: python
+
+            T = np.eye(2, k=1)
+            R = np.concatenate([np.ones(1,1), np.zeros((2, 1))], axis=0)
+
+            self.ssm['transition'] = T
+            self.ssm['selection'] = R
+
+        Next, placeholders need to be inserted for the random variables rho_1, rho_2, theta_1, theta_2, and sigma.
+        This can be done many ways: we could define two vectors, rho and theta, and a scalar for sigma, or five
+        scalars. Whatever is chosen, the choice needs to be consistent with the ``param_names`` property.
+
+        Suppose the ``param_names`` are ``[rho, theta, sigma]``, then we make one placeholder for each, and insert it
+        into the correct ``ssm`` matrix, at the correct location. To create placeholders, use the
+        ``make_and_register_variable`` helper method, which will maintain an internal registry of variables.
+
+        .. code:: python
+            rho_parmas = self.make_and_register_variable(name='rho', shape=(2,))
+            theta_params = self.make_and_register_variable(name='theta', shape=(2,))
+            sigma = self.make_and_register_variable(name='sigma', shape=(1,))
+
+            self.ssm['transition', :, 0] = rho_params
+            self.ssm['selection', 1:, 0] = theta_params
+            self.ssm['state_cov', 0, 0] = sigma
+        """
+        raise NotImplementedError("The make_symbolic_statespace method has not been implemented!")
+
+    def _get_matrix_shape_and_dims(self, name: str) -> tuple[tuple[int] | None, tuple[str] | None]:
+        """
+        Get the shape and dimensions of a matrix associated with the specified name.
+
+        This method retrieves the shape and dimensions of a matrix associated with the given name. Importantly,
+        it will only find named dimension if they are the "default" dimension names defined in the
+        statespace.utils.constant.py file.
+
+        Parameters
+        ----------
+        name : str
+            The name of the matrix whose shape and dimensions need to be retrieved.
+
+        Returns
+        -------
+        shape: tuple or None
+            If no named dimension are found, the shape of the requested matrix, otherwise None.
+
+        dims: tuple or None
+            If named dimensions are found, a tuple of strings, otherwise None
+        """
+
+        pm_mod = modelcontext(None)
+        dims = MATRIX_DIMS.get(name, None)
+        dims = dims if all([dim in pm_mod.coords.keys() for dim in dims]) else None
+        data_len = len(self._fit_data)
+
+        if name in self.kalman_filter.seq_names:
+            shape = (data_len, *self.ssm[SHORT_NAME_TO_LONG[name]].type.shape)
+            dims = (TIME_DIM, *dims)
+        else:
+            shape = self.ssm[SHORT_NAME_TO_LONG[name]].type.shape
+
+        shape = shape if dims is None else None
+
+        return shape, dims
+
+    def _save_exogenous_data_info(self):
+        """
+        Store exogenous data required by posterior sampling functions
+        """
+        pymc_mod = modelcontext(None)
+        for data_name in self.data_names:
+            data = pymc_mod[data_name]
+            self._exog_data_info[data_name] = {
+                "name": data_name,
+                "value": data.get_value(),
+                "dims": pymc_mod.named_vars_to_dims.get(data_name, None),
+            }
+
+    def _insert_random_variables(self):
+        """
+        Replace pytensor symbolic variables with PyMC random variables.
+
+        Examples
+        --------
+        .. code:: python
+
+            ss_mod = pmss.BayesianSARIMA(order=(2, 0, 2), verbose=False, stationary_initialization=True)
+            with pm.Model():
+                 x0 = pm.Normal('x0', size=ss_mod.k_states)
+                 ar_params = pm.Normal('ar_params', size=ss_mod.p)
+                 ma_parama = pm.Normal('ma_params', size=ss_mod.q)
+                 sigma_state = pm.Normal('sigma_state')
+
+                 ss_mod._insert_random_variables()
+                 matrics = ss_mod.unpack_statespace()
+
+            pm.draw(matrices['transition'], random_seed=RANDOM_SEED)
+            >>> array([[-0.90590386,  1.        ,  0.        ],
+            >>>        [ 1.25190143,  0.        ,  1.        ],
+            >>>        [ 0.        ,  0.        ,  0.        ]])
+
+            pm.draw(matrices['selection'], random_seed=RANDOM_SEED)
+            >>> array([[ 1.        ],
+            >>>        [-2.46741039],
+            >>>        [-0.28947689]])
+
+            pm.draw(matrices['state_cov'], random_seed=RANDOM_SEED)
+            >>> array([[-1.69353533]])
+        """
+
+        pymc_model = modelcontext(None)
+        found_params = []
+        with pymc_model:
+            for param_name in self.param_names:
+                param = getattr(pymc_model, param_name, None)
+                if param:
+                    found_params.append(param.name)
+
+        missing_params = list(set(self.param_names) - set(found_params))
+        if len(missing_params) > 0:
+            raise ValueError(
+                "The following required model parameters were not found in the PyMC model: "
+                + ", ".join(missing_params)
+            )
+
+        excess_params = list(set(found_params) - set(self.param_names))
+        if len(excess_params) > 0:
+            raise ValueError(
+                "The following parameters were found in the PyMC model but are not required by the statespace model: "
+                + ", ".join(excess_params)
+            )
+
+        matrices = list(self._unpack_statespace_with_placeholders())
+
+        replacement_dict = {var: pymc_model[name] for name, var in self._name_to_variable.items()}
+        self.subbed_ssm = graph_replace(matrices, replace=replacement_dict, strict=True)
+
+    def _insert_data_variables(self):
+        """
+        Replace symbolic pytensor variables with PyMC data containers.
+
+        Only used when models require exogenous data. The observed data is not added to the model using this method!
+        """
+
+        try:
+            data_names = self.data_names
+        except NotImplementedError:
+            return
+
+        pymc_model = modelcontext(None)
+        found_data = []
+        with pymc_model:
+            for data_name in data_names:
+                data = getattr(pymc_model, data_name, None)
+                if data:
+                    found_data.append(data.name)
+
+        missing_data = list(set(data_names) - set(found_data))
+        if len(missing_data) > 0:
+            raise ValueError(
+                "The following required exogenous data were not found in the PyMC model: "
+                + ", ".join(missing_data)
+            )
+
+        replacement_dict = {data: pymc_model[name] for name, data in self._name_to_data.items()}
+        self.subbed_ssm = graph_replace(self.subbed_ssm, replace=replacement_dict, strict=True)
+
+    def _register_matrices_with_pymc_model(self) -> list[pt.TensorVariable]:
+        """
+        Add all statespace matrices to the PyMC model currently on the context stack as pm.Deterministic nodes, and
+        adds named dimensions if they are found.
+
+        Returns
+        -------
+        registered_matrices: list of pt.TensorVariable
+            list of statespace matrices, wrapped in pm.Deterministic
+        """
+
+        pm_mod = modelcontext(None)
+        matrices = self.unpack_statespace()
+
+        registered_matrices = []
+        for i, (matrix, name) in enumerate(zip(matrices, MATRIX_NAMES)):
+            time_varying_ndim = 2 if name in VECTOR_VALUED else 3
+            if not getattr(pm_mod, name, None):
+                shape, dims = self._get_matrix_shape_and_dims(name)
+                has_dims = dims is not None
+
+                if matrix.ndim == time_varying_ndim and has_dims:
+                    dims = (TIME_DIM, *dims)
+
+                x = pm.Deterministic(name, matrix, dims=dims)
+                registered_matrices.append(x)
+            else:
+                registered_matrices.append(matrices[i])
+
+        return registered_matrices
+
+    @staticmethod
+    def _register_kalman_filter_outputs_with_pymc_model(outputs: tuple[pt.TensorVariable]) -> None:
+        mod = modelcontext(None)
+        coords = mod.coords
+
+        states, covs = outputs[:4], outputs[4:]
+
+        state_names = [
+            "filtered_state",
+            "predicted_state",
+            "predicted_observed_state",
+            "smoothed_state",
+        ]
+        cov_names = [
+            "filtered_covariance",
+            "predicted_covariance",
+            "predicted_observed_covariance",
+            "smoothed_covariance",
+        ]
+
+        with mod:
+            for var, name in zip(states + covs, state_names + cov_names):
+                dim_names = FILTER_OUTPUT_DIMS.get(name, None)
+                dims = tuple([dim if dim in coords.keys() else None for dim in dim_names])
+                pm.Deterministic(name, var, dims=dims)
+
+    def build_statespace_graph(
+        self,
+        data: np.ndarray | pd.DataFrame | pt.TensorVariable,
+        register_data: bool = True,
+        mode: str | None = None,
+        missing_fill_value: float | None = None,
+        cov_jitter: float | None = JITTER_DEFAULT,
+        save_kalman_filter_outputs_in_idata: bool = False,
+    ) -> None:
+        """
+        Given a parameter vector `theta`, constructs the full computational graph describing the state space model and
+        the associated log probability of the data. Hidden states and log probabilities are computed via the Kalman
+        Filter.
+
+        Parameters
+        ----------
+        data : Union[np.ndarray, pd.DataFrame, pt.TensorVariable]
+            The observed data used to fit the state space model. It can be a NumPy array, a Pandas DataFrame,
+            or a Pytensor tensor variable.
+
+        register_data : bool, optional, default=True
+            If True, the observed data will be registered with PyMC as a pm.Data variable. In addition,
+            a "time" dim will be created an added to the model's coords.
+
+        mode : Optional[str], optional, default=None
+            The Pytensor mode used for the computation graph construction. If None, the default mode will be used.
+            Other options include "JAX" and "NUMBA".
+
+        missing_fill_value: float, optional, default=-9999
+            A value to mask in missing values. NaN values in the data need to be filled with an arbitrary value to
+            avoid triggering PyMC's automatic imputation machinery (missing values are instead filled by treating them
+            as hidden states during Kalman filtering).
+
+            In general this never needs to be set. But if by a wild coincidence your data includes the value -9999.0,
+            you will need to change the missing_fill_value to something else, to avoid incorrectly mark in
+            data as missing.
+
+        cov_jitter: float, default 1e-8 or 1e-6 if pytensor.config.floatX is float32
+            The Kalman filter is known to be numerically unstable, especially at half precision. This value is added to
+            the diagonal of every covariance matrix -- predicted, filtered, and smoothed -- at every step, to ensure
+            all matrices are strictly positive semi-definite.
+
+            Obviously, if this can be zero, that's best. In general:
+                - Having measurement error makes Kalman Filters more robust. A large source of numerical errors come
+                  from the Filtered and Smoothed covariance matrices having a zero in the (0, 0) position, which always
+                  occurs when there is no measurement error. You can lower this value in the presence of measurement
+                  error.
+
+                - The Univariate Filter is more robust than other filters, and can tolerate a lower jitter value
+
+        save_kalman_filter_outputs_in_idata: bool, optional, default=False
+            If True, Kalman Filter outputs will be saved in the model as deterministics. Useful for debugging, but
+            should not be necessary for the majority of users.
+        """
+        pm_mod = modelcontext(None)
+
+        self._insert_random_variables()
+        self._save_exogenous_data_info()
+        self._insert_data_variables()
+
+        obs_coords = pm_mod.coords.get(OBS_STATE_DIM, None)
+        self._fit_data = data
+
+        data, nan_mask = register_data_with_pymc(
+            data,
+            n_obs=self.ssm.k_endog,
+            obs_coords=obs_coords,
+            register_data=register_data,
+            missing_fill_value=missing_fill_value,
+        )
+
+        filter_outputs = self.kalman_filter.build_graph(
+            pt.as_tensor_variable(data),
+            *self.unpack_statespace(),
+            mode=mode,
+            missing_fill_value=missing_fill_value,
+            cov_jitter=cov_jitter,
+        )
+
+        logp = filter_outputs.pop(-1)
+        states, covs = filter_outputs[:3], filter_outputs[3:]
+        filtered_states, predicted_states, observed_states = states
+        filtered_covariances, predicted_covariances, observed_covariances = covs
+        if save_kalman_filter_outputs_in_idata:
+            smooth_states, smooth_covariances = self._build_smoother_graph(
+                filtered_states, filtered_covariances, self.unpack_statespace(), mode=mode
+            )
+            all_kf_outputs = [*states, smooth_states, *covs, smooth_covariances]
+            self._register_kalman_filter_outputs_with_pymc_model(all_kf_outputs)
+
+        obs_dims = FILTER_OUTPUT_DIMS["predicted_observed_state"]
+        obs_dims = obs_dims if all([dim in pm_mod.coords.keys() for dim in obs_dims]) else None
+
+        SequenceMvNormal(
+            "obs",
+            mus=observed_states,
+            covs=observed_covariances,
+            logp=logp,
+            observed=data,
+            dims=obs_dims,
+        )
+
+        self._fit_coords = pm_mod.coords.copy()
+        self._fit_dims = pm_mod.named_vars_to_dims.copy()
+        self._fit_mode = mode
+
+    def _build_smoother_graph(
+        self,
+        filtered_states: pt.TensorVariable,
+        filtered_covariances: pt.TensorVariable,
+        matrices,
+        mode: str | None = None,
+        cov_jitter=JITTER_DEFAULT,
+    ):
+        """
+        Build the computation graph for the Kalman smoother.
+
+        This method constructs the computation graph for applying the Kalman smoother to the filtered states
+        and covariances obtained from the Kalman filter. The Kalman smoother is used to generate smoothed
+        estimates of the latent states and their covariances in a state space model.
+
+        The Kalman smoother provides a more accurate estimate of the latent states by incorporating future
+        information in the backward pass, resulting in smoothed state trajectories.
+
+        Parameters
+        ----------
+        filtered_states : pytensor.tensor.TensorVariable
+            The filtered states obtained from the Kalman filter. Returned by the `build_statespace_graph` method.
+
+        filtered_covariances : pytensor.tensor.TensorVariable
+            The filtered state covariances obtained from the Kalman filter. Returned by the `build_statespace_graph`
+            method.
+
+        mode : Optional[str], default=None
+            The mode used by pytensor for the construction of the logp graph. If None, the mode provided to
+            `build_statespace_graph` will be used.
+
+        Returns
+        -------
+        Tuple[pytensor.tensor.TensorVariable, pytensor.tensor.TensorVariable]
+            A tuple containing TensorVariables representing the smoothed states and smoothed state covariances
+            obtained from the Kalman smoother.
+        """
+
+        pymc_model = modelcontext(None)
+        with pymc_model:
+            *_, T, Z, R, H, Q = matrices
+
+            smooth_states, smooth_covariances = self.kalman_smoother.build_graph(
+                T, R, Q, filtered_states, filtered_covariances, mode=mode, cov_jitter=cov_jitter
+            )
+            smooth_states.name = "smooth_states"
+            smooth_covariances.name = "smooth_covariances"
+
+            return smooth_states, smooth_covariances
+
+    def _build_dummy_graph(self) -> None:
+        """
+        Build a dummy computation graph for the state space model matrices.
+
+        This method creates "dummy" pm.Flat variables representing the deep parameters used in the state space model.
+
+        Returns
+        -------
+        list[pm.Flat]
+            A list of pm.Flat variables representing all parameters estimated by the model.
+        """
+        for name in self.param_names:
+            pm.Flat(
+                name,
+                shape=self._name_to_variable[name].type.shape,
+                dims=self._fit_dims.get(name, None),
+            )
+
+    def _kalman_filter_outputs_from_dummy_graph(
+        self,
+        data: pt.TensorLike | None = None,
+        data_dims: str | tuple[str] | list[str] | None = None,
+        scenario: dict[str, pd.DataFrame] | pd.DataFrame | None = None,
+    ) -> tuple[list[pt.TensorVariable], list[tuple[pt.TensorVariable, pt.TensorVariable]]]:
+        """
+        Builds a Kalman filter graph using "dummy" pm.Flat distributions for the model variables and sorts the returns
+        into (mean, covariance) pairs for each of filtered, predicted, and smoothed output.
+
+        Parameters
+        ----------
+        data: pt.TensorLike, optional
+            Observed data on which to condition the model. If not provided, the function will use the data that was
+            provided when the model was built.
+        data_dims: str or tuple of str, optional
+            Dimension names associated with the model data. If None, defaults to ("time", "obs_state")
+
+        Returns
+        -------
+        matrices: list of tensors
+            Statespace matrices with dummy parameters.
+
+        grouped_outputs: list of tuple of tensors
+            A list of tuples, each containing the mean and covariance of the filtered, predicted, and smoothed states.
+        """
+        if scenario is None:
+            scenario = dict()
+
+        pm_mod = modelcontext(None)
+        self._build_dummy_graph()
+        self._insert_random_variables()
+
+        for name in self.data_names:
+            if name not in pm_mod:
+                pm.Data(**self._exog_data_info[name])
+
+        self._insert_data_variables()
+
+        for name in self.data_names:
+            if name in scenario.keys():
+                pm.set_data({name: scenario[name]})
+
+        x0, P0, c, d, T, Z, R, H, Q = self.unpack_statespace()
+
+        if data is None:
+            data = self._fit_data
+
+        obs_coords = pm_mod.coords.get(OBS_STATE_DIM, None)
+
+        data, nan_mask = register_data_with_pymc(
+            data,
+            n_obs=self.ssm.k_endog,
+            obs_coords=obs_coords,
+            data_dims=data_dims,
+            register_data=True,
+        )
+
+        filter_outputs = self.kalman_filter.build_graph(
+            data,
+            x0,
+            P0,
+            c,
+            d,
+            T,
+            Z,
+            R,
+            H,
+            Q,
+            mode=self._fit_mode,
+        )
+
+        filter_outputs.pop(-1)
+        states, covariances = filter_outputs[:3], filter_outputs[3:]
+
+        filtered_states, predicted_states, _ = states
+        filtered_covariances, predicted_covariances, _ = covariances
+
+        [smoothed_states, smoothed_covariances] = self.kalman_smoother.build_graph(
+            T, R, Q, filtered_states, filtered_covariances, mode=self._fit_mode
+        )
+
+        grouped_outputs = [
+            (filtered_states, filtered_covariances),
+            (predicted_states, predicted_covariances),
+            (smoothed_states, smoothed_covariances),
+        ]
+
+        return [x0, P0, c, d, T, Z, R, H, Q], grouped_outputs
+
+    def _sample_conditional(
+        self,
+        idata: InferenceData,
+        group: str,
+        random_seed: RandomState | None = None,
+        data: pt.TensorLike | None = None,
+        **kwargs,
+    ):
+        """
+        Common functionality shared between `sample_conditional_prior` and `sample_conditional_posterior`. See those
+        methods for details.
+
+        Parameters
+        ----------
+        idata : InferenceData
+            An Arviz InferenceData object containing the posterior distribution over model parameters.
+
+        group : str
+            InferenceData group from which to draw samples. Should be one of "prior" or "posterior".
+
+        random_seed : int, RandomState or Generator, optional
+            Seed for the random number generator.
+
+        data: pt.TensorLike, optional
+            Observed data on which to condition the model. If not provided, the function will use the data that was
+            provided when the model was built.
+
+        kwargs:
+            Additional keyword arguments are passed to pymc.sample_posterior_predictive
+
+        Returns
+        -------
+        InferenceData
+            An Arviz InferenceData object containing sampled trajectories from the requested conditional distribution,
+            with data variables "filtered_{group}", "predicted_{group}", and "smoothed_{group}".
+        """
+        if data is None and self._fit_data is None:
+            raise ValueError("No data provided to condition the model")
+
+        _verify_group(group)
+        group_idata = getattr(idata, group)
+
+        with pm.Model(coords=self._fit_coords) as forward_model:
+            (
+                [
+                    x0,
+                    P0,
+                    c,
+                    d,
+                    T,
+                    Z,
+                    R,
+                    H,
+                    Q,
+                ],
+                grouped_outputs,
+            ) = self._kalman_filter_outputs_from_dummy_graph(data=data)
+
+            for name, (mu, cov) in zip(FILTER_OUTPUT_TYPES, grouped_outputs):
+                dummy_ll = pt.zeros_like(mu)
+
+                state_dims = (
+                    (TIME_DIM, ALL_STATE_DIM)
+                    if all([dim in self._fit_coords for dim in [TIME_DIM, ALL_STATE_DIM]])
+                    else (None, None)
+                )
+                obs_dims = (
+                    (TIME_DIM, OBS_STATE_DIM)
+                    if all([dim in self._fit_coords for dim in [TIME_DIM, OBS_STATE_DIM]])
+                    else (None, None)
+                )
+
+                SequenceMvNormal(
+                    f"{name}_{group}",
+                    mus=mu,
+                    covs=cov,
+                    logp=dummy_ll,
+                    dims=state_dims,
+                )
+
+                obs_mu = (Z @ mu[..., None]).squeeze(-1)
+                obs_cov = Z @ cov @ pt.swapaxes(Z, -2, -1) + H
+
+                SequenceMvNormal(
+                    f"{name}_{group}_observed",
+                    mus=obs_mu,
+                    covs=obs_cov,
+                    logp=dummy_ll,
+                    dims=obs_dims,
+                )
+
+        # TODO: Remove this after pm.Flat initial values are fixed
+        forward_model.rvs_to_initial_values = {
+            rv: None for rv in forward_model.rvs_to_initial_values.keys()
+        }
+
+        frozen_model = freeze_dims_and_data(forward_model)
+        with frozen_model:
+            idata_conditional = pm.sample_posterior_predictive(
+                group_idata,
+                var_names=[
+                    f"{name}_{group}{suffix}"
+                    for name in FILTER_OUTPUT_TYPES
+                    for suffix in ["", "_observed"]
+                ],
+                compile_kwargs={"mode": get_mode(self._fit_mode)},
+                random_seed=random_seed,
+                **kwargs,
+            )
+
+        return idata_conditional.posterior_predictive
+
+    def _sample_unconditional(
+        self,
+        idata: InferenceData,
+        group: str,
+        steps: int | None = None,
+        use_data_time_dim: bool = False,
+        random_seed: RandomState | None = None,
+        **kwargs,
+    ):
+        """
+        Draw unconditional sample trajectories according to state space dynamics, using random samples from the
+        a provided trace. The state space update equations are:
+
+            X[t+1] = T @ X[t] + R @ eta[t], eta ~ N(0, Q)
+            Y[t] = Z @ X[t] + nu[t], nu ~ N(0, H)
+            x[0] ~ N(a0, P0)
+
+        Parameters
+        ----------
+        idata : InferenceData
+            An Arviz InferenceData object with a posterior group containing samples from the
+            posterior distribution over model parameters.
+
+        steps : Optional[int], default=None
+            The number of time steps to sample for the unconditional trajectories. If not provided (None),
+            the function will sample trajectories for the entire available time dimension in the posterior.
+            Otherwise, it will generate trajectories for the specified number of steps.
+
+        use_data_time_dim : bool, default=False
+            If True, the function uses the time dimension present in the provided `idata` object to sample
+            unconditional trajectories. If False, a custom time dimension is created based on the number of steps
+            specified, or if steps is None, it uses the entire available time dimension in the posterior.
+
+        random_seed : int, RandomState or Generator, optional
+            Seed for the random number generator.
+
+        kwargs:
+            Additional keyword arguments are passed to pymc.sample_posterior_predictive
+
+        Returns
+        -------
+        InferenceData
+            An Arviz InfereceData with two groups, posterior_latent and posterior_observed
+
+            - posterior_latent represents the latent state trajectories `X[t]`, which follows the dynamics:
+              `x[t+1] = T @ x[t] + R @ eta[t]`, where `eta ~ N(0, Q)`.
+
+            - posterior_observed represents the observed state trajectories `Y[t]`, which is obtained from
+              the latent state trajectories: `y[t] = Z @ x[t] + nu[t]`, where `nu ~ N(0, H)`.
+        """
+        _verify_group(group)
+        group_idata = getattr(idata, group)
+        dims = None
+        temp_coords = self._fit_coords.copy()
+
+        if not use_data_time_dim and steps is not None:
+            temp_coords.update({TIME_DIM: np.arange(1 + steps, dtype="int")})
+            steps = len(temp_coords[TIME_DIM]) - 1
+        elif steps is not None:
+            n_dimsteps = len(temp_coords[TIME_DIM])
+            if n_dimsteps != steps:
+                raise ValueError(
+                    f"Length of time dimension does not match specified number of steps, expected"
+                    f" {n_dimsteps} steps, or steps=None."
+                )
+        else:
+            steps = len(temp_coords[TIME_DIM]) - 1
+
+        if all([dim in self._fit_coords for dim in [TIME_DIM, ALL_STATE_DIM, OBS_STATE_DIM]]):
+            dims = [TIME_DIM, ALL_STATE_DIM, OBS_STATE_DIM]
+
+        with pm.Model(coords=temp_coords if dims is not None else None) as forward_model:
+            self._build_dummy_graph()
+            self._insert_random_variables()
+
+            for name in self.data_names:
+                pm.Data(**self._exog_data_info[name])
+
+            self._insert_data_variables()
+
+            matrices = [x0, P0, c, d, T, Z, R, H, Q] = self.unpack_statespace()
+
+            if not self.measurement_error:
+                H_jittered = pm.Deterministic(
+                    "H_jittered", pt.specify_shape(stabilize(H), (self.k_endog, self.k_endog))
+                )
+                matrices = [x0, P0, c, d, T, Z, R, H_jittered, Q]
+
+            LinearGaussianStateSpace(
+                group,
+                *matrices,
+                steps=steps,
+                dims=dims,
+                mode=self._fit_mode,
+                sequence_names=self.kalman_filter.seq_names,
+                k_endog=self.k_endog,
+            )
+
+        # TODO: Remove this after pm.Flat has its initial_value fixed
+        forward_model.rvs_to_initial_values = {
+            rv: None for rv in forward_model.rvs_to_initial_values.keys()
+        }
+        frozen_model = freeze_dims_and_data(forward_model)
+
+        with frozen_model:
+            idata_unconditional = pm.sample_posterior_predictive(
+                group_idata,
+                var_names=[f"{group}_latent", f"{group}_observed"],
+                compile_kwargs={"mode": self._fit_mode},
+                random_seed=random_seed,
+                **kwargs,
+            )
+
+        return idata_unconditional.posterior_predictive
+
+    def sample_conditional_prior(
+        self, idata: InferenceData, random_seed: RandomState | None = None, **kwargs
+    ) -> InferenceData:
+        """
+        Sample from the conditional prior; that is, given parameter draws from the prior distribution,
+        compute Kalman filtered trajectories. Trajectories are drawn from a single multivariate normal with mean and
+        covariance computed via either the Kalman filter, smoother, or predictions.
+
+        Parameters
+        ----------
+        idata : InferenceData
+            Arviz InferenceData with prior samples for state space matrices x0, P0, c, d, T, Z, R, H, Q.
+            Obtained from `pm.sample_prior_predictive` after calling PyMCStateSpace.build_statespace_graph().
+
+        random_seed : int, RandomState or Generator, optional
+            Seed for the random number generator.
+
+        kwargs:
+            Additional keyword arguments are passed to pymc.sample_posterior_predictive
+
+        Returns
+        -------
+        InferenceData
+            An Arviz InferenceData object containing sampled trajectories from the conditional prior.
+            The trajectories are stored in the posterior_predictive group with names "filtered_prior",
+             "predicted_prior", and "smoothed_prior".
+        """
+
+        return self._sample_conditional(idata, "prior", random_seed, **kwargs)
+
+    def sample_conditional_posterior(
+        self, idata: InferenceData, random_seed: RandomState | None = None, **kwargs
+    ):
+        """
+        Sample from the conditional posterior; that is, given parameter draws from the posterior distribution,
+        compute Kalman filtered trajectories. Trajectories are drawn from a single multivariate normal with mean and
+        covariance computed via either the Kalman filter, smoother, or predictions.
+
+        Parameters
+        ----------
+        idata : InferenceData
+            An Arviz InferenceData object containing the posterior distribution over model parameters.
+
+        random_seed : int, RandomState or Generator, optional
+            Seed for the random number generator.
+
+        kwargs:
+            Additional keyword arguments are passed to pymc.sample_posterior_predictive
+
+        Returns
+        -------
+        InferenceData
+            An Arviz InferenceData object containing sampled trajectories from the conditional posterior.
+            The trajectories are stored in the posterior_predictive group with names "filtered_posterior",
+             "predicted_posterior", and "smoothed_posterior".
+        """
+
+        return self._sample_conditional(idata, "posterior", random_seed, **kwargs)
+
+    def sample_unconditional_prior(
+        self,
+        idata: InferenceData,
+        steps: int | None = None,
+        use_data_time_dim: bool = False,
+        random_seed: RandomState | None = None,
+        **kwargs,
+    ) -> InferenceData:
+        """
+        Draw unconditional sample trajectories according to state space dynamics, using random samples from the prior
+        distribution over model parameters. The state space update equations are:
+
+            X[t+1] = T @ X[t] + R @ eta[t], eta ~ N(0, Q)
+            Y[t] = Z @ X[t] + nu[t], nu ~ N(0, H)
+
+        Parameters
+        ----------
+        idata: InferenceData
+            Arviz InferenceData with prior samples for state space matrices x0, P0, c, d, T, Z, R, H, Q.
+            Obtained from `pm.sample_prior_predictive` after calling PyMCStateSpace.build_statespace_graph().
+
+        steps : Optional[int], default=None
+            The number of time steps to sample for the unconditional trajectories. If not provided (None),
+            the function will sample trajectories for the entire available time dimension in the posterior.
+            Otherwise, it will generate trajectories for the specified number of steps.
+
+        use_data_time_dim : bool, default=False
+            If True, the function uses the time dimension present in the provided `idata` object to sample
+            unconditional trajectories. If False, a custom time dimension is created based on the number of steps
+            specified, or if steps is None, it uses the entire available time dimension in the posterior.
+
+        random_seed : int, RandomState or Generator, optional
+            Seed for the random number generator.
+
+        kwargs:
+            Additional keyword arguments are passed to pymc.sample_posterior_predictive
+
+        Returns
+        -------
+        InferenceData
+            An Arviz InfereceData with two data variables, prior_latent and prior_observed
+
+            - prior_latent represents the latent state trajectories `X[t]`, which follows the dynamics:
+              `x[t+1] = T @ x[t] + R @ eta[t]`, where `eta ~ N(0, Q)`.
+
+            - prior_observed represents the observed state trajectories `Y[t]`, which is obtained from
+              the observation equation: `y[t] = Z @ x[t] + nu[t]`, where `nu ~ N(0, H)`.
+        """
+
+        return self._sample_unconditional(
+            idata, "prior", steps, use_data_time_dim, random_seed, **kwargs
+        )
+
+    def sample_unconditional_posterior(
+        self,
+        idata: InferenceData,
+        steps: int | None = None,
+        use_data_time_dim: bool = False,
+        random_seed: RandomState | None = None,
+        **kwargs,
+    ) -> InferenceData:
+        """
+        Draw unconditional sample trajectories according to state space dynamics, using random samples from the
+        posterior distribution over model parameters. The state space update equations are:
+
+            X[t+1] = T @ X[t] + R @ eta[t], eta ~ N(0, Q)
+            Y[t] = Z @ X[t] + nu[t], nu ~ N(0, H)
+            x[0] ~ N(a0, P0)
+
+        Parameters
+        ----------
+        idata : InferenceData
+            An Arviz InferenceData object with a posterior group containing samples from the
+            posterior distribution over model parameters.
+
+        steps : Optional[int], default=None
+            The number of time steps to sample for the unconditional trajectories. If not provided (None),
+            the function will sample trajectories for the entire available time dimension in the posterior.
+            Otherwise, it will generate trajectories for the specified number of steps.
+
+        use_data_time_dim : bool, default=False
+            If True, the function uses the time dimension present in the provided `idata` object to sample
+            unconditional trajectories. If False, a custom time dimension is created based on the number of steps
+            specified, or if steps is None, it uses the entire available time dimension in the posterior.
+
+        random_seed : int, RandomState or Generator, optional
+            Seed for the random number generator.
+
+        Returns
+        -------
+        InferenceData
+            An Arviz InfereceData with two groups, posterior_latent and posterior_observed
+
+            - posterior_latent represents the latent state trajectories `X[t]`, which follows the dynamics:
+              `x[t+1] = T @ x[t] + R @ eta[t]`, where `eta ~ N(0, Q)`.
+
+            - posterior_observed represents the observed state trajectories `Y[t]`, which is obtained from
+              the latent state trajectories: `y[t] = Z @ x[t] + nu[t]`, where `nu ~ N(0, H)`.
+        """
+
+        return self._sample_unconditional(
+            idata, "posterior", steps, use_data_time_dim, random_seed, **kwargs
+        )
+
+    def sample_statespace_matrices(
+        self, idata, matrix_names: str | list[str] | None, group: str = "posterior"
+    ):
+        """
+        Draw samples of requested statespace matrices from provided idata
+
+        Parameters
+        ----------
+        matrix_names: str, list[str], optional
+            Statespace matrices to be sampled. Valid names are short names: x0, P0, c, d, T, Z, R, H, Q, or
+             "formal" names: initial_state, initial_state_cov, state_intercept, obs_intercept, transition, design,
+                             selection, obs_cov, state_cov
+        idata: az.InferenceData
+            InferenceData from which to sample
+
+        group: str, one of "posterior" or "prior"
+            Whether to sample from priors or posteriors
+
+        Returns
+        -------
+        idata_matrices: az.InterenceData
+        """
+        _verify_group(group)
+
+        if matrix_names is None:
+            matrix_names = MATRIX_NAMES
+        elif isinstance(matrix_names, str):
+            matrix_names = [matrix_names]
+
+        with pm.Model(coords=self._fit_coords) as forward_model:
+            self._build_dummy_graph()
+            self._insert_random_variables()
+
+            for name in self.data_names:
+                pm.Data(**self._exog_data_info[name])
+
+            self._insert_data_variables()
+            matrices = self.unpack_statespace()
+            for short_name, matrix in zip(MATRIX_NAMES, matrices):
+                long_name = SHORT_NAME_TO_LONG[short_name]
+                if (long_name in matrix_names) or (short_name in matrix_names):
+                    name = long_name if long_name in matrix_names else short_name
+                    dims = [x if x in self._fit_coords else None for x in MATRIX_DIMS[short_name]]
+                    pm.Deterministic(name, matrix, dims=dims)
+
+        # TODO: Remove this after pm.Flat has its initial_value fixed
+        forward_model.rvs_to_initial_values = {
+            rv: None for rv in forward_model.rvs_to_initial_values.keys()
+        }
+        frozen_model = freeze_dims_and_data(forward_model)
+        with frozen_model:
+            matrix_idata = pm.sample_posterior_predictive(
+                idata if group == "posterior" else idata.prior,
+                var_names=matrix_names,
+                compile_kwargs={"mode": self._fit_mode},
+                extend_inferencedata=False,
+            )
+
+        return matrix_idata
+
+    @staticmethod
+    def _validate_forecast_args(
+        time_index: pd.RangeIndex | pd.DatetimeIndex,
+        start: int | pd.Timestamp,
+        periods: int | None = None,
+        end: int | pd.Timestamp = None,
+        scenario: pd.DataFrame | np.ndarray | None = None,
+        use_scenario_index: bool = False,
+        verbose: bool = True,
+    ):
+        if isinstance(start, pd.Timestamp) and start not in time_index:
+            raise ValueError("Datetime start must be in the data index used to fit the model.")
+        elif isinstance(start, int):
+            if abs(start) > len(time_index):
+                raise ValueError(
+                    "Integer start must be within the range of the data index used to fit the model."
+                )
+        if periods is None and end is None:
+            raise ValueError("Must specify one of either periods or end")
+        if periods is not None and end is not None:
+            raise ValueError("Must specify exactly one of either periods or end")
+        if scenario is None and use_scenario_index:
+            raise ValueError("use_scenario_index=True requires a scenario to be provided.")
+        if scenario is not None and use_scenario_index:
+            if isinstance(scenario, dict):
+                first_df = next(
+                    (df for df in scenario.values() if isinstance(df, pd.DataFrame | pd.Series)),
+                    None,
+                )
+                if first_df is None:
+                    raise ValueError(
+                        "use_scenario_index=True requires a scenario to be a DataFrame or Series."
+                    )
+            elif not isinstance(scenario, pd.DataFrame | pd.Series):
+                raise ValueError(
+                    "use_scenario_index=True requires a scenario to be a DataFrame or Series."
+                )
+        if use_scenario_index and any(arg is not None for arg in [start, end, periods]) and verbose:
+            _log.warning(
+                "start, end, and periods arguments are ignored when use_scenario_index is True. Pass only "
+                "one or the other to avoid this warning, or pass verbose = False."
+            )
+
+    def _get_fit_time_index(self) -> pd.RangeIndex | pd.DatetimeIndex:
+        time_index = self._fit_coords.get(TIME_DIM, None) if self._fit_coords is not None else None
+        if time_index is None:
+            raise ValueError(
+                "No time dimension found on coordinates used to fit the model. Has this model been fit?"
+            )
+
+        if isinstance(time_index[0], pd.Timestamp):
+            time_index = pd.DatetimeIndex(time_index)
+            time_index.freq = time_index.inferred_freq
+        else:
+            time_index = np.array(time_index)
+
+        return time_index
+
+    def _validate_scenario_data(
+        self,
+        scenario: pd.DataFrame | np.ndarray | dict[str, pd.DataFrame | np.ndarray] | None,
+        name: str | None = None,
+        verbose=True,
+    ):
+        """
+        Validate the scenario data provided to the forecast method by checking that it has the correct shape and
+        dimensions.
+
+        Parameters
+        ----------
+        scenario
+        name
+        verbose
+
+        Returns
+        -------
+        scenario: pd.DataFrame | np.ndarray | dict[str, pd.DataFrame | np.ndarray]
+            Scenario data, validated and potentially modified.
+
+        """
+        if not self._needs_exog_data:
+            return scenario
+
+        var_to_dims = {key: info["dims"][1:] for key, info in self.data_info.items()}
+
+        if any(len(dims) > 1 for dims in var_to_dims.values()):
+            raise NotImplementedError(">2d exogenous data is not yet supported.")
+        coords = {
+            var: self._fit_coords[dim[0]]
+            for var, dim in var_to_dims.items()
+            if dim[0] in self._fit_coords
+        }
+
+        if self._needs_exog_data and scenario is None:
+            exog_str = ",".join(self._exog_names)
+            suffix = "s" if len(exog_str) > 1 else ""
+            raise ValueError(
+                f"This model was fit using exogenous data. Forecasting cannot be performed without "
+                f"providing scenario data for the following variable{suffix}: {exog_str}"
+            )
+
+        if isinstance(scenario, dict):
+            for name, data in scenario.items():
+                if name not in self._exog_names:
+                    raise ValueError(
+                        f"Scenario data provided for variable '{name}', which is not an exogenous variable "
+                        f"used to fit the model."
+                    )
+
+                # Recursively call this function to trigger the non-dictionary branch of the checks on each object
+                # inside the dictionary
+                scenario[name] = self._validate_scenario_data(data, name)
+
+            # The provided dictionary might be a mix of numpy arrays and dataframes if the user is truly horrible.
+            # For checking shapes, the first object will always be good enough. But we also need to make sure all the
+            # indices agree, so we grab the first dataframe (which might not exist, but that's OK)
+            first_scenario = next(iter(scenario.values()))
+            first_df = next((df for df in scenario.values() if isinstance(df, pd.DataFrame)), None)
+
+            if not all(data.shape[0] == first_scenario.shape[0] for data in scenario.values()):
+                raise ValueError(
+                    "Scenario data must have the same number of time steps for all variables."
+                )
+
+            if first_df is not None and not all(
+                df.index.equals(first_df.index)
+                for df in scenario.values()
+                if isinstance(df, pd.DataFrame)
+            ):
+                raise ValueError("Scenario data must have the same index for all variables.")
+
+            return scenario
+
+        elif isinstance(scenario, pd.Series | pd.DataFrame | np.ndarray | list | tuple):
+            # A user might be lazy and pass a simple list when there is only one exogenous variable.
+            if isinstance(scenario, list | tuple) or (
+                isinstance(scenario, np.ndarray) and scenario.ndim == 1
+            ):
+                scenario = np.array(scenario).reshape(-1, 1)
+
+            if name is None:
+                # name should only be None on the first non-recursive call. We only arrive to this branch in that case
+                # if a non-dictionary was passed, which in turn should only happen if only a single exogenous data
+                # needs to be set.
+                if len(self._exog_names) > 1:
+                    raise ValueError(
+                        "Multiple exogenous variables were used to fit the model. Provide a dictionary of "
+                        "scenario data instead."
+                    )
+                name = self._exog_names[0]
+
+            # Omit dataframe from this basic shape check so we can give more detailed information about missing columns
+            # in the next check
+            if not isinstance(scenario, pd.DataFrame | pd.Series) and scenario.shape[1] != len(
+                coords[name]
+            ):
+                raise ValueError(
+                    f"Scenario data for variable '{name}' has the wrong number of columns. Expected "
+                    f"{len(coords[name])}, got {scenario.shape[1]}"
+                )
+
+            if isinstance(scenario, pd.Series):
+                if len(coords[name]) > 1:
+                    raise ValueError(
+                        f"Scenario data for variable '{name}' has the wrong number of columns. Expected "
+                        f"{len(coords[name])}, got 1"
+                    )
+
+            if isinstance(scenario, pd.DataFrame):
+                expected_cols = coords[name]
+                cols = scenario.columns
+                missing_columns = sorted(list(set(expected_cols) - set(cols)))
+                if len(missing_columns) > 0:
+                    suffix = "s" if len(missing_columns) > 1 else ""
+                    raise ValueError(
+                        f"Scenario data for variable '{name}' is missing the following column{suffix}: "
+                        f"{', '.join(missing_columns)}"
+                    )
+
+                extra_columns = sorted(list(set(cols) - set(expected_cols)))
+                if len(extra_columns) > 0:
+                    suffix = "s" if len(extra_columns) > 1 else ""
+                    verb = "is" if len(extra_columns) == 1 else "are"
+                    raise ValueError(
+                        f"Scenario data for variable '{name}' contains the following extra column{suffix} "
+                        f"that {verb} not used by the model: "
+                        f"{', '.join(extra_columns)}"
+                    )
+
+                if not (a == b for a, b in zip(expected_cols, cols)) and verbose:
+                    _log.warning(
+                        f"Scenario data for {name} has a different column order than the data used to fit the "
+                        f"model. Columns will be automatically re-ordered. Ensure consistent ordering to avoid "
+                        f"silent errors."
+                    )
+                    scenario = scenario[expected_cols]
+
+            return scenario
+
+    @staticmethod
+    def _build_forecast_index(
+        time_index: pd.RangeIndex | pd.DatetimeIndex,
+        start: int | pd.Timestamp | None = None,
+        end: int | pd.Timestamp = None,
+        periods: int | None = None,
+        use_scenario_index: bool = False,
+        scenario: pd.DataFrame | np.ndarray | None = None,
+    ) -> tuple[int | pd.Timestamp, pd.RangeIndex | pd.DatetimeIndex]:
+        """
+        Construct a pandas Index for the requested forecast horizon.
+
+        Parameters
+        ----------
+        time_index: pd.RangeIndex or pd.DatetimeIndex
+            Index of the data used to fit the model
+        start: int or pd.Timestamp, optional
+            Date from which to begin forecasting. If using a datetime index, integer start will be interpreted
+            as a positional index. Otherwise, start must be found inside the time_index
+        end: int or pd.Timestamp, optional
+            Date at which to end forecasting. If using a datetime index, end must be a timestamp.
+        periods: int, optional
+            Number of periods to forecast
+        scenario:  pd.DataFrame, np.ndarray, optional
+            Scenario data to use for forecasting. If provided, the index of the scenario data will be used as the
+            forecast index. If provided, start, end, and periods will be ignored.
+        use_scenario_index: bool, default False
+            If True, the index of the scenario data will be used as the forecast index.
+
+
+        Returns
+        -------
+        start: int | pd.TimeStamp
+            The starting date index or time step from which to generate the forecasts.
+
+        forecast_index: pd.DatetimeIndex or pd.RangeIndex
+            Index for the forecast results
+        """
+
+        def get_or_create_index(x, time_index, start=None):
+            if isinstance(x, pd.DataFrame | pd.Series):
+                return x.index
+            elif isinstance(x, dict):
+                return get_or_create_index(next(iter(x.values())), time_index, start)
+            elif isinstance(x, np.ndarray | list | tuple):
+                if start is None:
+                    raise ValueError(
+                        "Provided scenario has no index and no start date was provided. This combination "
+                        "is ambiguous. Please provide a start date, or add an index to the scenario."
+                    )
+                is_datetime_index = isinstance(time_index, pd.DatetimeIndex)
+                n = x.shape[0] if isinstance(x, np.ndarray) else len(x)
+
+                if isinstance(start, int):
+                    start = time_index[start]
+                if is_datetime_index:
+                    return pd.date_range(start, periods=n, freq=time_index.freq)
+                return pd.RangeIndex(start, n + start, step=1, dtype="int")
+
+            else:
+                raise ValueError(f"{type(x)} is not a valid type for scenario data.")
+
+        x0_idx = None
+
+        if use_scenario_index:
+            forecast_index = get_or_create_index(scenario, time_index, start)
+            is_datetime = isinstance(forecast_index, pd.DatetimeIndex)
+
+            # If the user provided an index, we want to take it as-is (without removing the start value). Instead,
+            # step one back and use this as the start value.
+            delta = forecast_index.freq if is_datetime else 1
+            x0_idx = forecast_index[0] - delta
+
+        else:
+            # Otherwise, build an index. It will be a DateTime index if we have all the necessary information, otherwise
+            # use a range index.
+            is_datetime = isinstance(time_index, pd.DatetimeIndex)
+            forecast_index = None
+
+            if is_datetime:
+                freq = time_index.freq
+                if isinstance(start, int):
+                    start = time_index[start]
+                if isinstance(end, int):
+                    raise ValueError(
+                        "end must be a timestamp if using a datetime index. To specify a number of "
+                        "timesteps from the start date, use the periods argument instead."
+                    )
+                if end is not None:
+                    forecast_index = pd.date_range(start, end=end, freq=freq)
+                if periods is not None:
+                    # date_range includes both the start and end date, but we're going to pop off the start later
+                    # (it will be interpreted as x0). So we need to add 1 to the periods so the user gets "periods"
+                    # number of forecasts back
+                    forecast_index = pd.date_range(start, periods=periods + 1, freq=freq)
+
+            else:
+                # If the user provided a positive integer as start, directly interpret it as the start time. If its
+                # negative, interpret it as a positional index.
+                if start < 0:
+                    start = time_index[start]
+                if end is not None:
+                    forecast_index = pd.RangeIndex(start, end, step=1, dtype="int")
+                if periods is not None:
+                    forecast_index = pd.RangeIndex(start, start + periods + 1, step=1, dtype="int")
+
+        if is_datetime:
+            if forecast_index.freq != time_index.freq:
+                raise ValueError(
+                    "The frequency of the forecast index must match the frequency on the data used "
+                    f"to fit the model. Got {forecast_index.freq}, expected {time_index.freq}"
+                )
+
+        if x0_idx is None:
+            x0_idx, forecast_index = forecast_index[0], forecast_index[1:]
+        if x0_idx in forecast_index:
+            raise ValueError("x0_idx should not be in the forecast index")
+        if x0_idx not in time_index:
+            raise ValueError("start must be in the data index used to fit the model.")
+
+        # The starting value should not be included in the forecast index. It will be used only to define x0 and P0,
+        # and no forecast will be associated with it.
+        return x0_idx, forecast_index
+
+    def _finalize_scenario_initialization(
+        self,
+        scenario: pd.DataFrame | np.ndarray | dict[str, pd.DataFrame | np.ndarray] | None,
+        forecast_index: pd.RangeIndex | pd.DatetimeIndex,
+        name=None,
+    ):
+        try:
+            var_to_dims = {key: info["dims"][1:] for key, info in self.data_info.items()}
+        except NotImplementedError:
+            return scenario
+
+        if any(len(dims) > 1 for dims in var_to_dims.values()):
+            raise NotImplementedError(">2d exogenous data is not yet supported.")
+        coords = {
+            var: self._fit_coords[dim[0]]
+            for var, dim in var_to_dims.items()
+            if dim[0] in self._fit_coords
+        }
+
+        if scenario is None:
+            return scenario
+
+        if isinstance(scenario, dict):
+            for name, data in scenario.items():
+                scenario[name] = self._finalize_scenario_initialization(data, forecast_index, name)
+            return scenario
+
+        # This was already checked as valid
+        name = self._exog_names[0] if name is None else name
+
+        # Small tidying up in the case we just have a single scenario that's already a dataframe.
+        if isinstance(scenario, pd.DataFrame | pd.Series):
+            if isinstance(scenario, pd.Series):
+                scenario = scenario.to_frame(name=coords[name][0])
+            if not scenario.index.equals(forecast_index):
+                scenario.index = forecast_index
+
+        # lists and tuples were handled during validation, along with shape check, so just cast arrays to dataframes
+        # with the correct index and columns
+        if isinstance(scenario, np.ndarray):
+            scenario = pd.DataFrame(scenario, index=forecast_index, columns=coords[name])
+
+        return scenario
+
+    def forecast(
+        self,
+        idata: InferenceData,
+        start: int | pd.Timestamp | None = None,
+        periods: int | None = None,
+        end: int | pd.Timestamp = None,
+        scenario: pd.DataFrame | np.ndarray | dict[str, pd.DataFrame | np.ndarray] | None = None,
+        use_scenario_index: bool = False,
+        filter_output="smoothed",
+        random_seed: RandomState | None = None,
+        verbose: bool = True,
+        **kwargs,
+    ) -> InferenceData:
+        """
+        Generate forecasts of state space model trajectories into the future.
+
+        This function combines posterior parameter samples in the provided idata with model dynamics to generate
+        forecasts for out-of-sample data. The trajectory is initialized using the filter output specified in
+        the filter_output argument.
+
+        Parameters
+        ----------
+        idata : InferenceData
+            An Arviz InferenceData object containing the posterior distribution over model parameters.
+
+        start : int or pd.Timestamp, optional
+            The starting date index or time step from which to generate the forecasts. If the data provided to
+            `PyMCStateSpace.build_statespace_graph` had a datetime index, `start` should be a datetime.
+            If using integer time series, `start` should be an integer indicating the starting time step. In either
+            case, `start` should be in the data index used to build the statespace graph.
+
+            If start is None, the last value on the data's index will be used.
+
+        periods : int, optional
+            The number of time steps to forecast into the future. If `periods` is specified, the `end`
+            parameter will be ignored. If `None`, then the `end` parameter must be provided.
+
+        end : int or pd.Timestamp, optional
+            The ending date index or time step up to which to generate the forecasts. If the data provided to
+            `PyMCStateSpace.build_statespace_graph` had a datetime index, `start` should be a datetime.
+            If using integer time series, `end` should be an integer indicating the ending time step.
+            If `end` is provided, the `periods` parameter will be ignored.
+
+        scenario: pd.Dataframe or np.ndarray, optional
+            Exogenous variables to use for scenario-based forecasting. Must be a 2d array-like, with second dimension
+            equal to the number of exogenous variables. If start, end, or periods are specified, the first dimension
+            must conform with these settings. Otherwise, the index of the scenario data will be used to set the
+            number of forecast steps. If the index of the forecast scenairo is a pandas DateTimeIndex, its frequency
+            must match the frequency of the data used to fit the model. Otherwise, dates will be based on the number
+            of forecast steps and the data.
+
+        use_scenario_index: bool, default False
+            If True, the index of the scenario data will be used to determine the forecast period. In this case,
+            the start, end, and periods arguments will be ignored. If True, the scenario data must be a DataFrame,
+            otherwise an error will be raised.
+
+        filter_output : str, default="smoothed"
+            The type of Kalman Filter output used to initialize the forecasts. The 0th timestep of the forecast will
+            be sampled from x[0] ~ N(filter_output_mean[start], filter_output_covariance[start]). Default is "smoothed",
+            which uses past and future data to make the best possible hidden state estimate.
+
+        random_seed : int, RandomState or Generator, optional
+            Seed for the random number generator.
+
+        verbose: bool, default=True
+            Whether to print diagnostic information about forecasting.
+
+        kwargs:
+            Additional keyword arguments are passed to pymc.sample_posterior_predictive
+
+        Returns
+        -------
+        InferenceData
+            An Arviz InferenceData object containing forecast samples for the latent and observed state
+            trajectories of the state space model, named  "forecast_latent" and "forecast_observed".
+
+                - forecast_latent represents the latent state trajectories `X[t]`, which follows the dynamics:
+                  `x[t+1] = T @ x[t] + R @ eta[t]`, where `eta ~ N(0, Q)`.
+
+                - forecast_observed represents the observed state trajectories `Y[t]`, which is obtained from
+                  the latent state trajectories: `y[t] = Z @ x[t] + nu[t]`, where `nu ~ N(0, H)`.
+
+        """
+        filter_time_dim = TIME_DIM
+
+        _validate_filter_arg(filter_output)
+        time_index = self._get_fit_time_index()
+
+        if start is None and verbose:
+            _log.warning(
+                "No start date provided. Using the last date in the data index. To silence this warning, "
+                "explicitly pass a start date or set verbose = False"
+            )
+            start = time_index[-1]
+
+        if self._needs_exog_data and not isinstance(scenario, dict):
+            if len(self.data_names) > 1:
+                raise ValueError(
+                    "Model needs more than one exogenous data to do forecasting. In this case, you must "
+                    "pass a dictionary of scenario data."
+                )
+            [data_name] = self.data_names
+            scenario = {data_name: scenario}
+
+        scenario: dict = self._validate_scenario_data(scenario, verbose=verbose)
+
+        self._validate_forecast_args(
+            time_index=time_index,
+            start=start,
+            end=end,
+            periods=periods,
+            scenario=scenario,
+            use_scenario_index=use_scenario_index,
+            verbose=verbose,
+        )
+
+        t0, forecast_index = self._build_forecast_index(
+            time_index=time_index,
+            start=start,
+            end=end,
+            periods=periods,
+            scenario=scenario,
+            use_scenario_index=use_scenario_index,
+        )
+        scenario = self._finalize_scenario_initialization(scenario, forecast_index)
+        temp_coords = self._fit_coords.copy()
+
+        dims = None
+        if all([dim in temp_coords for dim in [filter_time_dim, ALL_STATE_DIM, OBS_STATE_DIM]]):
+            dims = [TIME_DIM, ALL_STATE_DIM, OBS_STATE_DIM]
+
+        t0_idx = np.flatnonzero(time_index == t0)[0]
+
+        temp_coords["data_time"] = time_index
+        temp_coords[TIME_DIM] = forecast_index
+
+        mu_dims, cov_dims = None, None
+        if all([dim in self._fit_coords for dim in [TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM]]):
+            mu_dims = ["data_time", ALL_STATE_DIM]
+            cov_dims = ["data_time", ALL_STATE_DIM, ALL_STATE_AUX_DIM]
+
+        with pm.Model(coords=temp_coords) as forecast_model:
+            (_, _, *matrices), grouped_outputs = self._kalman_filter_outputs_from_dummy_graph(
+                data_dims=["data_time", OBS_STATE_DIM],
+            )
+
+            group_idx = FILTER_OUTPUT_TYPES.index(filter_output)
+            mu, cov = grouped_outputs[group_idx]
+
+            x0 = pm.Deterministic(
+                "x0_slice", mu[t0_idx], dims=mu_dims[1:] if mu_dims is not None else None
+            )
+            P0 = pm.Deterministic(
+                "P0_slice", cov[t0_idx], dims=cov_dims[1:] if cov_dims is not None else None
+            )
+
+            if scenario is not None:
+                sub_dict = {
+                    forecast_model[data_name]: pt.as_tensor_variable(
+                        scenario.get(data_name), name=data_name
+                    )
+                    for data_name in self.data_names
+                }
+
+                matrices = graph_replace(matrices, replace=sub_dict, strict=True)
+                [setattr(matrix, "name", name) for name, matrix in zip(MATRIX_NAMES[2:], matrices)]
+
+            _ = LinearGaussianStateSpace(
+                "forecast",
+                x0,
+                P0,
+                *matrices,
+                steps=len(forecast_index),
+                dims=dims,
+                mode=self._fit_mode,
+                sequence_names=self.kalman_filter.seq_names,
+                k_endog=self.k_endog,
+                append_x0=False,
+            )
+
+        forecast_model.rvs_to_initial_values = {
+            k: None for k in forecast_model.rvs_to_initial_values.keys()
+        }
+        frozen_model = freeze_dims_and_data(forecast_model)
+
+        with frozen_model:
+            idata_forecast = pm.sample_posterior_predictive(
+                idata,
+                var_names=["forecast_latent", "forecast_observed"],
+                compile_kwargs={"mode": self._fit_mode},
+                random_seed=random_seed,
+                **kwargs,
+            )
+
+        return idata_forecast.posterior_predictive
+
+    def impulse_response_function(
+        self,
+        idata,
+        n_steps: int = 40,
+        use_posterior_cov: bool = True,
+        shock_size: float | np.ndarray | None = None,
+        shock_cov: np.ndarray | None = None,
+        shock_trajectory: np.ndarray | None = None,
+        orthogonalize_shocks: bool = False,
+        random_seed: RandomState | None = None,
+        **kwargs,
+    ):
+        """
+        Generate impulse response functions (IRF) from state space model dynamics.
+
+        An impulse response function represents the dynamic response of the state space model
+        to an instantaneous shock applied to the system. This function calculates the IRF
+        based on either provided shock specifications or the posterior state covariance matrix.
+
+        Parameters
+        ----------
+        idata : az.InferenceData
+            An Arviz InferenceData object containing the posterior distribution over model parameters.
+
+        n_steps: int
+            The number of time steps to calculate the impulse response. Default is 40.
+
+            If `shock_trajectory` is provided, the length of the shock trajectory will override this value.
+
+        use_posterior_cov: bool, default=True
+            Whether to use the covariance matrix of the posterior distribution to generate the impulse response.
+
+            Only one of `use_posterior_cov`, `shock_cov`, `shock_size`, or `shock_trajectory` can be specified.
+
+        shock_size : Optional[Union[float, np.ndarray]], default=None
+            The size of the shock applied to the system. If specified, it will create a covariance
+            matrix for the shock with diagonal elements equal to `shock_size`. If float, the diagonal will be filled
+            with `shock_size`. If an array, `shock_size` must match the number of shocks in the state space model.
+
+            Only one of `use_posterior_cov`, `shock_cov`, `shock_size`, or `shock_trajectory` can be specified.
+
+        shock_cov : Optional[np.ndarray], default=None
+            A user-specified covariance matrix for the shocks. It should be a 2D numpy array with
+            dimensions (n_shocks, n_shocks), where n_shocks is the number of shocks in the state space model.
+
+            Only one of `use_posterior_cov`, `shock_cov`, `shock_size`, or `shock_trajectory` can be specified.
+
+        shock_trajectory : Optional[np.ndarray], default=None
+            A pre-defined trajectory of shocks applied to the system. It should be a 2D numpy array
+            with dimensions (n, n_shocks), where n is the number of time steps and k_posdef is the
+            number of shocks in the state space model.
+
+            Only one of `use_posterior_cov`, `shock_cov`, `shock_size`, or `shock_trajectory` can be specified.
+
+        orthogonalize_shocks : bool, default=False
+            If True, orthogonalize the shocks using Cholesky decomposition when generating the impulse
+            response. This option is ignored if `shock_trajectory` or `shock_size` are used.
+
+        random_seed : int, RandomState or Generator, optional
+            Seed for the random number generator.
+
+        kwargs:
+            Additional keyword arguments are passed to pymc.sample_posterior_predictive
+
+        Returns
+        -------
+        pm.InferenceData
+            An Arviz InferenceData object containing impulse response function in a variable named "irf".
+        """
+        options = [shock_size, shock_cov, shock_trajectory]
+        n_options = sum(x is not None for x in options)
+        Q = None  # No covariance matrix needed if a trajectory is provided. Will be overwritten later if needed.
+
+        if n_options > 1:
+            raise ValueError("Specify exactly 0 or 1 of shock_size, shock_cov, or shock_trajectory")
+        elif n_options == 1:
+            # If the user passed an alternative parameterization for the shocks of the IRF, don't use the posterior
+            use_posterior_cov = False
+
+        if shock_trajectory is not None:
+            # Validate the shock trajectory
+            n, k = shock_trajectory.shape
+            steps = n
+
+            if k != self.k_posdef:
+                raise ValueError(
+                    "If shock_trajectory is provided, there must be a trajectory provided for each shock. "
+                    f"Model has {self.k_posdef} shocks, but shock_trajectory has only {k} columns"
+                )
+            if steps is not None and steps != n:
+                _log.warning(
+                    "Both steps and shock_trajectory were provided but do not agree. Length of "
+                    "shock_trajectory will take priority, and steps will be ignored."
+                )
+            n_steps = n  # Overwrite steps with the length of the shock trajectory
+            shock_trajectory = pt.as_tensor_variable(shock_trajectory)
+
+        simulation_coords = self._fit_coords.copy()
+        simulation_coords[TIME_DIM] = np.arange(n_steps, dtype="int")
+
+        with pm.Model(coords=simulation_coords):
+            self._build_dummy_graph()
+            self._insert_random_variables()
+
+            P0, _, c, d, T, Z, R, H, post_Q = self.unpack_statespace()
+            x0 = pm.Deterministic("x0_new", pt.zeros(self.k_states), dims=[ALL_STATE_DIM])
+
+            if use_posterior_cov:
+                Q = post_Q
+                if orthogonalize_shocks:
+                    Q = pt.linalg.cholesky(Q) / pt.diag(Q)
+            elif shock_cov is not None:
+                Q = pt.as_tensor_variable(shock_cov)
+                if orthogonalize_shocks:
+                    Q = pt.linalg.cholesky(Q) / pt.diag(Q)
+
+            if shock_trajectory is None:
+                shock_trajectory = pt.zeros((n_steps, self.k_posdef))
+                if Q is not None:
+                    init_shock = MvNormalSVD("initial_shock", mu=0, cov=Q, dims=[SHOCK_DIM])
+                else:
+                    init_shock = pm.Deterministic(
+                        "initial_shock",
+                        pt.as_tensor_variable(np.atleast_1d(shock_size)),
+                        dims=[SHOCK_DIM],
+                    )
+                shock_trajectory = pt.set_subtensor(shock_trajectory[0], init_shock)
+
+            else:
+                shock_trajectory = pt.as_tensor_variable(shock_trajectory)
+
+            def irf_step(shock, x, c, T, R):
+                next_x = c + T @ x + R @ shock
+                return next_x
+
+            irf, updates = pytensor.scan(
+                irf_step,
+                sequences=[shock_trajectory],
+                outputs_info=[x0],
+                non_sequences=[c, T, R],
+                n_steps=n_steps,
+                strict=True,
+                mode=self._fit_mode,
+            )
+
+            pm.Deterministic("irf", irf, dims=[TIME_DIM, ALL_STATE_DIM])
+
+            compile_kwargs = kwargs.get("compile_kwargs", {})
+            if "mode" not in compile_kwargs.keys():
+                compile_kwargs = {"mode": self._fit_mode}
+            else:
+                mode = compile_kwargs.get("mode")
+                if mode is not None and mode != self._fit_mode:
+                    raise ValueError(
+                        f"User provided compile mode ({mode}) does not match the compile mode used to "
+                        f"construct the model ({self._fit_mode})."
+                    )
+
+                compile_kwargs.update({"mode": self._fit_mode})
+
+            irf_idata = pm.sample_posterior_predictive(
+                idata,
+                var_names=["irf"],
+                compile_kwargs=compile_kwargs,
+                random_seed=random_seed,
+                **kwargs,
+            )
+
+            return irf_idata.posterior_predictive
+
+    def _sort_obs_inputs_by_time_varying(self, d, Z):
+        seqs = []
+        non_seqs = []
+
+        for matrix, name in zip([d, Z], ["d", "Z"]):
+            if name in self.kalman_filter.seq_names:
+                seqs.append(matrix)
+            else:
+                non_seqs.append(matrix)
+
+        return seqs, non_seqs
+
+    @staticmethod
+    def _sort_obs_scan_args(args):
+        args = list(args)
+
+        # If a matrix is time-varying, pytensor will put a [t] on the name
+        arg_names = [x.name.replace("[t]", "") for x in args]
+        ordered_args = []
+
+        for name in ["d", "Z"]:
+            idx = arg_names.index(name)
+            ordered_args.append(args[idx])
+
+        return ordered_args