pymc-extras 0.2.0__py3-none-any.whl

pymc_extras/statespace/models/structural.py

@@ -0,0 +1,1651 @@
+import functools as ft
+import logging
+
+from abc import ABC
+from collections.abc import Sequence
+from itertools import pairwise
+from typing import Any
+
+import numpy as np
+import pytensor
+import pytensor.tensor as pt
+import xarray as xr
+
+from pytensor import Variable
+
+from pymc_extras.statespace.core import PytensorRepresentation
+from pymc_extras.statespace.core.statespace import PyMCStateSpace
+from pymc_extras.statespace.models.utilities import (
+    conform_time_varying_and_time_invariant_matrices,
+    make_default_coords,
+)
+from pymc_extras.statespace.utils.constants import (
+    ALL_STATE_AUX_DIM,
+    ALL_STATE_DIM,
+    AR_PARAM_DIM,
+    LONG_MATRIX_NAMES,
+    POSITION_DERIVATIVE_NAMES,
+    TIME_DIM,
+)
+
+_log = logging.getLogger("pymc.experimental.statespace")
+
+floatX = pytensor.config.floatX
+
+
+def order_to_mask(order):
+    if isinstance(order, int):
+        return np.ones(order).astype(bool)
+    else:
+        return np.array(order).astype(bool)
+
+
+def _frequency_transition_block(s, j):
+    lam = 2 * np.pi * j / s
+
+    return pt.stack([[pt.cos(lam), pt.sin(lam)], [-pt.sin(lam), pt.cos(lam)]])
+
+
+class StructuralTimeSeries(PyMCStateSpace):
+    r"""
+    Structural Time Series Model
+
+    The structural time series model, named by [1] and presented in statespace form in [2], is a framework for
+    decomposing a univariate time series into level, trend, seasonal, and cycle components. It also admits the
+    possibility of exogenous regressors. Unlike the SARIMAX framework, the time series is not assumed to be stationary.
+
+    Notes
+    -----
+
+    .. math::
+         y_t = \mu_t + \gamma_t + c_t + \varepsilon_t
+
+    """
+
+    def __init__(
+        self,
+        ssm: PytensorRepresentation,
+        state_names: list[str],
+        data_names: list[str],
+        shock_names: list[str],
+        param_names: list[str],
+        exog_names: list[str],
+        param_dims: dict[str, tuple[int]],
+        coords: dict[str, Sequence],
+        param_info: dict[str, dict[str, Any]],
+        data_info: dict[str, dict[str, Any]],
+        component_info: dict[str, dict[str, Any]],
+        measurement_error: bool,
+        name_to_variable: dict[str, Variable],
+        name_to_data: dict[str, Variable] | None = None,
+        name: str | None = None,
+        verbose: bool = True,
+        filter_type: str = "standard",
+    ):
+        # Add the initial state covariance to the parameters
+        if name is None:
+            name = "data"
+        self._name = name
+
+        k_states, k_posdef, k_endog = ssm.k_states, ssm.k_posdef, ssm.k_endog
+        param_names, param_dims, param_info = self._add_inital_state_cov_to_properties(
+            param_names, param_dims, param_info, k_states
+        )
+        self._state_names = state_names.copy()
+        self._data_names = data_names.copy()
+        self._shock_names = shock_names.copy()
+        self._param_names = param_names.copy()
+        self._param_dims = param_dims.copy()
+
+        default_coords = make_default_coords(self)
+        coords.update(default_coords)
+
+        self._coords = coords
+        self._param_info = param_info.copy()
+        self._data_info = data_info.copy()
+        self.measurement_error = measurement_error
+
+        super().__init__(
+            k_endog,
+            k_states,
+            max(1, k_posdef),
+            filter_type=filter_type,
+            verbose=verbose,
+            measurement_error=measurement_error,
+        )
+        self.ssm = ssm.copy()
+
+        if k_posdef == 0:
+            # If there is no randomness in the model, add dummy matrices to the representation to avoid errors
+            # when we go to construct random variables from the matrices
+            self.ssm.k_posdef = self.k_posdef
+            self.ssm.shapes["state_cov"] = (1, 1, 1)
+            self.ssm["state_cov"] = pt.zeros((1, 1, 1))
+
+            self.ssm.shapes["selection"] = (1, self.k_states, 1)
+            self.ssm["selection"] = pt.zeros((1, self.k_states, 1))
+
+        self._component_info = component_info.copy()
+
+        self._name_to_variable = name_to_variable.copy()
+        self._name_to_data = name_to_data.copy()
+
+        self._exog_names = exog_names.copy()
+        self._needs_exog_data = len(exog_names) > 0
+
+        P0 = self.make_and_register_variable("P0", shape=(self.k_states, self.k_states))
+        self.ssm["initial_state_cov"] = P0
+
+    @staticmethod
+    def _add_inital_state_cov_to_properties(param_names, param_dims, param_info, k_states):
+        param_names += ["P0"]
+        param_dims["P0"] = (ALL_STATE_DIM, ALL_STATE_AUX_DIM)
+        param_info["P0"] = {
+            "shape": (k_states, k_states),
+            "constraints": "Positive semi-definite",
+            "dims": param_dims["P0"],
+        }
+
+        return param_names, param_dims, param_info
+
+    @property
+    def param_names(self):
+        return self._param_names
+
+    @property
+    def data_names(self) -> list[str]:
+        return self._data_names
+
+    @property
+    def state_names(self):
+        return self._state_names
+
+    @property
+    def observed_states(self):
+        return [self._name]
+
+    @property
+    def shock_names(self):
+        return self._shock_names
+
+    @property
+    def param_dims(self):
+        return self._param_dims
+
+    @property
+    def coords(self) -> dict[str, Sequence]:
+        return self._coords
+
+    @property
+    def param_info(self) -> dict[str, dict[str, Any]]:
+        return self._param_info
+
+    @property
+    def data_info(self) -> dict[str, dict[str, Any]]:
+        return self._data_info
+
+    def make_symbolic_graph(self) -> None:
+        """
+        Assign placeholder pytensor variables among statespace matrices in positions where PyMC variables will go.
+
+        Notes
+        -----
+        This assignment is handled by the components, so this function is implemented only to avoid the
+        NotImplementedError raised by the base class.
+        """
+
+        pass
+
+    def _state_slices_from_info(self):
+        info = self._component_info.copy()
+        comp_states = np.cumsum([0] + [info["k_states"] for info in info.values()])
+        state_slices = [slice(i, j) for i, j in pairwise(comp_states)]
+
+        return state_slices
+
+    def _hidden_states_from_data(self, data):
+        state_slices = self._state_slices_from_info()
+        info = self._component_info
+        names = info.keys()
+        result = []
+
+        for i, (name, s) in enumerate(zip(names, state_slices)):
+            obs_idx = info[name]["obs_state_idx"]
+            if obs_idx is None:
+                continue
+
+            X = data[..., s]
+            if info[name]["combine_hidden_states"]:
+                sum_idx = np.flatnonzero(obs_idx)
+                result.append(X[..., sum_idx].sum(axis=-1)[..., None])
+            else:
+                comp_names = self.state_names[s]
+                for j, state_name in enumerate(comp_names):
+                    result.append(X[..., j, None])
+
+        return np.concatenate(result, axis=-1)
+
+    def _get_subcomponent_names(self):
+        state_slices = self._state_slices_from_info()
+        info = self._component_info
+        names = info.keys()
+        result = []
+
+        for i, (name, s) in enumerate(zip(names, state_slices)):
+            if info[name]["combine_hidden_states"]:
+                result.append(name)
+            else:
+                comp_names = self.state_names[s]
+                result.extend([f"{name}[{comp_name}]" for comp_name in comp_names])
+        return result
+
+    def extract_components_from_idata(self, idata: xr.Dataset) -> xr.Dataset:
+        r"""
+        Extract interpretable hidden states from an InferenceData returned by a PyMCStateSpace sampling method
+
+        Parameters
+        ----------
+        idata: Dataset
+            A Dataset object, returned by a PyMCStateSpace sampling method
+
+        Returns
+        -------
+        idata: Dataset
+            An Dataset object with hidden states transformed to represent only the "interpretable" subcomponents
+            of the structural model.
+
+        Notes
+        -----
+        In general, a structural statespace model can be represented as:
+
+        .. math::
+            y_t = \mu_t + \nu_t + \cdots + \gamma_t + c_t + \xi_t + \epsilon_t \tag{1}
+
+        Where:
+
+            - :math:`\mu_t` is the level of the data at time t
+            - :math:`\nu_t` is the slope of the data at time t
+            - :math:`\cdots` are higher time derivatives of the position (acceleration, jerk, etc) at time t
+            - :math:`\gamma_t` is the seasonal component at time t
+            - :math:`c_t` is the cycle component at time t
+            - :math:`\xi_t` is the autoregressive error at time t
+            - :math:`\varepsilon_t` is the measurement error at time t
+
+        In state space form, some or all of these components are represented as linear combinations of other
+        subcomponents, making interpretation of the outputs of the outputs difficult. The purpose of this function is
+        to take the expended statespace representation and return a "reduced form" of only the components shown in
+        equation (1).
+        """
+
+        def _extract_and_transform_variable(idata, new_state_names):
+            *_, time_dim, state_dim = idata.dims
+            state_func = ft.partial(self._hidden_states_from_data)
+            new_idata = xr.apply_ufunc(
+                state_func,
+                idata,
+                input_core_dims=[[time_dim, state_dim]],
+                output_core_dims=[[time_dim, state_dim]],
+                exclude_dims={state_dim},
+            )
+            new_idata.coords.update({state_dim: new_state_names})
+            return new_idata
+
+        var_names = list(idata.data_vars.keys())
+        is_latent = [idata[name].shape[-1] == self.k_states for name in var_names]
+        new_state_names = self._get_subcomponent_names()
+
+        latent_names = [name for latent, name in zip(is_latent, var_names) if latent]
+        dropped_vars = set(var_names) - set(latent_names)
+        if len(dropped_vars) > 0:
+            _log.warning(
+                f'Variables {", ".join(dropped_vars)} do not contain all hidden states (their last dimension '
+                f"is not {self.k_states}). They will not be present in the modified idata."
+            )
+        if len(dropped_vars) == len(var_names):
+            raise ValueError(
+                "Provided idata had no variables with all hidden states; cannot extract components."
+            )
+
+        idata_new = xr.Dataset(
+            {
+                name: _extract_and_transform_variable(idata[name], new_state_names)
+                for name in latent_names
+            }
+        )
+        return idata_new
+
+
+class Component(ABC):
+    r"""
+    Base class for a component of a structural timeseries model.
+
+    This base class contains a subset of the class attributes of the PyMCStateSpace class, and none of the class
+    methods. The purpose of a component is to allow the partial definition of a structural model. Components are
+    assembled into a full model by the StructuralTimeSeries class.
+
+    Parameters
+    ----------
+    name: str
+        The name of the component
+    k_endog: int
+        Number of endogenous variables being modeled. Currently, must be one because structural models only support
+        univariate data.
+    k_states: int
+        Number of hidden states in the component model
+    k_posdef: int
+        Rank of the state covariance matrix, or the number of sources of innovations in the component model
+    measurement_error: bool
+        Whether the observation associated with the component has measurement error. Default is False.
+    combine_hidden_states: bool
+        Flag for the ``extract_hidden_states_from_data`` method. When ``True``, hidden states from the component model
+        are extracted as ``hidden_states[:, np.flatnonzero(Z)]``. Should be True in models where hidden states
+        individually have no interpretation, such as seasonal or autoregressive components.
+    """
+
+    def __init__(
+        self,
+        name,
+        k_endog,
+        k_states,
+        k_posdef,
+        state_names=None,
+        data_names=None,
+        shock_names=None,
+        param_names=None,
+        exog_names=None,
+        representation: PytensorRepresentation | None = None,
+        measurement_error=False,
+        combine_hidden_states=True,
+        component_from_sum=False,
+        obs_state_idxs=None,
+    ):
+        self.name = name
+        self.k_endog = k_endog
+        self.k_states = k_states
+        self.k_posdef = k_posdef
+        self.measurement_error = measurement_error
+
+        self.state_names = state_names if state_names is not None else []
+        self.data_names = data_names if data_names is not None else []
+        self.shock_names = shock_names if shock_names is not None else []
+        self.param_names = param_names if param_names is not None else []
+        self.exog_names = exog_names if exog_names is not None else []
+
+        self.needs_exog_data = len(self.exog_names) > 0
+        self.coords = {}
+        self.param_dims = {}
+
+        self.param_info = {}
+        self.data_info = {}
+
+        self.param_counts = {}
+
+        if representation is None:
+            self.ssm = PytensorRepresentation(k_endog=k_endog, k_states=k_states, k_posdef=k_posdef)
+        else:
+            self.ssm = representation
+
+        self._name_to_variable = {}
+        self._name_to_data = {}
+
+        if not component_from_sum:
+            self.populate_component_properties()
+            self.make_symbolic_graph()
+
+        self._component_info = {
+            self.name: {
+                "k_states": self.k_states,
+                "k_enodg": self.k_endog,
+                "k_posdef": self.k_posdef,
+                "combine_hidden_states": combine_hidden_states,
+                "obs_state_idx": obs_state_idxs,
+            }
+        }
+
+    def make_and_register_variable(self, name, shape, dtype=floatX) -> Variable:
+        r"""
+        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, shape, dtype=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:
+        raise NotImplementedError
+
+    def populate_component_properties(self):
+        raise NotImplementedError
+
+    def _get_combined_shapes(self, other):
+        k_states = self.k_states + other.k_states
+        k_posdef = self.k_posdef + other.k_posdef
+        if self.k_endog != other.k_endog:
+            raise NotImplementedError(
+                "Merging elements with different numbers of observed states is not supported.>"
+            )
+        k_endog = self.k_endog
+
+        return k_states, k_posdef, k_endog
+
+    def _combine_statespace_representations(self, other):
+        def make_slice(name, x, o_x):
+            ndim = max(x.ndim, o_x.ndim)
+            return (name,) + (slice(None, None, None),) * ndim
+
+        k_states, k_posdef, k_endog = self._get_combined_shapes(other)
+
+        self_matrices = [self.ssm[name] for name in LONG_MATRIX_NAMES]
+        other_matrices = [other.ssm[name] for name in LONG_MATRIX_NAMES]
+
+        x0, P0, c, d, T, Z, R, H, Q = (
+            self.ssm[make_slice(name, x, o_x)]
+            for name, x, o_x in zip(LONG_MATRIX_NAMES, self_matrices, other_matrices)
+        )
+        o_x0, o_P0, o_c, o_d, o_T, o_Z, o_R, o_H, o_Q = (
+            other.ssm[make_slice(name, x, o_x)]
+            for name, x, o_x in zip(LONG_MATRIX_NAMES, self_matrices, other_matrices)
+        )
+
+        initial_state = pt.concatenate(conform_time_varying_and_time_invariant_matrices(x0, o_x0))
+        initial_state.name = x0.name
+
+        initial_state_cov = pt.linalg.block_diag(P0, o_P0)
+        initial_state_cov.name = P0.name
+
+        state_intercept = pt.concatenate(conform_time_varying_and_time_invariant_matrices(c, o_c))
+        state_intercept.name = c.name
+
+        obs_intercept = d + o_d
+        obs_intercept.name = d.name
+
+        transition = pt.linalg.block_diag(T, o_T)
+        transition.name = T.name
+
+        design = pt.concatenate(conform_time_varying_and_time_invariant_matrices(Z, o_Z), axis=-1)
+        design.name = Z.name
+
+        selection = pt.linalg.block_diag(R, o_R)
+        selection.name = R.name
+
+        obs_cov = H + o_H
+        obs_cov.name = H.name
+
+        state_cov = pt.linalg.block_diag(Q, o_Q)
+        state_cov.name = Q.name
+
+        new_ssm = PytensorRepresentation(
+            k_endog=k_endog,
+            k_states=k_states,
+            k_posdef=k_posdef,
+            initial_state=initial_state,
+            initial_state_cov=initial_state_cov,
+            state_intercept=state_intercept,
+            obs_intercept=obs_intercept,
+            transition=transition,
+            design=design,
+            selection=selection,
+            obs_cov=obs_cov,
+            state_cov=state_cov,
+        )
+
+        return new_ssm
+
+    def _combine_property(self, other, name):
+        self_prop = getattr(self, name)
+        if isinstance(self_prop, list):
+            return self_prop + getattr(other, name)
+        elif isinstance(self_prop, dict):
+            new_prop = self_prop.copy()
+            new_prop.update(getattr(other, name))
+            return new_prop
+
+    def _combine_component_info(self, other):
+        combined_info = {}
+        for key, value in self._component_info.items():
+            if not key.startswith("StateSpace"):
+                if key in combined_info.keys():
+                    raise ValueError(f"Found duplicate component named {key}")
+                combined_info[key] = value
+
+        for key, value in other._component_info.items():
+            if not key.startswith("StateSpace"):
+                if key in combined_info.keys():
+                    raise ValueError(f"Found duplicate component named {key}")
+                combined_info[key] = value
+
+        return combined_info
+
+    def _make_combined_name(self):
+        components = self._component_info.keys()
+        name = f'StateSpace[{", ".join(components)}]'
+        return name
+
+    def __add__(self, other):
+        state_names = self._combine_property(other, "state_names")
+        data_names = self._combine_property(other, "data_names")
+        param_names = self._combine_property(other, "param_names")
+        shock_names = self._combine_property(other, "shock_names")
+        param_info = self._combine_property(other, "param_info")
+        data_info = self._combine_property(other, "data_info")
+        param_dims = self._combine_property(other, "param_dims")
+        coords = self._combine_property(other, "coords")
+        exog_names = self._combine_property(other, "exog_names")
+
+        _name_to_variable = self._combine_property(other, "_name_to_variable")
+        _name_to_data = self._combine_property(other, "_name_to_data")
+
+        measurement_error = any([self.measurement_error, other.measurement_error])
+
+        k_states, k_posdef, k_endog = self._get_combined_shapes(other)
+        ssm = self._combine_statespace_representations(other)
+
+        new_comp = Component(
+            name="",
+            k_endog=1,
+            k_states=k_states,
+            k_posdef=k_posdef,
+            measurement_error=measurement_error,
+            representation=ssm,
+            component_from_sum=True,
+        )
+        new_comp._component_info = self._combine_component_info(other)
+        new_comp.name = new_comp._make_combined_name()
+
+        names_and_props = [
+            ("state_names", state_names),
+            ("data_names", data_names),
+            ("param_names", param_names),
+            ("shock_names", shock_names),
+            ("param_dims", param_dims),
+            ("coords", coords),
+            ("param_dims", param_dims),
+            ("param_info", param_info),
+            ("data_info", data_info),
+            ("exog_names", exog_names),
+            ("_name_to_variable", _name_to_variable),
+            ("_name_to_data", _name_to_data),
+        ]
+
+        for prop, value in names_and_props:
+            setattr(new_comp, prop, value)
+
+        return new_comp
+
+    def build(self, name=None, filter_type="standard", verbose=True):
+        """
+        Build a StructuralTimeSeries statespace model from the current component(s)
+
+        Parameters
+        ----------
+        name: str, optional
+            Name of the exogenous data being modeled. Default is "data"
+
+        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.
+
+        Returns
+        -------
+        PyMCStateSpace
+            An initialized instance of a PyMCStateSpace, constructed using the system matrices contained in the
+            components.
+        """
+
+        return StructuralTimeSeries(
+            self.ssm,
+            name=name,
+            state_names=self.state_names,
+            data_names=self.data_names,
+            shock_names=self.shock_names,
+            param_names=self.param_names,
+            param_dims=self.param_dims,
+            coords=self.coords,
+            param_info=self.param_info,
+            data_info=self.data_info,
+            component_info=self._component_info,
+            measurement_error=self.measurement_error,
+            exog_names=self.exog_names,
+            name_to_variable=self._name_to_variable,
+            name_to_data=self._name_to_data,
+            filter_type=filter_type,
+            verbose=verbose,
+        )
+
+
+class LevelTrendComponent(Component):
+    r"""
+    Level and trend component of a structural time series model
+
+    Parameters
+    ----------
+    __________
+    order : int
+
+        Number of time derivatives of the trend to include in the model. For example, when order=3, the trend will
+        be of the form ``y = a + b * t + c * t ** 2``, where the coefficients ``a, b, c`` come from the initial
+        state values.
+
+    innovations_order : int or sequence of int, optional
+
+        The number of stochastic innovations to include in the model. By default, ``innovations_order = order``
+
+    Notes
+    -----
+    This class implements the level and trend components of the general structural time series model. In the most
+    general form, the level and trend is described by a system of two time-varying equations.
+
+    .. math::
+        \begin{align}
+            \mu_{t+1} &= \mu_t + \nu_t + \zeta_t \\
+            \nu_{t+1} &= \nu_t + \xi_t
+            \zeta_t &\sim N(0, \sigma_\zeta) \\
+            \xi_t &\sim N(0, \sigma_\xi)
+        \end{align}
+
+    Where :math:`\mu_{t+1}` is the mean of the timeseries at time t, and :math:`\nu_t` is the drift or the slope of
+    the process. When both innovations :math:`\zeta_t` and :math:`\xi_t` are included in the model, it is known as a
+    *local linear trend* model. This system of two equations, corresponding to ``order=2``, can be expanded or
+    contracted by adding or removing equations. ``order=3`` would add an acceleration term to the sytsem:
+
+    .. math::
+        \begin{align}
+            \mu_{t+1} &= \mu_t + \nu_t + \zeta_t \\
+            \nu_{t+1} &= \nu_t + \eta_t + \xi_t \\
+            \eta_{t+1} &= \eta_{t-1} + \omega_t \\
+            \zeta_t &\sim N(0, \sigma_\zeta) \\
+            \xi_t &\sim N(0, \sigma_\xi) \\
+            \omega_t &\sim N(0, \sigma_\omega)
+        \end{align}
+
+    After setting all innovation terms to zero and defining initial states :math:`\mu_0, \nu_0, \eta_0`, these equations
+    can be collapsed to:
+
+    .. math::
+        \mu_t = \mu_0 + \nu_0 \cdot t + \eta_0 \cdot t^2
+
+    Which clarifies how the order and initial states influence the model. In particular, the initial states are the
+    coefficients on the intercept, slope, acceleration, and so on.
+
+    In this light, allowing for innovations can be understood as allowing these coefficients to vary over time. Each
+    component can be individually selected for time variation by passing a list to the ``innovations_order`` argument.
+    For example, a constant intercept with time varying trend and acceleration is specified as ``order=3,
+    innovations_order=[0, 1, 1]``.
+
+    By choosing the ``order`` and ``innovations_order``, a large variety of models can be obtained. Notable
+    models include:
+
+    * Constant intercept, ``order=1, innovations_order=0``
+
+    .. math::
+        \mu_t = \mu
+
+    * Constant linear slope, ``order=2, innovations_order=0``
+
+    .. math::
+        \mu_t = \mu_{t-1} + \nu
+
+    * Gaussian Random Walk, ``order=1, innovations_order=1``
+
+    .. math::
+        \mu_t = \mu_{t-1} + \zeta_t
+
+    * Gaussian Random Walk with Drift, ``order=2, innovations_order=1``
+
+    .. math::
+        \mu_t = \mu_{t-1} + \nu + \zeta_t
+
+    * Smooth Trend, ``order=2, innovations_order=[0, 1]``
+
+    .. math::
+        \begin{align}
+            \mu_t &= \mu_{t-1} + \nu_{t-1} \\
+            \nu_t &= \nu_{t-1} + \xi_t
+        \end{align}
+
+    * Local Level, ``order=2, innovations_order=2``
+
+    [1] notes that the smooth trend model produces more gradually changing slopes than the full local linear trend
+    model, and is equivalent to an "integrated trend model".
+
+    References
+    ----------
+    .. [1] Durbin, James, and Siem Jan Koopman. 2012.
+        Time Series Analysis by State Space Methods: Second Edition.
+        Oxford University Press.
+
+    """
+
+    def __init__(
+        self,
+        order: int | list[int] = 2,
+        innovations_order: int | list[int] | None = None,
+        name: str = "LevelTrend",
+    ):
+        if innovations_order is None:
+            innovations_order = order
+
+        self._order_mask = order_to_mask(order)
+        max_state = np.flatnonzero(self._order_mask)[-1].item() + 1
+
+        # If the user passes excess zeros, raise an error. The alternative is to prune them, but this would cause
+        # the shape of the state to be different to what the user expects.
+        if len(self._order_mask) > max_state:
+            raise ValueError(
+                f"order={order} is invalid. The highest derivative should not be set to zero. If you want a "
+                f"lower order model, explicitly omit the zeros."
+            )
+        k_states = max_state
+
+        if isinstance(innovations_order, int):
+            n = innovations_order
+            innovations_order = order_to_mask(k_states)
+            if n > 0:
+                innovations_order[n:] = False
+            else:
+                innovations_order[:] = False
+        else:
+            innovations_order = order_to_mask(innovations_order)
+
+        self.innovations_order = innovations_order[:max_state]
+        k_posdef = int(sum(innovations_order))
+
+        super().__init__(
+            name,
+            k_endog=1,
+            k_states=k_states,
+            k_posdef=k_posdef,
+            measurement_error=False,
+            combine_hidden_states=False,
+            obs_state_idxs=np.array([1.0] + [0.0] * (k_states - 1)),
+        )
+
+    def populate_component_properties(self):
+        name_slice = POSITION_DERIVATIVE_NAMES[: self.k_states]
+        self.param_names = ["initial_trend"]
+        self.state_names = [name for name, mask in zip(name_slice, self._order_mask) if mask]
+        self.param_dims = {"initial_trend": ("trend_state",)}
+        self.coords = {"trend_state": self.state_names}
+        self.param_info = {"initial_trend": {"shape": (self.k_states,), "constraints": None}}
+
+        if self.k_posdef > 0:
+            self.param_names += ["sigma_trend"]
+            self.shock_names = [
+                name for name, mask in zip(name_slice, self.innovations_order) if mask
+            ]
+            self.param_dims["sigma_trend"] = ("trend_shock",)
+            self.coords["trend_shock"] = self.shock_names
+            self.param_info["sigma_trend"] = {"shape": (self.k_posdef,), "constraints": "Positive"}
+
+        for name in self.param_names:
+            self.param_info[name]["dims"] = self.param_dims[name]
+
+    def make_symbolic_graph(self) -> None:
+        initial_trend = self.make_and_register_variable("initial_trend", shape=(self.k_states,))
+        self.ssm["initial_state", :] = initial_trend
+        triu_idx = np.triu_indices(self.k_states)
+        self.ssm[np.s_["transition", triu_idx[0], triu_idx[1]]] = 1
+
+        R = np.eye(self.k_states)
+        R = R[:, self.innovations_order]
+        self.ssm["selection", :, :] = R
+
+        self.ssm["design", 0, :] = np.array([1.0] + [0.0] * (self.k_states - 1))
+
+        if self.k_posdef > 0:
+            sigma_trend = self.make_and_register_variable("sigma_trend", shape=(self.k_posdef,))
+            diag_idx = np.diag_indices(self.k_posdef)
+            idx = np.s_["state_cov", diag_idx[0], diag_idx[1]]
+            self.ssm[idx] = sigma_trend**2
+
+
+class MeasurementError(Component):
+    r"""
+    Measurement error term for a structural timeseries model
+
+    Parameters
+    ----------
+    name: str, optional
+
+        Name of the observed data. Default is "obs".
+
+    Notes
+    -----
+    This component should only be used in combination with other components, because it has no states. It's only use
+    is to add a variance parameter to the model, associated with the observation noise matrix H.
+
+    Examples
+    --------
+    Create and estimate a deterministic linear trend with measurement error
+
+    .. code:: python
+
+        from pymc_extras.statespace import structural as st
+        import pymc as pm
+        import pytensor.tensor as pt
+
+        trend = st.LevelTrendComponent(order=2, innovations_order=0)
+        error = st.MeasurementError()
+        ss_mod = (trend + error).build()
+
+        with pm.Model(coords=ss_mod.coords) as model:
+            P0 = pm.Deterministic('P0', pt.eye(ss_mod.k_states) * 10, dims=ss_mod.param_dims['P0'])
+            intitial_trend = pm.Normal('initial_trend', sigma=10, dims=ss_mod.param_dims['initial_trend'])
+            sigma_obs = pm.Exponential('sigma_obs', 1, dims=ss_mod.param_dims['sigma_obs'])
+
+            ss_mod.build_statespace_graph(data, mode='JAX')
+            idata = pm.sample(nuts_sampler='numpyro')
+    """
+
+    def __init__(self, name: str = "MeasurementError"):
+        k_endog = 1
+        k_states = 0
+        k_posdef = 0
+
+        super().__init__(
+            name, k_endog, k_states, k_posdef, measurement_error=True, combine_hidden_states=False
+        )
+
+    def populate_component_properties(self):
+        self.param_names = [f"sigma_{self.name}"]
+        self.param_dims = {}
+        self.param_info = {
+            f"sigma_{self.name}": {
+                "shape": (),
+                "constraints": "Positive",
+                "dims": None,
+            }
+        }
+
+    def make_symbolic_graph(self) -> None:
+        sigma_shape = ()
+        error_sigma = self.make_and_register_variable(f"sigma_{self.name}", shape=sigma_shape)
+        diag_idx = np.diag_indices(self.k_endog)
+        idx = np.s_["obs_cov", diag_idx[0], diag_idx[1]]
+        self.ssm[idx] = error_sigma**2
+
+
+class AutoregressiveComponent(Component):
+    r"""
+    Autoregressive timeseries component
+
+    Parameters
+    ----------
+    order: int or sequence of int
+
+        If int, the number of lags to include in the model.
+        If a sequence, an array-like of zeros and ones indicating which lags to include in the model.
+
+    Notes
+    -----
+    An autoregressive component can be thought of as a way o introducing serially correlated errors into the model.
+    The process is modeled:
+
+    .. math::
+        x_t = \sum_{i=1}^p \rho_i x_{t-i}
+
+    Where ``p``, the number of autoregressive terms to model, is the order of the process. By default, all lags up to
+    ``p`` are included in the model. To disable lags, pass a list of zeros and ones to the ``order`` argumnet. For
+    example, ``order=[1, 1, 0, 1]`` would become:
+
+    .. math::
+        x_t = \rho_1 x_{t-1} + \rho_2 x_{t-1} + \rho_4 x_{t-1}
+
+    The coefficient :math:`\rho_3` has been constrained to zero.
+
+    .. warning:: This class is meant to be used as a component in a structural time series model. For modeling of
+              stationary processes with ARIMA, use ``statespace.BayesianSARIMA``.
+
+    Examples
+    --------
+    Model a timeseries as an AR(2) process with non-zero mean:
+
+    .. code:: python
+
+        from pymc_extras.statespace import structural as st
+        import pymc as pm
+        import pytensor.tensor as pt
+
+        trend = st.LevelTrendComponent(order=1, innovations_order=0)
+        ar = st.AutoregressiveComponent(2)
+        ss_mod = (trend + ar).build()
+
+        with pm.Model(coords=ss_mod.coords) as model:
+            P0 = pm.Deterministic('P0', pt.eye(ss_mod.k_states) * 10, dims=ss_mod.param_dims['P0'])
+            intitial_trend = pm.Normal('initial_trend', sigma=10, dims=ss_mod.param_dims['initial_trend'])
+            ar_params = pm.Normal('ar_params', dims=ss_mod.param_dims['ar_params'])
+            sigma_ar = pm.Exponential('sigma_ar', 1, dims=ss_mod.param_dims['sigma_ar'])
+
+            ss_mod.build_statespace_graph(data, mode='JAX')
+            idata = pm.sample(nuts_sampler='numpyro')
+
+    """
+
+    def __init__(self, order: int = 1, name: str = "AutoRegressive"):
+        order = order_to_mask(order)
+        ar_lags = np.flatnonzero(order).ravel().astype(int) + 1
+        k_states = len(order)
+
+        self.order = order
+        self.ar_lags = ar_lags
+
+        super().__init__(
+            name=name,
+            k_endog=1,
+            k_states=k_states,
+            k_posdef=1,
+            measurement_error=True,
+            combine_hidden_states=True,
+            obs_state_idxs=np.r_[[1.0], np.zeros(k_states - 1)],
+        )
+
+    def populate_component_properties(self):
+        self.state_names = [f"L{i + 1}.data" for i in range(self.k_states)]
+        self.shock_names = [f"{self.name}_innovation"]
+        self.param_names = ["ar_params", "sigma_ar"]
+        self.param_dims = {"ar_params": (AR_PARAM_DIM,)}
+        self.coords = {AR_PARAM_DIM: self.ar_lags.tolist()}
+
+        self.param_info = {
+            "ar_params": {
+                "shape": (self.k_states,),
+                "constraints": None,
+                "dims": (AR_PARAM_DIM,),
+            },
+            "sigma_ar": {"shape": (), "constraints": "Positive", "dims": None},
+        }
+
+    def make_symbolic_graph(self) -> None:
+        k_nonzero = int(sum(self.order))
+        ar_params = self.make_and_register_variable("ar_params", shape=(k_nonzero,))
+        sigma_ar = self.make_and_register_variable("sigma_ar", shape=())
+
+        T = np.eye(self.k_states, k=-1)
+        self.ssm["transition", :, :] = T
+        self.ssm["selection", 0, 0] = 1
+        self.ssm["design", 0, 0] = 1
+
+        ar_idx = ("transition", np.zeros(k_nonzero, dtype="int"), np.nonzero(self.order)[0])
+        self.ssm[ar_idx] = ar_params
+
+        cov_idx = ("state_cov", *np.diag_indices(1))
+        self.ssm[cov_idx] = sigma_ar**2
+
+
+class TimeSeasonality(Component):
+    r"""
+    Seasonal component, modeled in the time domain
+
+    Parameters
+    ----------
+    season_length: int
+        The number of periods in a single seasonal cycle, e.g. 12 for monthly data with annual seasonal pattern, 7 for
+        daily data with weekly seasonal pattern, etc.
+
+    innovations: bool, default True
+        Whether to include stochastic innovations in the strength of the seasonal effect
+
+    name: str, default None
+        A name for this seasonal component. Used to label dimensions and coordinates. Useful when multiple seasonal
+        components are included in the same model. Default is ``f"Seasonal[s={season_length}]"``
+
+    state_names: list of str, default None
+        List of strings for seasonal effect labels. If provided, it must be of length ``season_length``. An example
+        would be ``state_names = ['Mon', 'Tue', 'Wed', 'Thur', 'Fri', 'Sat', 'Sun']`` when data is daily with a weekly
+        seasonal pattern (``season_length = 7``).
+
+        If None, states will be numbered ``[State_0, ..., State_s]``
+
+    Notes
+    -----
+    A seasonal effect is any pattern that repeats every fixed interval. Although there are many possible ways to
+    model seasonal effects, the implementation used here is the one described by [1] as the "canonical" time domain
+    representation. The seasonal component can be expressed:
+
+    .. math::
+        \gamma_t = -\sum_{i=1}^{s-1} \gamma_{t-i} + \omega_t, \quad \omega_t \sim N(0, \sigma_\gamma)
+
+    Where :math:`s` is the ``seasonal_length`` parameter and :math:`\omega_t` is the (optional) stochastic innovation.
+    To give interpretation to the :math:`\gamma` terms, it is helpful to work  through the algebra for a simple
+    example. Let :math:`s=4`, and omit the shock term. Define initial conditions :math:`\gamma_0, \gamma_{-1},
+    \gamma_{-2}`. The value of the seasonal component for the first 5 timesteps will be:
+
+    .. math::
+        \begin{align}
+            \gamma_1 &= -\gamma_0 - \gamma_{-1} - \gamma_{-2} \\
+             \gamma_2 &= -\gamma_1 - \gamma_0 - \gamma_{-1} \\
+                       &= -(-\gamma_0 - \gamma_{-1} - \gamma_{-2}) - \gamma_0 - \gamma_{-1}  \\
+                       &= (\gamma_0 - \gamma_0 )+ (\gamma_{-1} - \gamma_{-1}) + \gamma_{-2} \\
+                       &= \gamma_{-2} \\
+              \gamma_3 &= -\gamma_2 - \gamma_1 - \gamma_0  \\
+                       &= -\gamma_{-2} - (-\gamma_0 - \gamma_{-1} - \gamma_{-2}) - \gamma_0 \\
+                       &=  (\gamma_{-2} - \gamma_{-2}) + \gamma_{-1} + (\gamma_0 - \gamma_0) \\
+                       &= \gamma_{-1} \\
+              \gamma_4 &= -\gamma_3 - \gamma_2 - \gamma_1 \\
+                       &= -\gamma_{-1} - \gamma_{-2} -(-\gamma_0 - \gamma_{-1} - \gamma_{-2}) \\
+                       &= (\gamma_{-2} - \gamma_{-2}) + (\gamma_{-1} - \gamma_{-1}) + \gamma_0 \\
+                       &= \gamma_0 \\
+              \gamma_5 &= -\gamma_4 - \gamma_3 - \gamma_2 \\
+                       &= -\gamma_0 - \gamma_{-1} - \gamma_{-2} \\
+                       &= \gamma_1
+        \end{align}
+
+    This exercise shows that, given a list ``initial_conditions`` of length ``s-1``, the effects of this model will be:
+
+        - Period 1: ``-sum(initial_conditions)``
+        - Period 2: ``initial_conditions[-1]``
+        - Period 3: ``initial_conditions[-2]``
+        - ...
+        - Period s: ``initial_conditions[0]``
+        - Period s+1: ``-sum(initial_condition)``
+
+    And so on. So for interpretation, the ``season_length - 1`` initial states are, when reversed, the coefficients
+    associated with ``state_names[1:]``.
+
+    .. warning::
+        Although the ``state_names`` argument expects a list of length ``season_length``, only ``state_names[1:]``
+        will be saved as model dimensions, since the 1st coefficient is not identified (it is defined as
+        :math:`-\sum_{i=1}^{s} \gamma_{t-i}`).
+
+    Examples
+    --------
+    Estimate monthly with a model with a gaussian random walk trend and monthly seasonality:
+
+    .. code:: python
+
+        from pymc_extras.statespace import structural as st
+        import pymc as pm
+        import pytensor.tensor as pt
+        import pandas as pd
+
+        # Get month names
+        state_names = pd.date_range('1900-01-01', '1900-12-31', freq='MS').month_name().tolist()
+
+        # Build the structural model
+        grw = st.LevelTrendComponent(order=1, innovations_order=1)
+        annual_season = st.TimeSeasonality(season_length=12, name='annual', state_names=state_names, innovations=False)
+        ss_mod = (grw + annual_season).build()
+
+        # Estimate with PyMC
+        with pm.Model(coords=ss_mod.coords) as model:
+            P0 = pm.Deterministic('P0', pt.eye(ss_mod.k_states) * 10, dims=ss_mod.param_dims['P0'])
+            intitial_trend = pm.Deterministic('initial_trend', pt.zeros(1), dims=ss_mod.param_dims['initial_trend'])
+            annual_coefs = pm.Normal('annual_coefs', sigma=1e-2, dims=ss_mod.param_dims['annual_coefs'])
+            trend_sigmas = pm.HalfNormal('trend_sigmas', sigma=1e-6, dims=ss_mod.param_dims['trend_sigmas'])
+            ss_mod.build_statespace_graph(data, mode='JAX')
+            idata = pm.sample(nuts_sampler='numpyro')
+
+    References
+    ----------
+    .. [1] Durbin, James, and Siem Jan Koopman. 2012.
+        Time Series Analysis by State Space Methods: Second Edition.
+        Oxford University Press.
+    """
+
+    def __init__(
+        self,
+        season_length: int,
+        innovations: bool = True,
+        name: str | None = None,
+        state_names: list | None = None,
+        pop_state: bool = True,
+    ):
+        if name is None:
+            name = f"Seasonal[s={season_length}]"
+        if state_names is None:
+            state_names = [f"{name}_{i}" for i in range(season_length)]
+        else:
+            if len(state_names) != season_length:
+                raise ValueError(
+                    f"state_names must be a list of length season_length, got {len(state_names)}"
+                )
+            state_names = state_names.copy()
+        self.innovations = innovations
+        self.pop_state = pop_state
+
+        if self.pop_state:
+            # In traditional models, the first state isn't identified, so we can help out the user by automatically
+            # discarding it.
+            # TODO: Can this be stashed and reconstructed automatically somehow?
+            state_names.pop(0)
+            k_states = season_length - 1
+
+        super().__init__(
+            name=name,
+            k_endog=1,
+            k_states=k_states,
+            k_posdef=int(innovations),
+            state_names=state_names,
+            measurement_error=False,
+            combine_hidden_states=True,
+            obs_state_idxs=np.r_[[1.0], np.zeros(k_states - 1)],
+        )
+
+    def populate_component_properties(self):
+        self.param_names = [f"{self.name}_coefs"]
+        self.param_info = {
+            f"{self.name}_coefs": {
+                "shape": (self.k_states,),
+                "constraints": None,
+                "dims": (f"{self.name}_state",),
+            }
+        }
+        self.param_dims = {f"{self.name}_coefs": (f"{self.name}_state",)}
+        self.coords = {f"{self.name}_state": self.state_names}
+
+        if self.innovations:
+            self.param_names += [f"sigma_{self.name}"]
+            self.param_info[f"sigma_{self.name}"] = {
+                "shape": (),
+                "constraints": "Positive",
+                "dims": None,
+            }
+            self.shock_names = [f"{self.name}"]
+
+    def make_symbolic_graph(self) -> None:
+        T = np.eye(self.k_states, k=-1)
+        T[0, :] = -1
+
+        self.ssm["transition", :, :] = T
+        self.ssm["design", 0, 0] = 1
+
+        initial_states = self.make_and_register_variable(
+            f"{self.name}_coefs", shape=(self.k_states,)
+        )
+        self.ssm["initial_state", np.arange(self.k_states, dtype=int)] = initial_states
+
+        if self.innovations:
+            self.ssm["selection", 0, 0] = 1
+            season_sigma = self.make_and_register_variable(f"sigma_{self.name}", shape=())
+            cov_idx = ("state_cov", *np.diag_indices(1))
+            self.ssm[cov_idx] = season_sigma**2
+
+
+class FrequencySeasonality(Component):
+    r"""
+    Seasonal component, modeled in the frequency domain
+
+    Parameters
+    ----------
+    season_length: float
+        The number of periods in a single seasonal cycle, e.g. 12 for monthly data with annual seasonal pattern, 7 for
+        daily data with weekly seasonal pattern, etc. Non-integer seasonal_length is also permitted, for example
+        365.2422 days in a (solar) year.
+
+    n: int
+        Number of fourier features to include in the seasonal component. Default is ``season_length // 2``, which
+        is the maximum possible. A smaller number can be used for a more wave-like seasonal pattern.
+
+    name: str, default None
+        A name for this seasonal component. Used to label dimensions and coordinates. Useful when multiple seasonal
+        components are included in the same model. Default is ``f"Seasonal[s={season_length}, n={n}]"``
+
+    innovations: bool, default True
+        Whether to include stochastic innovations in the strength of the seasonal effect
+
+    Notes
+    -----
+    A seasonal effect is any pattern that repeats every fixed interval. Although there are many possible ways to
+    model seasonal effects, the implementation used here is the one described by [1] as the "canonical" frequency domain
+    representation. The seasonal component can be expressed:
+
+    .. math::
+        \begin{align}
+            \gamma_t &= \sum_{j=1}^{2n} \gamma_{j,t} \\
+            \gamma_{j, t+1} &= \gamma_{j,t} \cos \lambda_j + \gamma_{j,t}^\star \sin \lambda_j + \omega_{j, t} \\
+            \gamma_{j, t}^\star &= -\gamma_{j,t} \sin \lambda_j + \gamma_{j,t}^\star \cos \lambda_j + \omega_{j,t}^\star
+            \lambda_j &= \frac{2\pi j}{s}
+        \end{align}
+
+    Where :math:`s` is the ``seasonal_length``.
+
+    Unlike a ``TimeSeasonality`` component, a ``FrequencySeasonality`` component does not require integer season
+    length. In addition, for long seasonal periods, it is possible to obtain a more compact state space representation
+    by choosing ``n << s // 2``. Using ``TimeSeasonality``, an annual seasonal pattern in daily data requires 364
+    states, whereas ``FrequencySeasonality`` always requires ``2 * n`` states, regardless of the ``seasonal_length``.
+    The price of this compactness is less representational power. At ``n = 1``, the seasonal pattern will be a pure
+    sine wave. At ``n = s // 2``, any arbitrary pattern can be represented.
+
+    One cost of the added flexibility of ``FrequencySeasonality`` is reduced interpretability. States of this model are
+    coefficients :math:`\gamma_1, \gamma^\star_1, \gamma_2, \gamma_2^\star ..., \gamma_n, \gamma^\star_n` associated
+    with different frequencies in the fourier representation of the seasonal pattern. As a result, it is not possible
+    to isolate and identify a "Monday" effect, for instance.
+    """
+
+    def __init__(self, season_length, n=None, name=None, innovations=True):
+        if n is None:
+            n = int(season_length // 2)
+        if name is None:
+            name = f"Frequency[s={season_length}, n={n}]"
+
+        k_states = n * 2
+        self.n = n
+        self.season_length = season_length
+        self.innovations = innovations
+
+        # If the model is completely saturated (n = s // 2), the last state will not be identified, so it shouldn't
+        # get a parameter assigned to it and should just be fixed to zero.
+        # Test this way (rather than n == s // 2) to catch cases when n is non-integer.
+        self.last_state_not_identified = self.season_length / self.n == 2.0
+        self.n_coefs = k_states - int(self.last_state_not_identified)
+
+        obs_state_idx = np.zeros(k_states)
+        obs_state_idx[slice(0, k_states, 2)] = 1
+
+        super().__init__(
+            name=name,
+            k_endog=1,
+            k_states=k_states,
+            k_posdef=k_states * int(self.innovations),
+            measurement_error=False,
+            combine_hidden_states=True,
+            obs_state_idxs=obs_state_idx,
+        )
+
+    def make_symbolic_graph(self) -> None:
+        self.ssm["design", 0, slice(0, self.k_states, 2)] = 1
+
+        init_state = self.make_and_register_variable(f"{self.name}", shape=(self.n_coefs,))
+
+        init_state_idx = np.arange(self.n_coefs, dtype=int)
+        self.ssm["initial_state", init_state_idx] = init_state
+
+        T_mats = [_frequency_transition_block(self.season_length, j + 1) for j in range(self.n)]
+        T = pt.linalg.block_diag(*T_mats)
+        self.ssm["transition", :, :] = T
+
+        if self.innovations:
+            sigma_season = self.make_and_register_variable(f"sigma_{self.name}", shape=())
+            self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * sigma_season**2
+            self.ssm["selection", :, :] = np.eye(self.k_states)
+
+    def populate_component_properties(self):
+        self.state_names = [f"{self.name}_{f}_{i}" for i in range(self.n) for f in ["Cos", "Sin"]]
+        self.param_names = [f"{self.name}"]
+
+        self.param_dims = {self.name: (f"{self.name}_state",)}
+        self.param_info = {
+            f"{self.name}": {
+                "shape": (self.k_states - int(self.last_state_not_identified),),
+                "constraints": None,
+                "dims": (f"{self.name}_state",),
+            }
+        }
+
+        init_state_idx = np.arange(self.k_states, dtype=int)
+        if self.last_state_not_identified:
+            init_state_idx = init_state_idx[:-1]
+        self.coords = {f"{self.name}_state": [self.state_names[i] for i in init_state_idx]}
+
+        if self.innovations:
+            self.shock_names = self.state_names.copy()
+            self.param_names += [f"sigma_{self.name}"]
+            self.param_info[f"sigma_{self.name}"] = {
+                "shape": (),
+                "constraints": "Positive",
+                "dims": None,
+            }
+
+
+class CycleComponent(Component):
+    r"""
+    A component for modeling longer-term cyclical effects
+
+    Parameters
+    ----------
+    name: str
+        Name of the component. Used in generated coordinates and state names. If None, a descriptive name will be
+        used.
+
+    cycle_length: int, optional
+        The length of the cycle, in the calendar units of your data. For example, if your data is monthly, and you
+        want to model a 12-month cycle, use ``cycle_length=12``. You cannot specify both ``cycle_length`` and
+        ``estimate_cycle_length``.
+
+    estimate_cycle_length: bool, default False
+        Whether to estimate the cycle length. If True, an additional parameter, ``cycle_length`` will be added to the
+        model. You cannot specify both ``cycle_length`` and ``estimate_cycle_length``.
+
+    dampen: bool, default False
+        Whether to dampen the cycle by multiplying by a dampening factor :math:`\rho` at every timestep. If true,
+        an additional parameter, ``dampening_factor`` will be added to the model.
+
+    innovations: bool, default True
+        Whether to include stochastic innovations in the strength of the seasonal effect. If True, an additional
+        parameter, ``sigma_{name}`` will be added to the model.
+
+    Notes
+    -----
+    The cycle component is very similar in implementation to the frequency domain seasonal component, expect that it
+    is restricted to n=1. The cycle component can be expressed:
+
+    .. math::
+        \begin{align}
+            \gamma_t &= \rho \gamma_{t-1} \cos \lambda + \rho \gamma_{t-1}^\star \sin \lambda + \omega_{t} \\
+            \gamma_{t}^\star &= -\rho \gamma_{t-1} \sin \lambda + \rho \gamma_{t-1}^\star \cos \lambda + \omega_{t}^\star \\
+            \lambda &= \frac{2\pi}{s}
+        \end{align}
+
+    Where :math:`s` is the ``cycle_length``. [1] recommend that this component be used for longer term cyclical
+    effects, such as business cycles, and that the seasonal component be used for shorter term effects, such as
+    weekly or monthly seasonality.
+
+    Unlike a FrequencySeasonality component, the length of a CycleComponent can be estimated.
+
+    Examples
+    --------
+    Estimate a business cycle with length between 6 and 12 years:
+
+    .. code:: python
+
+        from pymc_extras.statespace import structural as st
+        import pymc as pm
+        import pytensor.tensor as pt
+        import pandas as pd
+        import numpy as np
+
+        data = np.random.normal(size=(100, 1))
+
+        # Build the structural model
+        grw = st.LevelTrendComponent(order=1, innovations_order=1)
+        cycle = st.CycleComponent('business_cycle', estimate_cycle_length=True, dampen=False)
+        ss_mod = (grw + cycle).build()
+
+        # Estimate with PyMC
+        with pm.Model(coords=ss_mod.coords) as model:
+            P0 = pm.Deterministic('P0', pt.eye(ss_mod.k_states), dims=ss_mod.param_dims['P0'])
+            intitial_trend = pm.Normal('initial_trend', dims=ss_mod.param_dims['initial_trend'])
+            sigma_trend = pm.HalfNormal('sigma_trend', dims=ss_mod.param_dims['sigma_trend'])
+
+            cycle_strength = pm.Normal('business_cycle')
+            cycle_length = pm.Uniform('business_cycle_length', lower=6, upper=12)
+
+            sigma_cycle = pm.HalfNormal('sigma_business_cycle', sigma=1)
+            ss_mod.build_statespace_graph(data, mode='JAX')
+
+            idata = pm.sample(nuts_sampler='numpyro')
+
+    References
+    ----------
+    .. [1] Durbin, James, and Siem Jan Koopman. 2012.
+        Time Series Analysis by State Space Methods: Second Edition.
+        Oxford University Press.
+    """
+
+    def __init__(
+        self,
+        name: str | None = None,
+        cycle_length: int | None = None,
+        estimate_cycle_length: bool = False,
+        dampen: bool = False,
+        innovations: bool = True,
+    ):
+        if cycle_length is None and not estimate_cycle_length:
+            raise ValueError("Must specify cycle_length if estimate_cycle_length is False")
+        if cycle_length is not None and estimate_cycle_length:
+            raise ValueError("Cannot specify cycle_length if estimate_cycle_length is True")
+        if name is None:
+            cycle = int(cycle_length) if cycle_length is not None else "Estimate"
+            name = f"Cycle[s={cycle}, dampen={dampen}, innovations={innovations}]"
+
+        self.estimate_cycle_length = estimate_cycle_length
+        self.cycle_length = cycle_length
+        self.innovations = innovations
+        self.dampen = dampen
+        self.n_coefs = 1
+
+        k_states = 2
+        k_endog = 1
+        k_posdef = 2
+
+        obs_state_idx = np.zeros(k_states)
+        obs_state_idx[slice(0, k_states, 2)] = 1
+
+        super().__init__(
+            name=name,
+            k_endog=k_endog,
+            k_states=k_states,
+            k_posdef=k_posdef,
+            measurement_error=False,
+            combine_hidden_states=True,
+            obs_state_idxs=obs_state_idx,
+        )
+
+    def make_symbolic_graph(self) -> None:
+        self.ssm["design", 0, slice(0, self.k_states, 2)] = 1
+        self.ssm["selection", :, :] = np.eye(self.k_states)
+        self.param_dims = {self.name: (f"{self.name}_state",)}
+        self.coords = {f"{self.name}_state": self.state_names}
+
+        init_state = self.make_and_register_variable(f"{self.name}", shape=(self.k_states,))
+
+        self.ssm["initial_state", :] = init_state
+
+        if self.estimate_cycle_length:
+            lamb = self.make_and_register_variable(f"{self.name}_length", shape=())
+        else:
+            lamb = self.cycle_length
+
+        if self.dampen:
+            rho = self.make_and_register_variable(f"{self.name}_dampening_factor", shape=())
+        else:
+            rho = 1
+
+        T = rho * _frequency_transition_block(lamb, j=1)
+        self.ssm["transition", :, :] = T
+
+        if self.innovations:
+            sigma_cycle = self.make_and_register_variable(f"sigma_{self.name}", shape=())
+            self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * sigma_cycle**2
+
+    def populate_component_properties(self):
+        self.state_names = [f"{self.name}_{f}" for f in ["Cos", "Sin"]]
+        self.param_names = [f"{self.name}"]
+
+        self.param_info = {
+            f"{self.name}": {
+                "shape": (2,),
+                "constraints": None,
+                "dims": (f"{self.name}_state",),
+            }
+        }
+
+        if self.estimate_cycle_length:
+            self.param_names += [f"{self.name}_length"]
+            self.param_info[f"{self.name}_length"] = {
+                "shape": (),
+                "constraints": "Positive, non-zero",
+                "dims": None,
+            }
+
+        if self.dampen:
+            self.param_names += [f"{self.name}_dampening_factor"]
+            self.param_info[f"{self.name}_dampening_factor"] = {
+                "shape": (),
+                "constraints": "0 < x ≤ 1",
+                "dims": None,
+            }
+
+        if self.innovations:
+            self.param_names += [f"sigma_{self.name}"]
+            self.param_info[f"sigma_{self.name}"] = {
+                "shape": (),
+                "constraints": "Positive",
+                "dims": None,
+            }
+            self.shock_names = self.state_names.copy()
+
+
+class RegressionComponent(Component):
+    def __init__(
+        self,
+        k_exog: int | None = None,
+        name: str | None = "Exogenous",
+        state_names: list[str] | None = None,
+        innovations=False,
+    ):
+        self.innovations = innovations
+        k_exog = self._handle_input_data(k_exog, state_names, name)
+
+        k_states = k_exog
+        k_endog = 1
+        k_posdef = k_exog
+
+        super().__init__(
+            name=name,
+            k_endog=k_endog,
+            k_states=k_states,
+            k_posdef=k_posdef,
+            state_names=self.state_names,
+            measurement_error=False,
+            combine_hidden_states=False,
+            exog_names=[f"data_{name}"],
+            obs_state_idxs=np.ones(k_states),
+        )
+
+    @staticmethod
+    def _get_state_names(k_exog: int | None, state_names: list[str] | None, name: str):
+        if k_exog is None and state_names is None:
+            raise ValueError("Must specify at least one of k_exog or state_names")
+        if state_names is not None and k_exog is not None:
+            if len(state_names) != k_exog:
+                raise ValueError(f"Expected {k_exog} state names, found {len(state_names)}")
+        elif k_exog is None:
+            k_exog = len(state_names)
+        else:
+            state_names = [f"{name}_{i + 1}" for i in range(k_exog)]
+
+        return k_exog, state_names
+
+    def _handle_input_data(self, k_exog: int, state_names: list[str] | None, name) -> int:
+        k_exog, state_names = self._get_state_names(k_exog, state_names, name)
+        self.state_names = state_names
+
+        return k_exog
+
+    def make_symbolic_graph(self) -> None:
+        betas = self.make_and_register_variable(f"beta_{self.name}", shape=(self.k_states,))
+        regression_data = self.make_and_register_data(
+            f"data_{self.name}", shape=(None, self.k_states)
+        )
+
+        self.ssm["initial_state", :] = betas
+        self.ssm["transition", :, :] = np.eye(self.k_states)
+        self.ssm["selection", :, :] = np.eye(self.k_states)
+        self.ssm["design"] = pt.expand_dims(regression_data, 1)
+
+        if self.innovations:
+            sigma_beta = self.make_and_register_variable(
+                f"sigma_beta_{self.name}", (self.k_states,)
+            )
+            row_idx, col_idx = np.diag_indices(self.k_states)
+            self.ssm["state_cov", row_idx, col_idx] = sigma_beta**2
+
+    def populate_component_properties(self) -> None:
+        self.shock_names = self.state_names
+
+        self.param_names = [f"beta_{self.name}"]
+        self.data_names = [f"data_{self.name}"]
+        self.param_dims = {
+            f"beta_{self.name}": ("exog_state",),
+        }
+
+        self.param_info = {
+            f"beta_{self.name}": {
+                "shape": (self.k_states,),
+                "constraints": None,
+                "dims": ("exog_state",),
+            },
+        }
+
+        self.data_info = {
+            f"data_{self.name}": {
+                "shape": (None, self.k_states),
+                "dims": (TIME_DIM, "exog_state"),
+            },
+        }
+        self.coords = {"exog_state": self.state_names}
+
+        if self.innovations:
+            self.param_names += [f"sigma_beta_{self.name}"]
+            self.param_dims[f"sigma_beta_{self.name}"] = "exog_state"
+            self.param_info[f"sigma_beta_{self.name}"] = {
+                "shape": (),
+                "constraints": "Positive",
+                "dims": ("exog_state",),
+            }